# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging
import re
from dataclasses import dataclass, field

from .embeddings import embed_texts, local_embed
from .llm import complete_chat
from .store import Chunk, VectorStore

logger = logging.getLogger(__name__)

CITATION_MARKER = re.compile(r"\[(\d+)\]")

CHAT_TOP_K = 10
MAX_CONTEXT_CHARS = 16000
CHAT_MAX_TOKENS = 1400

SYSTEM_PROMPT = (
    "You are the DeepSearch research assistant. Answer the user's question using ONLY "
    "the numbered SOURCES below, which were gathered during a web research session. "
    "Never use outside knowledge or guess. If the sources do not contain the answer, "
    "say so plainly. Cite every claim inline with the bracket marker of the source it "
    "comes from, like [1] or [2]. Keep the answer focused and well structured in "
    "markdown."
)


@dataclass
class ChatAnswer:
    text: str
    citations: list[dict] = field(default_factory=list)


def _build_context(chunks: list[Chunk]) -> tuple[str, list[dict]]:
    blocks: list[str] = []
    citations: list[dict] = []
    used = 0
    for index, chunk in enumerate(chunks, start=1):
        snippet = chunk.text.strip()
        if not snippet:
            continue
        header = f"[{index}] {chunk.title or chunk.url} ({chunk.url})"
        block = f"{header}\n{snippet}"
        if used + len(block) > MAX_CONTEXT_CHARS and blocks:
            break
        used += len(block)
        blocks.append(block)
        citations.append(
            {
                "index": index,
                "url": chunk.url,
                "title": chunk.title or chunk.url,
                "score": round(chunk.score, 4),
            }
        )
    return "\n\n".join(blocks), citations


class DeepsearchChat:
    def __init__(self, collection_name: str, api_key: str) -> None:
        self.store = VectorStore(collection_name)
        self.api_key = api_key

    async def _embed_query(self, question: str) -> list[float]:
        result = await embed_texts([question], self.api_key)
        if not result.vectors or not result.vectors[0]:
            result = local_embed([question])
        return result.vectors[0]

    async def retrieve(self, question: str) -> list[Chunk]:
        query_vector = await self._embed_query(question)
        stored_dim = self.store.dims
        if stored_dim is not None and len(query_vector) != stored_dim:
            query_vector = local_embed([question]).vectors[0]
            if len(query_vector) != stored_dim:
                logger.warning(
                    "deepsearch query embedding dim %d != stored %d",
                    len(query_vector),
                    stored_dim,
                )
                return []
        return self.store.hybrid_search(question, query_vector, top_k=CHAT_TOP_K)

    async def answer(self, question: str, history: list[dict] | None = None) -> ChatAnswer:
        chunks = await self.retrieve(question)
        if not chunks:
            return ChatAnswer(
                text=(
                    "The research session did not capture anything relevant to that "
                    "question. Try rephrasing or running a deeper search."
                ),
                citations=[],
            )
        context, citations = _build_context(chunks)
        messages = [{"role": "system", "content": SYSTEM_PROMPT}]
        for turn in (history or [])[-6:]:
            role = turn.get("role")
            content = turn.get("content")
            if role in ("user", "assistant") and content:
                messages.append({"role": role, "content": content})
        messages.append(
            {
                "role": "user",
                "content": f"SOURCES:\n{context}\n\nQUESTION: {question}",
            }
        )
        try:
            text = await complete_chat(
                messages, self.api_key, max_tokens=CHAT_MAX_TOKENS
            )
        except Exception as exc:
            logger.warning("deepsearch chat synthesis failed: %s", exc)
            text = (
                "I could not reach the language model to synthesise an answer, but the "
                "most relevant sources are listed below."
            )
        valid = {citation["index"] for citation in citations}
        text = self._strip_unmatched_markers(text, valid)
        return ChatAnswer(text=text, citations=citations)

    def _strip_unmatched_markers(self, text: str, valid: set[int]) -> str:
        def replace(match: re.Match) -> str:
            index = int(match.group(1))
            return match.group(0) if index in valid else ""

        return CITATION_MARKER.sub(replace, text)