"""RAG answer generation and LLM-as-Judge scoring for LoCoMo evaluation.

Includes confidence-aware refusal (#1275): when all retrieved claims have
similarity scores below a configurable threshold, the system refuses to
answer rather than generating a potentially hallucinated response.
"""

from typing import List, Dict, Any, Optional
from dataclasses import dataclass

# Default confidence threshold for refusal (#1275)
DEFAULT_CONFIDENCE_THRESHOLD: float = 0.3

# Standard refusal response
REFUSAL_RESPONSE: str = "I don't have information about that."


@dataclass
class GeneratedAnswer:
    question: str
    answer: str
    retrieved_claims: List[str]
    retrieval_scores: List[float]
    refused: bool = False  # True if confidence-aware refusal triggered (#1275)


@dataclass
class JudgeResult:
    score: float  # 1.0 or 0.0
    reasoning: str
    question: str
    prediction: str
    ground_truth: str


def should_refuse(
    retrieval_scores: List[float],
    confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
) -> bool:
    """Check if the system should refuse to answer based on retrieval scores.

    Returns True when ALL top-k retrieved claims have similarity scores below
    the confidence threshold, indicating the retrieved context is likely
    irrelevant to the query.

    Args:
        retrieval_scores: Similarity scores for top-k retrieved claims.
        confidence_threshold: Minimum score required (default 0.3).

    Returns:
        True if the system should refuse (all scores below threshold).
    """
    if not retrieval_scores:
        return True  # No claims retrieved → refuse
    return all(score < confidence_threshold for score in retrieval_scores)


async def generate_answer(
    question: str,
    retrieved_claims: List[str],
    llm,
    model: Optional[str] = None,
    retrieval_scores: Optional[List[float]] = None,
    confidence_threshold: float = DEFAULT_CONFIDENCE_THRESHOLD,
) -> GeneratedAnswer:
    """Given a question and retrieved claims, generate a natural language answer.

    Confidence-aware refusal (#1275): If retrieval_scores are provided and all
    are below confidence_threshold, returns a refusal response without calling
    the LLM. If no scores are provided, proceeds normally (backward compatible).
    """
    # Confidence-aware refusal (#1275)
    scores = retrieval_scores or []
    if retrieval_scores is not None and should_refuse(scores, confidence_threshold):
        return GeneratedAnswer(
            question=question,
            answer=REFUSAL_RESPONSE,
            retrieved_claims=retrieved_claims,
            retrieval_scores=scores,
            refused=True,
        )

    context = "\n".join(f"- {c}" for c in retrieved_claims)
    prompt = f"""Based ONLY on the following memories, answer the question concisely.
If the memories don't contain enough information to answer, say "No information available."

Memories:
{context}

Question: {question}
Answer:"""

    kwargs = {}
    if model:
        kwargs["model"] = model

    resp = await llm.complete(
        system="You are a factual QA system. Answer based only on provided memories. Be concise.",
        user=prompt,
        **kwargs,
    )
    text = resp.text if hasattr(resp, "text") else str(resp)
    return GeneratedAnswer(
        question=question,
        answer=text.strip(),
        retrieved_claims=retrieved_claims,
        retrieval_scores=scores,
    )


async def judge_answer(
    question: str,
    prediction: str,
    ground_truth: str,
    llm,
    model: Optional[str] = None,
) -> JudgeResult:
    """Binary LLM-as-Judge: does the prediction correctly answer the question?

    Matches Mem0/LoCoMo paper methodology — binary 1.0 or 0.0.
    """
    prompt = f"""You are a strict factual accuracy judge. Compare the system's answer to the ground truth.

Question: {question}
Ground truth answer: {ground_truth}
System answer: {prediction}

Does the system answer correctly address the question with factually consistent information?
The system answer doesn't need to be word-for-word identical, but must convey the same factual content.

Reply with EXACTLY one word: "correct" or "incorrect"
Then on a new line, briefly explain why (one sentence)."""

    kwargs = {}
    if model:
        kwargs["model"] = model

    resp = await llm.complete(
        system="You are a strict factual judge. Output 'correct' or 'incorrect' then a brief reason.",
        user=prompt,
        **kwargs,
    )
    text = resp.text if hasattr(resp, "text") else str(resp)
    lines = text.strip().split("\n", 1)
    verdict = lines[0].strip().lower()
    reasoning = lines[1].strip() if len(lines) > 1 else ""

    # Check for "incorrect" first since it contains "correct" as substring
    if "incorrect" in verdict:
        score = 0.0
    elif "correct" in verdict:
        score = 1.0
    else:
        score = 0.0
    return JudgeResult(
        score=score,
        reasoning=reasoning,
        question=question,
        prediction=prediction,
        ground_truth=ground_truth,
    )
