import numpy as np

from src import config
from src.chunking import split_text
from src.embeddings import build_cosine_index, embed_texts


class EvidenceError(Exception):
    """Raised when there is no readable text to build evidence from."""


class EvidenceStore:
    """All source text, chunked and indexed in an in-memory FAISS cosine index.

    Used three ways: retrieve passages per sub-question for the Writer (`gather`),
    and pull the passages that back a specific claim for the Checker (`support_for`).
    """

    def __init__(self, sources: list, embedder=None):
        self.sources = sources
        self.embedder = embedder
        self.records = []
        for source_idx, source in enumerate(sources):
            for piece in split_text(source["content"], config.CHUNK_SIZE, config.CHUNK_OVERLAP):
                self.records.append({"id": len(self.records), "text": piece, "source_idx": source_idx})
        if not self.records:
            raise EvidenceError("The sources contained no readable text.")

        self.vectors = embed_texts([r["text"] for r in self.records], embedder)
        self.index = build_cosine_index(self.vectors)
        self.numbering = {}   # source_idx -> citation number, assigned by gather()
        self.selected = []

    def _search(self, query: str, k: int):
        query_vector = embed_texts([query], self.embedder)
        scores, ids = self.index.search(query_vector, min(len(self.records), k))
        return [(self.records[i], float(s)) for s, i in zip(scores[0], ids[0]) if i >= 0]

    def gather(self, questions: list, chunks_per_question: int) -> tuple:
        """Selects the best passages for each sub-question; returns (selected_chunks, numbered_sources)."""
        chosen, chosen_ids, per_source_total = [], set(), {}

        for q_idx, question in enumerate(questions):
            kept, per_source_q = 0, {}
            for record, score in self._search(question, chunks_per_question * 6):
                if kept >= chunks_per_question or len(chosen) >= config.MAX_EVIDENCE_CHUNKS:
                    break
                sid = record["source_idx"]
                if per_source_q.get(sid, 0) >= config.MAX_PER_SOURCE_PER_QUESTION:
                    continue
                if record["id"] in chosen_ids:
                    # Already picked for an earlier question: it still covers this one, so count it
                    # instead of padding this question with weaker passages.
                    if score >= config.MIN_SCORE:
                        per_source_q[sid] = per_source_q.get(sid, 0) + 1
                        kept += 1
                    continue
                if per_source_total.get(sid, 0) >= config.MAX_CHUNKS_PER_SOURCE:
                    continue
                if score < config.MIN_SCORE and kept >= config.MIN_KEEP_PER_QUESTION:
                    continue
                chosen_ids.add(record["id"])
                chosen.append({**record, "score": score, "question_idx": q_idx})
                per_source_q[sid] = per_source_q.get(sid, 0) + 1
                per_source_total[sid] = per_source_total.get(sid, 0) + 1
                kept += 1

        if not chosen:
            raise EvidenceError("No relevant passages were found in the sources.")

        # Number sources by first appearance so citations read [1], [2], [3]... in order
        self.numbering = {}
        for chunk in chosen:
            self.numbering.setdefault(chunk["source_idx"], len(self.numbering) + 1)
        for chunk in chosen:
            chunk["cite"] = self.numbering[chunk["source_idx"]]
        self.selected = chosen

        numbered = [
            {
                "n": n,
                "title": self.sources[sid]["title"],
                "url": self.sources[sid].get("url", ""),
                "kind": self.sources[sid].get("kind", "web"),
            }
            for sid, n in sorted(self.numbering.items(), key=lambda kv: kv[1])
        ]
        return chosen, numbered

    def support_for(self, claim: str, cites: list, k: int = 2) -> list:
        """The k passages from the cited sources that are closest to a claim (what the Checker reads)."""
        wanted = {sid for sid, n in self.numbering.items() if n in set(cites)}
        candidates = [r["id"] for r in self.records if r["source_idx"] in wanted]
        if not candidates:
            return []
        claim_vector = embed_texts([claim], self.embedder)[0]
        sims = self.vectors[candidates] @ claim_vector
        best = np.argsort(-sims)[:k]
        return [self.records[candidates[i]]["text"] for i in best]