from __future__ import annotations

import logging
import re
from typing import Any, Callable, Dict, List, Optional, Sequence

from pydantic import BaseModel, Field

from .dense import DenseRetriever
from .fusion import UnifiedCandidate, reciprocal_rank_fusion
from ..indexing.models import SearchResult

logger = logging.getLogger(__name__)


class MultiQueryRetrievalResponse(BaseModel):
    """Structured response from multi-query retriever."""
    original_query: str
    generated_queries: List[str] = Field(default_factory=list)
    candidate_chunks: int
    returned_chunks: int
    candidate_document_ids: List[str] = Field(default_factory=list)
    results: List[UnifiedCandidate] = Field(default_factory=list)


class MultiQueryRetriever:
    """
    Generates multiple query perspectives to overcome vocabulary mismatch and
    retrieves/fuses candidate chunks across all formulations using RRF.
    """

    def __init__(
        self,
        base_retriever: Optional[DenseRetriever] = None,
        generator_fn: Optional[Callable[[str, int], List[str]]] = None,
        num_queries: int = 3,
        default_top_k: int = 10,
    ) -> None:
        self.base_retriever = base_retriever or DenseRetriever()
        self.generator_fn = generator_fn
        self.num_queries = num_queries
        self.default_top_k = default_top_k
        self._query_cache: Dict[str, List[str]] = {}

    def _generate_queries_heuristic(self, query: str) -> List[str]:
        """
        High-precision rule-based query expansion and perspective generation.
        Extracts key noun phrases, relational clauses, and focused entity variants
        without adding diluting conversational fluff.
        """
        q = query.strip()
        words = q.split()
        queries = [q]

        # 1. Multi-hop & relational clause splitting
        relation_patterns = [
            r"\b(?:connect(?:s|ed)?\s+to|interact(?:s|ed)?\s+with|jointly\s+impact|influence(?:s|d)?\s+whether)\b",
            r"\b(?:compared\s+to|in\s+contrast\s+to|versus|vs\.?)\b",
            r"\b(?:prevent(?:s|ed)?\s+.*?\s+in\s+the|because\s+of|due\s+to)\b",
        ]
        split_done = False
        for pat in relation_patterns:
            parts = re.split(pat, q, flags=re.IGNORECASE)
            if len(parts) >= 2:
                for part in parts:
                    cleaned_part = re.sub(r"^(?:how|why|what|explain|does|do|the)\s+", "", part.strip(), flags=re.IGNORECASE).strip(" ?.,")
                    if len(cleaned_part.split()) >= 2:
                        queries.append(cleaned_part)
                split_done = True
                break

        # 2. Key phrase and entity extraction (keyword-dense core)
        stopwords = {
            "how", "why", "what", "when", "where", "which", "does", "do", "explain",
            "tell", "about", "describe", "the", "a", "an", "is", "are", "was", "were",
            "in", "on", "at", "by", "for", "with", "from", "and", "or", "to", "of",
        }
        content_tokens = [w for w in re.findall(r"[a-zA-Z0-9_\-\.]+", q) if w.lower() not in stopwords]
        if len(content_tokens) >= 3:
            keyword_core = " ".join(content_tokens)
            if keyword_core not in queries:
                queries.append(keyword_core)

        # 3. Domain-specific synonym and terminology expansion
        synonym_map = {
            "convs2s": "convolutional sequence to sequence",
            "bleu": "bilingual evaluation understudy translation score",
            "stipend": "financial assistance remuneration honorarium",
            "positional encoding": "sinusoidal position embeddings order",
            "label smoothing": "regularization penalizing confident predictions",
            "d_model": "transformer embedding hidden dimension",
            "attendance": "minimum attendance 75 percent requirement",
        }
        q_lower = q.lower()
        for term, expansion in synonym_map.items():
            if term in q_lower:
                expanded_variant = f"{q} {expansion}"
                if expanded_variant not in queries:
                    queries.append(expanded_variant)
                break

        # Deduplicate while preserving order
        seen = set()
        deduped = []
        for cand in queries:
            c_norm = cand.strip().lower()
            if c_norm and c_norm not in seen:
                seen.add(c_norm)
                deduped.append(cand.strip())

        # If still only 1 query generated, provide a focused perspective query
        if len(deduped) < 2:
            deduped.append(f"{q} overview")

        return deduped[: self.num_queries + 1]

    def generate_queries(self, query: str) -> List[str]:
        if query in self._query_cache:
            return self._query_cache[query]

        if self.generator_fn:
            try:
                queries = self.generator_fn(query, self.num_queries)
                if queries:
                    if query not in queries:
                        queries = [query] + queries
                    result = queries[: self.num_queries + 1]
                    self._query_cache[query] = result
                    return result
            except Exception as exc:
                logger.warning(f"Error in multi-query LLM generator: {exc}, using heuristics.")

        result = self._generate_queries_heuristic(query)
        self._query_cache[query] = result
        return result

    def retrieve(
        self,
        query: str,
        *,
        candidate_document_ids: Optional[Sequence[str]] = None,
        filter_criteria: Optional[Dict[str, Any]] = None,
        top_k: Optional[int] = None,
    ) -> MultiQueryRetrievalResponse:
        k = top_k or self.default_top_k
        formulations = self.generate_queries(query)

        candidate_runs: List[List[SearchResult]] = []
        for q in formulations:
            dense_resp = self.base_retriever.retrieve(
                query=q,
                candidate_document_ids=candidate_document_ids,
                filter_criteria=filter_criteria,
                top_k=k,
            )
            run_items = [
                SearchResult(
                    chunk_id=r.chunk_id,
                    document_id=r.document_id,
                    text=r.text,
                    score=r.score,
                    rank=r.rank,
                    parent_id=r.parent_id,
                    source_locator=r.source_locator,
                    page=r.page,
                    section=r.section,
                    retriever_name="multi_query_dense",
                    metadata=r.metadata,
                )
                for r in dense_resp.results
            ]
            candidate_runs.append(run_items)

        fused = reciprocal_rank_fusion(candidate_runs, k=60, top_n=k)

        return MultiQueryRetrievalResponse(
            original_query=query,
            generated_queries=formulations,
            candidate_chunks=len(fused),
            returned_chunks=len(fused),
            candidate_document_ids=list(candidate_document_ids or []),
            results=fused,
        )