"""
Hallucination Guard Module - NLI-based hallucination detection for RAG outputs.

Primary method: Cross-encoder NLI model (DeBERTa-v3) scores (context, answer) pairs.
Secondary method: LLM self-check via OpenRouter API when DeBERTa is unavailable.

Used AFTER LLM generation to verify answers are grounded in retrieved context.
"""

from __future__ import annotations

import asyncio
import logging
import os
import re
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any

import httpx

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

NLI_MODEL_PRIMARY = "cross-encoder/nli-deberta-v3-small"
NLI_MODEL_FALLBACK = "moritz23/nli-deberta-v3-xsmall"

OPENROUTER_API_KEY_ENV = "OPENROUTER_API_KEY"
LLM_API_KEY_ENV = "LLM_API_KEY"
# Decode base64 LLM key if present, otherwise use plain LLM_API_KEY
# (safety net: ensures key is decoded even when imported without main.py)
if os.getenv("LLM_API_KEY_B64"):
    import base64 as _b64

    os.environ["LLM_API_KEY"] = _b64.b64decode(os.getenv("LLM_API_KEY_B64")).decode()
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://api.deepseek.com/v1/chat/completions")
LLM_MODEL = os.getenv("LLM_MODEL", "deepseek-v4-flash")
OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
OPENROUTER_MODEL = "openai/gpt-4.1-mini"

MAX_CHECK_TIMEOUT_S = 5.0

# NLI label indices (standard for cross-encoder NLI models)
LABEL_ENTAILMENT = 0
LABEL_NEUTRAL = 1
LABEL_CONTRADICTION = 2

# Confidence thresholds
CONFIDENCE_HIGH = 0.85
CONFIDENCE_MEDIUM = 0.60


class ConfidenceLevel(StrEnum):
    HIGH = "high"
    MEDIUM = "medium"
    LOW = "low"


# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------


@dataclass
class CheckResult:
    """Result of a single hallucination check."""

    is_faithful: bool
    confidence: float
    method: str
    contradiction_details: list[str] = field(default_factory=list)
    flagged_claims: list[str] = field(default_factory=list)

    @property
    def confidence_level(self) -> ConfidenceLevel:
        if self.confidence > CONFIDENCE_HIGH:
            return ConfidenceLevel.HIGH
        elif self.confidence >= CONFIDENCE_MEDIUM:
            return ConfidenceLevel.MEDIUM
        return ConfidenceLevel.LOW


@dataclass
class CitationVerificationResult:
    """Result of verifying inline citations against sources."""

    total_citations: int
    verified_citations: int
    unsupported_citations: list[dict[str, Any]] = field(default_factory=list)
    citation_details: list[dict[str, Any]] = field(default_factory=list)
    all_supported: bool = True


# ---------------------------------------------------------------------------
# Claim splitter - simple sentence-level extraction
# ---------------------------------------------------------------------------

_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")


def _split_claims(text: str) -> list[str]:
    """Split answer into individual claim sentences for granular checking."""
    sentences = _SENTENCE_RE.split(text.strip())
    return [s.strip() for s in sentences if s.strip()]


# ---------------------------------------------------------------------------
# Citation parser
# ---------------------------------------------------------------------------

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


def _extract_citations(text: str) -> list[tuple[int, str]]:
    """
    Parse inline citations like [1], [2] from the answer.
    Returns list of (citation_number, sentence_containing_citation).
    """
    results: list[tuple[int, str]] = []
    sentences = _split_claims(text)
    for sentence in sentences:
        nums = _CITATION_RE.findall(sentence)
        for n in nums:
            results.append((int(n), sentence))
    return results


# ---------------------------------------------------------------------------
# HallucinationGuard
# ---------------------------------------------------------------------------


