"""Hybrid search — FTS5 keyword + embedding similarity + Claude synthesis."""

from __future__ import annotations

import hashlib
import re
import time
from collections import defaultdict

import numpy as np
from rich.console import Console

from .config import Config
from .db import Database

console = Console()

ASK_PROMPT = """\
You are an expert on IETF Internet-Drafts related to AI agents and autonomous systems.

Based on the following IETF drafts, answer this question:

**{question}**

## Source Drafts

{sources_block}

Instructions:
- Answer concisely but thoroughly (3-8 sentences).
- Cite specific drafts by name (e.g. draft-xyz-...) when referencing their contributions.
- If the drafts don't contain enough information to answer, say so clearly.
- Focus on concrete technical mechanisms, not general observations."""


def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    dot = np.dot(a, b)
    norm = np.linalg.norm(a) * np.linalg.norm(b)
    if norm == 0:
        return 0.0
    return float(dot / norm)


def _prompt_hash(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()[:16]


class HybridSearch:
    def __init__(self, config: Config, db: Database, embedder=None):
        self.config = config
        self.db = db
        self._embedder = embedder
        self._ollama_available: bool | None = None
        self._embeddings_cache: dict[str, np.ndarray] | None = None
        self._embeddings_cache_time: float = 0
        self._EMBEDDINGS_TTL: float = 300  # 5 minutes

    @property
    def embedder(self):
        """Lazy-load embedder to avoid import errors when Ollama is unavailable."""
        if self._embedder is None:
            try:
                from .embeddings import Embedder
                self._embedder = Embedder(self.config, self.db)
                self._ollama_available = True
            except Exception:
                self._ollama_available = False
        return self._embedder

    def _check_ollama(self) -> bool:
        """Check if Ollama is available for embedding queries."""
        if self._ollama_available is not None:
            return self._ollama_available
        try:
            embedder = self.embedder
            if embedder is None:
                self._ollama_available = False
                return False
            # Try a tiny embedding to verify connectivity
            embedder.embed_text("test")
            self._ollama_available = True
        except Exception:
            self._ollama_available = False
        return self._ollama_available

    def _get_all_embeddings(self) -> dict[str, np.ndarray]:
        """Return all embeddings, cached with TTL to avoid reloading on every query."""
        now = time.monotonic()
        if (self._embeddings_cache is not None
                and now - self._embeddings_cache_time < self._EMBEDDINGS_TTL):
            return self._embeddings_cache
        self._embeddings_cache = self.db.all_embeddings()
        self._embeddings_cache_time = now
        return self._embeddings_cache

    @staticmethod
    def _base_draft_name(name: str) -> str:
        """Strip revision suffix to get the base draft name for dedup.

        e.g. 'draft-wang-cats-odsi-02' → 'draft-wang-cats-odsi'
        """
        return re.sub(r'-\d{2,3}$', '', name)

    def search(self, query: str, top_k: int = 10) -> list[dict]:
        """Combine FTS5 keyword search + embedding similarity search.

        Returns ranked list of {name, title, score, excerpt, match_type}.
        Falls back to FTS5-only if Ollama is unavailable.
        Deduplicates draft revisions, keeping the highest-scoring version.
        """
        fts_results = self._fts_search(query, limit=top_k * 3)
        embed_results = self._embedding_search(query, limit=top_k * 3)

        if embed_results:
            merged = self._reciprocal_rank_fusion(fts_results, embed_results)
        else:
            merged = fts_results

        # Deduplicate revisions of the same draft, keep best score
        seen_bases: dict[str, int] = {}
        deduped = []
        for r in merged:
            base = self._base_draft_name(r["name"])
            if base not in seen_bases:
                seen_bases[base] = len(deduped)
                deduped.append(r)
            # else: skip lower-ranked revision

        return deduped[:top_k]

    @staticmethod
    def sanitize_fts_query(query: str) -> str:
        """Sanitize a query string for safe use in FTS5 MATCH expressions.

        Strips special FTS5 characters and operators (NEAR, OR, AND, NOT)
        that could cause injection or query syntax errors.
        """
        # Strip all non-alphanumeric, non-whitespace characters
        cleaned = re.sub(r'[^\w\s]', '', query)
        # Remove FTS5 boolean operators when used as standalone words
        cleaned = re.sub(r'\b(NEAR|OR|AND|NOT)\b', '', cleaned, flags=re.IGNORECASE)
        # Collapse whitespace
        cleaned = re.sub(r'\s+', ' ', cleaned).strip()
        return cleaned

    def _fts_search(self, query: str, limit: int = 20) -> list[dict]:
        """Run FTS5 keyword search, return ranked results."""
        safe_query = self.sanitize_fts_query(query)
        if not safe_query:
            return []
        try:
            drafts = self.db.search_drafts(safe_query, limit=limit)
        except Exception:
            return []

        results = []
        for rank, draft in enumerate(drafts):
            excerpt = draft.abstract[:200] if draft.abstract else ""
            results.append({
                "name": draft.name,
                "title": draft.title,
                "score": 0.0,  # Will be set by fusion
                "excerpt": excerpt,
                "match_type": "keyword",
                "rank": rank,
            })
        return results

    def _embedding_search(self, query: str, limit: int = 20) -> list[dict]:
        """Run embedding similarity search against all stored embeddings."""
        if not self._check_ollama():
            return []

        try:
            query_vec = self.embedder.embed_text(query)
        except Exception:
            self._ollama_available = False
            return []

        all_embeddings = self._get_all_embeddings()
        if not all_embeddings:
            return []

        similarities: list[tuple[str, float]] = []
        for name, vec in all_embeddings.items():
            sim = _cosine_similarity(query_vec, vec)
            similarities.append((name, sim))

        similarities.sort(key=lambda x: x[1], reverse=True)

        results = []
        for rank, (name, sim) in enumerate(similarities[:limit]):
            draft = self.db.get_draft(name)
            if draft is None:
                continue
            excerpt = draft.abstract[:200] if draft.abstract else ""
            results.append({
                "name": name,
                "title": draft.title if draft else name,
                "score": round(sim, 4),
                "excerpt": excerpt,
                "match_type": "semantic",
                "rank": rank,
                "similarity": round(sim, 4),
            })
        return results

    def _reciprocal_rank_fusion(
        self,
        fts_results: list[dict],
        embed_results: list[dict],
        k: int = 60,
    ) -> list[dict]:
        """Merge two ranked lists using reciprocal rank fusion (RRF).

        RRF score = sum(1 / (k + rank)) across all lists where the item appears.
        """
        scores: dict[str, float] = defaultdict(float)
        items: dict[str, dict] = {}

        for result in fts_results:
            name = result["name"]
            scores[name] += 1.0 / (k + result["rank"])
            if name not in items:
                items[name] = result.copy()
                items[name]["match_type"] = "keyword"

        for result in embed_results:
            name = result["name"]
            scores[name] += 1.0 / (k + result["rank"])
            if name in items:
                items[name]["match_type"] = "both"
                items[name]["similarity"] = result.get("similarity", 0)
            else:
                items[name] = result.copy()

        # Sort by RRF score
        ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)

        results = []
        for name, rrf_score in ranked:
            item = items[name]
            item["score"] = round(rrf_score, 4)
            results.append(item)

        return results

    def _build_sources(self, search_results: list[dict]) -> tuple[list[dict], str]:
        """Build source list and context block from search results."""
        sources_block = ""
        sources = []
        for r in search_results:
            draft = self.db.get_draft(r["name"])
            if draft is None:
                continue
            text_preview = draft.full_text[:500] if draft.full_text else ""
            sources_block += f"\n---\n**{draft.name}** — {draft.title}\n"
            sources_block += f"Abstract: {draft.abstract[:500]}\n"
            if text_preview:
                sources_block += f"Content excerpt: {text_preview}\n"
            sources.append({
                "name": draft.name,
                "title": draft.title,
                "similarity": r.get("similarity", r.get("score", 0)),
                "excerpt": r.get("excerpt", ""),
                "match_type": r.get("match_type", ""),
            })
        return sources, sources_block

    def search_only(self, question: str, top_k: int = 5) -> dict:
        """Search without Claude synthesis (free — only FTS5 + Ollama).

        Returns {sources: [...], has_cached_answer: bool, answer: str|None}.
        """
        search_results = self.search(question, top_k=top_k)
        if not search_results:
            return {"sources": [], "has_cached_answer": False, "answer": None}

        sources, sources_block = self._build_sources(search_results)

        # Check if we already have a cached answer for this question
        prompt = ASK_PROMPT.format(question=question, sources_block=sources_block)
        phash = _prompt_hash(prompt)
        cached = self.db.get_cached_response("_ask_", phash)

        return {
            "sources": sources,
            "has_cached_answer": cached is not None,
            "answer": cached,
        }

    def ask(self, question: str, top_k: int = 5, cheap: bool = True) -> dict:
        """Answer a natural language question using search + Claude synthesis.

        Returns {answer: str, sources: [{name, title, similarity, excerpt}]}.
        Caches Claude responses via llm_cache permanently.
        """
        search_results = self.search(question, top_k=top_k)

        if not search_results:
            return {
                "answer": "No relevant drafts found for your question.",
                "sources": [],
            }

        sources, sources_block = self._build_sources(search_results)

        prompt = ASK_PROMPT.format(
            question=question,
            sources_block=sources_block,
        )
        phash = _prompt_hash(prompt)

        # Check cache — cached answers are served forever (no TTL)
        cached = self.db.get_cached_response("_ask_", phash)
        if cached:
            return {
                "answer": cached,
                "sources": sources,
            }

        # Call Claude
        try:
            from .analyzer import Analyzer
            analyzer = Analyzer(self.config, self.db)
            text, in_tok, out_tok = analyzer._call_claude(
                prompt, max_tokens=1024, cheap=cheap
            )

            # Cache permanently
            self.db.cache_response(
                "_ask_", phash,
                self.config.claude_model_cheap if cheap else self.config.claude_model,
                prompt, text, in_tok, out_tok,
            )

            return {
                "answer": text,
                "sources": sources,
            }
        except Exception as e:
            return {
                "answer": f"Error generating answer: {e}",
                "sources": sources,
            }
