from rank_bm25 import BM25Okapi
import numpy as np

class HybridRetriever:
    def __init__(self, chunks, embeddings, embed_fn, alpha=0.5):
        """
        alpha: weight for dense scores (0 = pure BM25, 1 = pure dense)
        """
        self.chunks = chunks
        self.embeddings = embeddings
        self.embed_fn = embed_fn
        self.alpha = alpha
        
        # BM25 index
        tokenized = [chunk.split() for chunk in chunks]
        self.bm25 = BM25Okapi(tokenized)
    
    def retrieve(self, query: str, top_k: int = 5) -> list[str]:
        # Dense scores
        query_emb = self.embed_fn(query)
        dense_scores = np.dot(self.embeddings, query_emb)
        dense_ranks = np.argsort(-dense_scores)
        
        # Sparse scores
        sparse_scores = self.bm25.get_scores(query.split())
        sparse_ranks = np.argsort(-sparse_scores)
        
        # RRF merge
        rrf_scores = {}
        k = 60  # RRF constant
        
        for rank, idx in enumerate(dense_ranks):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + self.alpha / (k + rank + 1)
        
        for rank, idx in enumerate(sparse_ranks):
            rrf_scores[idx] = rrf_scores.get(idx, 0) + (1 - self.alpha) / (k + rank + 1)
        
        sorted_indices = sorted(rrf_scores, key=rrf_scores.get, reverse=True)
        return [self.chunks[i] for i in sorted_indices[:top_k]]

async def rewrite_query_for_retrieval(
    conversation_history: list[dict],
    current_query: str,
    llm_client
) -> str:
    """Rewrite current_query into a standalone question with full context."""
    
    if len(conversation_history) < 2:
        return current_query  # No context to add
    
    history_text = "\n".join([
        f"{msg['role'].capitalize()}: {msg['content']}"
        for msg in conversation_history[-4:]  # Last 2 turns
    ])
    
    response = await llm_client.chat.completions.create(
        model="gpt-4o-mini",  # Use cheap model for rewriting
        messages=[
            {
                "role": "system",
                "content": "Rewrite the user's follow-up question as a complete, "
                           "standalone question that includes all necessary context "
                           "from the conversation. Output only the rewritten question."
            },
            {
                "role": "user",
                "content": f"Conversation:\n{history_text}\n\n"
                           f"Follow-up question: {current_query}\n\n"
                           "Standalone question:"
            }
        ],
        max_tokens=150
    )
    
    return response.choices[0].message.content.strip()

class SentenceWindowIndex:
    """
    Index individual sentences for precise retrieval,
    but return surrounding context window when a match is found.
    """
    def __init__(self, documents: list[str], window_size: int = 3):
        self.window_size = window_size
        self.sentences = []
        self.doc_sentence_map = []  # (doc_idx, sent_idx) for each stored sentence
        
        for doc_idx, doc in enumerate(documents):
            sents = self._split_sentences(doc)
            for sent_idx, sent in enumerate(sents):
                self.sentences.append(sent)
                self.doc_sentence_map.append((doc_idx, sent_idx, len(sents)))
    
    def retrieve_with_context(self, query_embedding, doc_sentences_by_doc, top_k=5):
        """Returns sentence matches expanded to their surrounding window."""
        # ... vector search to find matching sentence indices
        results = []
        for sent_idx in top_matching_indices:
            doc_idx, s_idx, total_sents = self.doc_sentence_map[sent_idx]
            
            # Expand to window
            start = max(0, s_idx - self.window_size)
            end = min(total_sents, s_idx + self.window_size + 1)
            
            window_sentences = doc_sentences_by_doc[doc_idx][start:end]
            results.append(" ".join(window_sentences))
        
        return results
    
    def _split_sentences(self, text: str) -> list[str]:
        import re
        return re.split(r'(?<=[.!?])\s+', text)

GROUNDED_SYSTEM_PROMPT = """
You are a helpful assistant. Answer using ONLY the provided context.
For every claim, cite the source passage in [brackets].
If the context doesn't contain enough information, say so explicitly.
Do not add information from your training data.
"""

async def answer_with_grounding_check(
    query: str,
    context_chunks: list[str],
    llm_client
) -> dict:
    
    context_text = "\n\n".join([
        f"[Source {i+1}]: {chunk}"
        for i, chunk in enumerate(context_chunks)
    ])
    
    # Get grounded answer
    response = await llm_client.chat.completions.create(
        model="gpt-5",
        messages=[
            {"role": "system", "content": GROUNDED_SYSTEM_PROMPT},
            {"role": "user", "content": f"Context:\n{context_text}\n\nQuestion: {query}"}
        ]
    )
    
    answer = response.choices[0].message.content
    
    # Faithfulness check: ask the model to verify its own claims
    faithfulness_check = await llm_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {
                "role": "user",
                "content": f"Context:\n{context_text}\n\n"
                           f"Answer: {answer}\n\n"
                           "Does every claim in the answer appear in the context? "
                           "Reply with JSON: "
            }
        ],
        response_format={"type": "json_object"}
    )
    
    import json
    check_result = json.loads(faithfulness_check.choices[0].message.content)
    
    return {
        "answer": answer,
        "faithful": check_result["faithful"],
        "issues": check_result.get("issues", [])
    }

eval_dataset = [
    {
        "question": "What is the refund policy for annual plans?",
        "expected_answer": "30-day money-back guarantee",
        "source_chunk_id": "pricing-faq-chunk-12",
        "difficulty": "easy"
    },
    {
        "question": "Can I downgrade mid-billing cycle?",
        "expected_answer": "Yes, but changes take effect next billing period",
        "source_chunk_id": "billing-chunk-7",
        "difficulty": "medium"
    },
    {
        "question": "What's the best plan for a team of 5?",
        "expected_answer": None,  # Opinion question, no "correct" answer
        "source_chunk_id": None,
        "difficulty": "hard"
    }
]