class HallucinationGuard:
    """NLI-based hallucination detector with LLM self-check fallback."""

    _instance: HallucinationGuard | None = None

    def __init__(
        self,
        nli_model_name: str = NLI_MODEL_PRIMARY,
        fallback_model_name: str = NLI_MODEL_FALLBACK,
        timeout_s: float = MAX_CHECK_TIMEOUT_S,
    ) -> None:
        self._nli_model_name = nli_model_name
        self._fallback_model_name = fallback_model_name
        self._timeout_s = timeout_s

        # Lazy-loaded model artefacts
        self._tokenizer = None
        self._model = None
        self._model_loaded = False
        self._model_load_failed = False
        self._active_model_name: str | None = None

    # -----------------------------------------------------------------------
    # Singleton
    # -----------------------------------------------------------------------

    @classmethod
    async def get_guard(cls) -> HallucinationGuard:
        """Return the singleton HallucinationGuard, creating if needed."""
        if cls._instance is None:
            cls._instance = cls()
        if not cls._instance._model_loaded and not cls._instance._model_load_failed:
            await cls._instance._load_model()
        return cls._instance

    # -----------------------------------------------------------------------
    # Model loading
    # -----------------------------------------------------------------------

    async def _load_model(self) -> None:
        """Attempt to load the primary NLI model, then the fallback."""
        for model_name in (self._nli_model_name, self._fallback_model_name):
            try:
                logger.info("Loading NLI model: %s", model_name)
                from transformers import AutoModelForSequenceClassification, AutoTokenizer

                loop = asyncio.get_event_loop()
                tokenizer, model = await loop.run_in_executor(
                    None,
                    lambda mn=model_name: (
                        AutoTokenizer.from_pretrained(mn),
                        AutoModelForSequenceClassification.from_pretrained(mn),
                    ),
                )
                model.eval()
                self._tokenizer = tokenizer
                self._model = model
                self._model_loaded = True
                self._active_model_name = model_name
                logger.info("Successfully loaded NLI model: %s", model_name)
                return
            except Exception as exc:
                logger.warning("Failed to load NLI model %s: %s", model_name, exc)
                continue

        self._model_load_failed = True
        logger.warning("All NLI models failed to load. Falling back to LLM self-check.")

    # -----------------------------------------------------------------------
    # Warm-up & health
    # -----------------------------------------------------------------------

    async def warm_up(self) -> None:
        """Pre-load the NLI model so first real check is fast."""
        if not self._model_loaded and not self._model_load_failed:
            await self._load_model()
        logger.info(
            "HallucinationGuard warm-up complete. model_loaded=%s active=%s",
            self._model_loaded,
            self._active_model_name,
        )

    def health_check(self) -> dict[str, Any]:
        """Return current model status for health endpoints."""
        return {
            "nli_model_loaded": self._model_loaded,
            "nli_model_name": self._active_model_name,
            "fallback_active": self._model_load_failed,
            "primary_model": self._nli_model_name,
            "fallback_model": self._fallback_model_name,
            "method": "nli" if self._model_loaded else "llm_self_check",
        }

    # -----------------------------------------------------------------------
    # Core NLI scoring
    # -----------------------------------------------------------------------

    async def _nli_score(self, premise: str, hypothesis: str) -> dict[str, float]:
        """
        Run NLI scoring on a (premise, hypothesis) pair.
        Returns dict with keys: entailment, neutral, contradiction.
        """
        import torch

        def _sync_score() -> dict[str, float]:
            inputs = self._tokenizer(
                premise,
                hypothesis,
                truncation=True,
                max_length=512,
                return_tensors="pt",
            )
            with torch.no_grad():
                logits = self._model(**inputs).logits
            probs = torch.softmax(logits, dim=-1).squeeze().tolist()
            # Ensure probs is a list even for single-entry batches
            if isinstance(probs, float):
                probs = [probs]
            return {
                "entailment": float(probs[LABEL_ENTAILMENT]),
                "neutral": float(probs[LABEL_NEUTRAL]),
                "contradiction": float(probs[LABEL_CONTRADICTION]),
            }

        loop = asyncio.get_event_loop()
        result = await asyncio.wait_for(
            loop.run_in_executor(None, _sync_score),
            timeout=self._timeout_s,
        )
        return result

    # -----------------------------------------------------------------------
    # LLM self-check fallback
    # -----------------------------------------------------------------------

    @staticmethod
    def _build_llm_prompt(context: str, answer: str) -> str:
        return (
            "Given ONLY the following retrieved context, determine whether the answer "
            "is faithful to the context. An answer is faithful if every claim it makes "
            "is supported by the context. If the answer contradicts the context or adds "
            "information not present in the context, it is not faithful.\n\n"
            f"Context:\n{context}\n\n"
            f"Answer:\n{answer}\n\n"
            "Is this answer faithful to the context? Answer Yes or No, then provide "
            "brief reasoning.\n\nFormat:\nFaithful: Yes/No\nReasoning: ..."
        )

    async def _llm_self_check(self, answer: str, context: str) -> CheckResult:
        """Use an LLM (DeepSeek primary, OpenRouter fallback) to check faithfulness when NLI is unavailable."""
        api_key = os.environ.get(LLM_API_KEY_ENV, "") or os.environ.get(OPENROUTER_API_KEY_ENV, "")
        if not api_key:
            logger.error("LLM_API_KEY / OPENROUTER_API_KEY not set; LLM self-check unavailable.")
            return CheckResult(
                is_faithful=False,
                confidence=0.0,
                method="llm_self_check_unavailable",
                contradiction_details=["LLM API key not configured"],
                flagged_claims=[],
            )

        # Use DeepSeek URL/model if LLM_API_KEY is set, else fall back to OpenRouter
        use_llm = bool(os.environ.get(LLM_API_KEY_ENV, ""))
        api_url = LLM_BASE_URL if use_llm else OPENROUTER_API_URL
        model = LLM_MODEL if use_llm else OPENROUTER_MODEL

        prompt = self._build_llm_prompt(context, answer)
        payload = {
            "model": model,
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.0,
            "max_tokens": 256,
        }
        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        }

        try:
            async with httpx.AsyncClient(timeout=self._timeout_s) as client:
                resp = await client.post(api_url, json=payload, headers=headers)
                resp.raise_for_status()
                data = resp.json()

            content = data["choices"][0]["message"]["content"].strip()
            is_faithful = content.lower().startswith("faithful: yes") or ("faithful: yes" in content.lower()[:80])

            # Extract reasoning
            reasoning_parts = content.split("Reasoning:")
            reasoning = reasoning_parts[1].strip() if len(reasoning_parts) > 1 else content

            confidence = 0.7 if is_faithful else 0.3

            contradiction_details: list[str] = []
            flagged_claims: list[str] = []
            if not is_faithful:
                contradiction_details.append(reasoning)
                flagged_claims.append(answer[:200])

            return CheckResult(
                is_faithful=is_faithful,
                confidence=confidence,
                method="llm_self_check",
                contradiction_details=contradiction_details,
                flagged_claims=flagged_claims,
            )
        except TimeoutError:
            logger.warning("LLM self-check timed out after %ss", self._timeout_s)
            return CheckResult(
                is_faithful=False,
                confidence=0.0,
                method="llm_self_check_timeout",
                contradiction_details=["LLM self-check timed out"],
                flagged_claims=[],
            )
        except Exception as exc:
            logger.error("LLM self-check failed: %s", exc)
            return CheckResult(
                is_faithful=False,
                confidence=0.0,
                method="llm_self_check_error",
                contradiction_details=[str(exc)],
                flagged_claims=[],
            )

    # -----------------------------------------------------------------------
    # Public API - single check
    # -----------------------------------------------------------------------

    async def check_answer(self, answer: str, context: str) -> CheckResult:
        """
        Check a single (answer, context) pair for hallucinations.

        Uses NLI model if loaded, otherwise falls back to LLM self-check.
        Checks each claim sentence individually and aggregates.
        """
        if self._model_loaded:
            return await self._check_answer_nli(answer, context)
        else:
            return await self._llm_self_check(answer, context)

    async def _check_answer_nli(self, answer: str, context: str) -> CheckResult:
        """NLI-based check: score each claim sentence against the context."""
        claims = _split_claims(answer)
        if not claims:
            return CheckResult(
                is_faithful=True,
                confidence=1.0,
                method="nli",
                contradiction_details=[],
                flagged_claims=[],
            )

        flagged_claims: list[str] = []
        contradiction_details: list[str] = []
        total_entailment = 0.0
        total_neutral = 0.0
        total_contradiction = 0.0

        for claim in claims:
            try:
                scores = await asyncio.wait_for(
                    self._nli_score(context, claim),
                    timeout=self._timeout_s,
                )
            except TimeoutError:
                logger.warning("NLI scoring timed out for claim: %.80s", claim)
                flagged_claims.append(claim)
                contradiction_details.append(f"Timeout checking: {claim[:80]}")
                total_contradiction += 1.0
                continue
            except Exception as exc:
                logger.error("NLI scoring error: %s", exc)
                flagged_claims.append(claim)
                contradiction_details.append(f"Error checking: {claim[:80]}")
                total_contradiction += 1.0
                continue

            entailment = scores["entailment"]
            neutral = scores["neutral"]
            contradiction = scores["contradiction"]

            total_entailment += entailment
            total_neutral += neutral
            total_contradiction += contradiction

            if contradiction > entailment and contradiction > neutral:
                flagged_claims.append(claim)
                contradiction_details.append(f"Contradiction (p={contradiction:.2f}): {claim[:120]}")

        n = len(claims)
        avg_entailment = total_entailment / n
        avg_contradiction = total_contradiction / n

        # Faithful = no flagged contradictions. A claim is only flagged when
        # contradiction > entailment AND contradiction > neutral (line 411).
        # High neutral scores just mean "not enough info to confirm" - not a hallucination.
        is_faithful = len(flagged_claims) == 0
        confidence = (avg_entailment + (1.0 - avg_contradiction)) / 2.0 if is_faithful else 1.0 - avg_contradiction

        return CheckResult(
            is_faithful=is_faithful,
            confidence=round(confidence, 4),
            method="nli",
            contradiction_details=contradiction_details,
            flagged_claims=flagged_claims,
        )

    # -----------------------------------------------------------------------
    # Public API - multi-source check
    # -----------------------------------------------------------------------

    async def check_answer_with_sources(self, answer: str, sources: list[dict[str, Any]]) -> CheckResult:
        """
        Check an answer against each source individually and aggregate.

        Each source dict should have a 'text' or 'content' key with the
        source content. The answer is evaluated against each source; if ANY
        source entails a claim, it is considered supported.
        """
        if not sources:
            return CheckResult(
                is_faithful=False,
                confidence=0.0,
                method="nli" if self._model_loaded else "llm_self_check",
                contradiction_details=["No sources provided for verification"],
                flagged_claims=[],
            )

        claims = _split_claims(answer)
        if not claims:
            return CheckResult(
                is_faithful=True,
                confidence=1.0,
                method="nli" if self._model_loaded else "llm_self_check",
            )

        # Extract text from each source
        source_texts: list[str] = []
        for src in sources:
            text = src.get("text", src.get("content", src.get("page_content", "")))
            if text:
                source_texts.append(str(text))

        if not source_texts:
            return CheckResult(
                is_faithful=False,
                confidence=0.0,
                method="nli" if self._model_loaded else "llm_self_check",
                contradiction_details=["Sources contain no usable text"],
                flagged_claims=[],
            )

        if self._model_loaded:
            return await self._check_answer_with_sources_nli(claims, source_texts)
        else:
            # Combine sources for LLM self-check
            combined_context = "\n\n---\n\n".join(source_texts)
            return await self._llm_self_check(answer, combined_context)

    async def _check_answer_with_sources_nli(self, claims: list[str], source_texts: list[str]) -> CheckResult:
        """NLI-based multi-source check: a claim is supported if ANY source entails it."""
        flagged_claims: list[str] = []
        contradiction_details: list[str] = []

        claim_verdicts: list[float] = []  # best entailment score per claim

        for claim in claims:
            best_entailment = 0.0
            best_contradiction = 0.0

            for source in source_texts:
                try:
                    scores = await asyncio.wait_for(
                        self._nli_score(source, claim),
                        timeout=self._timeout_s,
                    )
                except TimeoutError:
                    logger.warning("NLI timeout for claim against source: %.60s", claim[:60])
                    continue
                except Exception as exc:
                    logger.error("NLI error in multi-source check: %s", exc)
                    continue

                best_entailment = max(best_entailment, scores["entailment"])
                best_contradiction = max(best_contradiction, scores["contradiction"])

            claim_verdicts.append(best_entailment)

            if best_contradiction > best_entailment and best_contradiction > 0.5:
                flagged_claims.append(claim)
                contradiction_details.append(f"Contradiction (best p={best_contradiction:.2f}): {claim[:120]}")

        n = len(claims)
        avg_entailment = sum(claim_verdicts) / n if n else 0.0
        has_contradiction = len(flagged_claims) > 0
        is_faithful = not has_contradiction
        confidence = avg_entailment if is_faithful else 1.0 - (len(flagged_claims) / n)

        return CheckResult(
            is_faithful=is_faithful,
            confidence=round(confidence, 4),
            method="nli",
            contradiction_details=contradiction_details,
            flagged_claims=flagged_claims,
        )

    # -----------------------------------------------------------------------
    # Citation verification
    # -----------------------------------------------------------------------

    async def verify_citations(self, answer: str, sources: list[dict[str, Any]]) -> CitationVerificationResult:
        """
        Parse inline citations [1], [2] from the answer and verify each cited
        source actually supports the claim it is attached to.

        Each source dict is expected to have a numeric index or be ordered
        such that sources[0] corresponds to [1], sources[1] to [2], etc.
        """
        citations = _extract_citations(answer)
        if not citations:
            return CitationVerificationResult(
                total_citations=0,
                verified_citations=0,
                unsupported_citations=[],
                citation_details=[],
                all_supported=True,
            )

        citation_details: list[dict[str, Any]] = []
        unsupported: list[dict[str, Any]] = []
        verified = 0

        for cit_num, sentence in citations:
            src_idx = cit_num - 1  # [1] -> index 0
            source_text = ""

            if 0 <= src_idx < len(sources):
                src = sources[src_idx]
                source_text = src.get("text", src.get("content", src.get("page_content", "")))
                source_text = str(source_text) if source_text else ""

            detail: dict[str, Any] = {
                "citation_number": cit_num,
                "sentence": sentence,
                "source_found": bool(source_text),
            }

            if not source_text:
                detail["supported"] = False
                detail["reason"] = "Source not found or empty"
                unsupported.append(detail)
                citation_details.append(detail)
                continue

            # Verify citation supports the claim
            if self._model_loaded:
                try:
                    scores = await asyncio.wait_for(
                        self._nli_score(source_text, sentence),
                        timeout=self._timeout_s,
                    )
                except TimeoutError:
                    detail["supported"] = False
                    detail["reason"] = "NLI timeout"
                    unsupported.append(detail)
                    citation_details.append(detail)
                    continue
                except Exception as exc:
                    detail["supported"] = False
                    detail["reason"] = f"NLI error: {exc}"
                    unsupported.append(detail)
                    citation_details.append(detail)
                    continue

                entailment = scores["entailment"]
                contradiction = scores["contradiction"]
                neutral = scores["neutral"]

                if entailment >= contradiction and entailment >= neutral:
                    detail["supported"] = True
                    detail["score"] = scores
                    verified += 1
                elif contradiction > entailment:
                    detail["supported"] = False
                    detail["score"] = scores
                    detail["reason"] = f"Contradiction detected (p={contradiction:.2f})"
                    unsupported.append(detail)
                else:
                    # Neutral - the source doesn't contradict but doesn't entail either
                    detail["supported"] = False
                    detail["score"] = scores
                    detail["reason"] = "Citation does not support claim (neutral)"
                    unsupported.append(detail)
            else:
                # Fallback: LLM self-check for citation
                try:
                    llm_result = await asyncio.wait_for(
                        self._llm_self_check(sentence, source_text),
                        timeout=self._timeout_s,
                    )
                    detail["supported"] = llm_result.is_faithful
                    detail["method"] = "llm_self_check"
                    if not llm_result.is_faithful:
                        detail["reason"] = "; ".join(llm_result.contradiction_details)
                        unsupported.append(detail)
                    else:
                        verified += 1
                except TimeoutError:
                    detail["supported"] = False
                    detail["reason"] = "LLM check timed out"
                    unsupported.append(detail)

            citation_details.append(detail)

        all_supported = len(unsupported) == 0

        return CitationVerificationResult(
            total_citations=len(citations),
            verified_citations=verified,
            unsupported_citations=unsupported,
            citation_details=citation_details,
            all_supported=all_supported,
        )

    # -----------------------------------------------------------------------
    # Batch checking
    # -----------------------------------------------------------------------

    async def check_multiple(self, answers_and_contexts: list[tuple[str, str]]) -> list[CheckResult]:
        """
        Check multiple (answer, context) pairs concurrently.

        Each call is bounded by the timeout; results are returned in order.
        """
        if not answers_and_contexts:
            return []

        tasks = [self.check_answer(answer, context) for answer, context in answers_and_contexts]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        output: list[CheckResult] = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error("Batch check %d failed: %s", i, result)
                output.append(
                    CheckResult(
                        is_faithful=False,
                        confidence=0.0,
                        method="error",
                        contradiction_details=[f"Batch check error: {result}"],
                        flagged_claims=[],
                    )
                )
            else:
                output.append(result)

        return output

    # -----------------------------------------------------------------------
    # Reset (mainly for testing)
    # -----------------------------------------------------------------------

    @classmethod
    def _reset_singleton(cls) -> None:
        """Reset the singleton instance (for testing purposes only)."""
        cls._instance = None


# ---------------------------------------------------------------------------
# Module-level convenience
# ---------------------------------------------------------------------------


async def get_guard() -> HallucinationGuard:
    """Async singleton accessor - the recommended entry point."""
    return await HallucinationGuard.get_guard()
