import json
import random
import time
import uuid

from sqlalchemy.orm import Session

from app.models.dataset import DatasetExample, ExampleStatus
from app.models.stubs import EvaluationResult, EvaluationStatus
from app.services import ollama_client
from app.services.rag import semantic_search

JUDGE_SYSTEM_PROMPT = """You are a strict evaluator for a domain-specific AI assistant. Given a \
question, a reference answer, and a candidate answer, score the candidate.

Return ONLY a JSON object: {"accuracy": <1-5>, "groundedness": <1-5>, "hallucination": <true|false>, "notes": "<short reason>"}

- accuracy: how factually correct the candidate is relative to the reference (5 = fully correct)
- groundedness: how well the candidate sticks to information implied by the reference rather than
  inventing unrelated claims (5 = fully grounded)
- hallucination: true if the candidate states something not supported by the reference or the
  provided context
"""


def _build_answer_messages(instruction: str, input_text: str | None, context_chunks: list[dict]) -> list[dict]:
    system = "You are a helpful, accurate assistant. Answer concisely and only state what you're confident about."
    if context_chunks:
        context_block = "\n\n".join(f"[{i+1}] {c['chunk_text']}" for i, c in enumerate(context_chunks))
        system += f"\n\nUse this retrieved context if relevant, and cite it as [1], [2], etc:\n{context_block}"

    user = instruction if not input_text else f"{instruction}\n\nContext: {input_text}"
    return [{"role": "system", "content": system}, {"role": "user", "content": user}]


def _judge(instruction: str, reference: str, candidate: str, judge_model: str) -> dict:
    user = f"Question: {instruction}\n\nReference answer: {reference}\n\nCandidate answer: {candidate}"
    result = ollama_client.chat(
        judge_model,
        [{"role": "system", "content": JUDGE_SYSTEM_PROMPT}, {"role": "user", "content": user}],
        json_mode=True,
        temperature=0,
    )
    try:
        parsed = json.loads(result["content"])
        return {
            "accuracy": float(parsed.get("accuracy", 0)),
            "groundedness": float(parsed.get("groundedness", 0)),
            "hallucination": bool(parsed.get("hallucination", False)),
            "notes": str(parsed.get("notes", "")),
        }
    except (json.JSONDecodeError, TypeError, ValueError):
        return {"accuracy": 0.0, "groundedness": 0.0, "hallucination": True, "notes": "judge response unparseable"}


def run_evaluation(
    db: Session,
    evaluation_id: uuid.UUID,
    dataset_id: uuid.UUID,
    project_id: uuid.UUID,
    base_model: str,
    judge_model: str,
    use_rag: bool,
    sample_size: int | None,
) -> None:
    evaluation = db.get(EvaluationResult, evaluation_id)
    evaluation.status = EvaluationStatus.running
    db.commit()

    examples = (
        db.query(DatasetExample)
        .filter(DatasetExample.dataset_id == dataset_id, DatasetExample.status == ExampleStatus.approved)
        .all()
    )
    if sample_size and sample_size < len(examples):
        examples = random.sample(examples, sample_size)

    if not examples:
        evaluation.status = EvaluationStatus.failed
        evaluation.metrics = {"error": "No approved examples in this dataset to evaluate against"}
        db.commit()
        return

    per_example = []
    for example in examples:
        context_chunks = []
        if use_rag:
            try:
                context_chunks = semantic_search(db, project_id, example.instruction, top_k=3)
            except Exception:  # noqa: BLE001
                context_chunks = []

        messages = _build_answer_messages(example.instruction, example.input, context_chunks)
        start = time.perf_counter()
        answer = ollama_client.chat(base_model, messages)
        latency_ms = (time.perf_counter() - start) * 1000

        judgement = _judge(example.instruction, example.output, answer["content"], judge_model)

        per_example.append(
            {
                "example_id": str(example.id),
                "instruction": example.instruction,
                "reference": example.output,
                "candidate": answer["content"],
                "latency_ms": round(latency_ms, 1),
                "output_tokens": answer["eval_count"],
                "input_tokens": answer["prompt_eval_count"],
                "used_rag": bool(context_chunks),
                **judgement,
            }
        )

    n = len(per_example)
    metrics = {
        "sample_size": n,
        "avg_accuracy": round(sum(e["accuracy"] for e in per_example) / n, 2),
        "avg_groundedness": round(sum(e["groundedness"] for e in per_example) / n, 2),
        "hallucination_rate": round(sum(1 for e in per_example if e["hallucination"]) / n, 3),
        "avg_latency_ms": round(sum(e["latency_ms"] for e in per_example) / n, 1),
        "avg_output_tokens": round(sum(e["output_tokens"] for e in per_example) / n, 1),
        "avg_input_tokens": round(sum(e["input_tokens"] for e in per_example) / n, 1),
        "rag_used": use_rag,
        "judge_model": judge_model,
        "examples": per_example,
    }

    evaluation.metrics = metrics
    evaluation.status = EvaluationStatus.completed
    db.commit()