"""Evaluation helpers for workflow-vs-agent comparisons."""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
from typing import Any

import pandas as pd
from langsmith import traceable

from src.config import build_langchain_llm
from src.schemas import EvalExample, EvalResult, JudgeResult, SystemOutput
from src.utils import read_jsonl

PROMPTS_DIR = Path(__file__).resolve().parent.parent / "prompts"


@traceable(name="evals.load_eval_dataset", run_type="chain")
def load_eval_dataset(path: str | Path) -> list[EvalExample]:
    """Load and validate the JSONL evaluation dataset."""
    return [EvalExample.model_validate(row) for row in read_jsonl(path)]


def _response_mentions_missing_info(example: EvalExample, response: str) -> bool:
    """Check whether a response asks for missing critical identifiers."""
    if example.expected.intent == "billing_issue" and not example.metadata.order_id:
        return "order" in response.lower() or "missing" in response.lower()
    if example.expected.intent == "account_access" and not example.metadata.account_id:
        return "account" in response.lower() or "verify" in response.lower()
    return True


def _unnecessary_tool_call(example: EvalExample, tool_call_count: int) -> bool:
    """Flag examples where a system likely used more tools than needed."""
    input_lc = example.input.lower()
    cheap_cases = [
        "cancel my subscription",
        "what is your refund policy",
        "how long do refunds take",
    ]
    if any(case in input_lc for case in cheap_cases):
        return tool_call_count > 2
    return False


def score_prediction(example: EvalExample, prediction: SystemOutput) -> dict[str, Any]:
    """Score one structured prediction against the expected labels."""
    return {
        "intent_correct": prediction.intent == example.expected.intent,
        "route_correct": prediction.route == example.expected.route,
        "escalation_correct": prediction.should_escalate == example.expected.should_escalate,
        "response_has_customer_action": any(
            phrase in prediction.response.lower()
            for phrase in ["share", "check", "reply", "confirm", "send"]
        ),
        "response_mentions_missing_info_when_needed": _response_mentions_missing_info(
            example, prediction.response
        ),
        "tool_call_count": prediction.tool_call_count,
        "unnecessary_tool_call_flag": _unnecessary_tool_call(
            example, prediction.tool_call_count
        ),
    }


@traceable(name="evals.run_eval_examples", run_type="chain")
def run_eval_examples(
    system_name: str,
    runner: Callable[[str, dict[str, Any]], SystemOutput],
    examples: list[EvalExample],
) -> list[EvalResult]:
    """Run one system over a list of eval examples."""
    results: list[EvalResult] = []
    for example in examples:
        prediction = runner(example.input, example.metadata.model_dump())
        results.append(
            EvalResult(
                example_id=example.id,
                system_name=system_name,
                input_text=example.input,
                prediction=prediction,
                expected=example.expected,
                reference_response=example.reference_response,
                scores=score_prediction(example, prediction),
            )
        )
    return results


def aggregate_metrics(results: list[EvalResult]) -> dict[str, float | int]:
    """Aggregate deterministic evaluation metrics across examples."""
    total = len(results)
    if total == 0:
        return {}
    return {
        "n_examples": total,
        "intent_accuracy": sum(r.scores["intent_correct"] for r in results) / total,
        "route_accuracy": sum(r.scores["route_correct"] for r in results) / total,
        "escalation_accuracy": sum(r.scores["escalation_correct"] for r in results) / total,
        "avg_tool_calls": sum(r.scores["tool_call_count"] for r in results) / total,
        "unnecessary_tool_call_rate": sum(
            r.scores["unnecessary_tool_call_flag"] for r in results
        )
        / total,
    }


def build_results_dataframe(results: list[EvalResult]) -> pd.DataFrame:
    """Convert evaluation results into a DataFrame for inspection."""
    rows = []
    for result in results:
        rows.append(
            {
                "example_id": result.example_id,
                "input_text": result.input_text,
                "system_name": result.system_name,
                "expected_intent": result.expected.intent,
                "predicted_intent": result.prediction.intent,
                "expected_route": result.expected.route,
                "predicted_route": result.prediction.route,
                "expected_escalation": result.expected.should_escalate,
                "predicted_escalation": result.prediction.should_escalate,
                "tool_call_count": result.prediction.tool_call_count,
                **result.scores,
            }
        )
    return pd.DataFrame(rows)


def compare_systems(
    workflow_results: list[EvalResult],
    agent_results: list[EvalResult],
) -> pd.DataFrame:
    """Join workflow and agent result tables side by side."""
    workflow_df = build_results_dataframe(workflow_results).add_prefix("workflow_")
    agent_df = build_results_dataframe(agent_results).add_prefix("agent_")
    return workflow_df.merge(
        agent_df,
        left_on="workflow_example_id",
        right_on="agent_example_id",
    )


@traceable(name="evals.judge_responses", run_type="chain")
def judge_responses(
    results: list[EvalResult],
    *,
    sample_size: int | None = None,
) -> list[dict[str, Any]]:
    """Run the optional LLM-as-a-judge pass over evaluation results."""
    llm = build_langchain_llm(temperature=0).with_structured_output(JudgeResult)
    judge_prompt = (PROMPTS_DIR / "judge_prompt.txt").read_text()
    selected = results[:sample_size] if sample_size else results
    judged: list[dict] = []
    for result in selected:
        judge_output = llm.invoke(
            [
                ("system", judge_prompt),
                (
                    "human",
                    f"Ticket input: {result.input_text}\n"
                    f"Expected labels: {result.expected.model_dump()}\n"
                    f"Reference response: {result.reference_response}\n"
                    f"Candidate response: {result.prediction.response}\n",
                ),
            ]
        )
        judged.append(
            {
                "example_id": result.example_id,
                "system_name": result.system_name,
                **judge_output.model_dump(),
            }
        )
    return judged


@traceable(name="evals.aggregate_judge_results", run_type="chain")
def aggregate_judge_results(judged_rows: list[dict[str, Any]]) -> dict[str, float | int]:
    """Aggregate optional judge-based evaluation scores."""
    if not judged_rows:
        return {}
    total = len(judged_rows)
    return {
        "n_judged": total,
        "avg_helpfulness": sum(row["helpfulness_score"] for row in judged_rows) / total,
        "avg_correctness": sum(row["correctness_score"] for row in judged_rows) / total,
        "avg_policy_alignment": sum(row["policy_alignment_score"] for row in judged_rows) / total,
        "pass_rate": sum(row["overall_pass"] for row in judged_rows) / total,
    }
