"""Run RAG quality evaluation with the Azure AI Evaluation SDK.

Evaluates an input dataset against three core RAG metrics on a 1 to 5 scale:
- Groundedness: assesses whether the response is fully derived from retrieved context.
- Relevance: assesses whether the response directly addresses the user query.
- Retrieval: assesses whether retrieved context contains the facts needed to answer.

Usage:
    python snippets/eval_release.py \\
        [--input eval/cases.jsonl] \\
        [--limit 5] \\
        [--threshold 3.0] \\
        [--min-pass-rate 0.9] \\
        [--is-reasoning-model] \\
        [--baseline baselines/previous_release.json] \\
        [--output reports/results.json]

Environment:
    AZURE_OPENAI_ENDPOINT     Endpoint of the Azure OpenAI judge resource.
    JUDGE_DEPLOYMENT          Deployment name of the judge model (e.g. gpt-4o).
    AZURE_OPENAI_API_VERSION  Optional API version (e.g. 2024-10-21).
    AZURE_AI_PROJECT          Optional Azure AI Foundry project string (subscription/resource-group/project).

Sign-in uses DefaultAzureCredential. No API keys are read or stored.
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import sys
from pathlib import Path
from typing import Any, NoReturn

from azure.ai.evaluation import (
    AzureOpenAIModelConfiguration,
    GroundednessEvaluator,
    RelevanceEvaluator,
    RetrievalEvaluator,
    evaluate,
)
from azure.identity import DefaultAzureCredential

logging.basicConfig(level=logging.INFO, format="%(message)s")
log = logging.getLogger("eval_release")

DEFAULT_THRESHOLD = 3.0
EVALUATOR_METRICS = {
    "groundedness": {
        "avg_key": "groundedness.groundedness",
        "row_key": "outputs.groundedness.groundedness",
        "reason_key": "outputs.groundedness.groundedness_reason",
    },
    "relevance": {
        "avg_key": "relevance.relevance",
        "row_key": "outputs.relevance.relevance",
        "reason_key": "outputs.relevance.relevance_reason",
    },
    "retrieval": {
        "avg_key": "retrieval.retrieval",
        "row_key": "outputs.retrieval.retrieval",
        "reason_key": "outputs.retrieval.retrieval_reason",
    },
}


def fail(message: str, exit_code: int = 2) -> NoReturn:
    """Log an error and terminate execution."""
    log.error("ERROR: %s", message)
    sys.exit(exit_code)


def parse_args() -> argparse.Namespace:
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Run RAG quality evaluation with the Azure AI Evaluation SDK."
    )
    parser.add_argument(
        "--input",
        type=Path,
        default=Path("eval/cases.jsonl"),
        help="Path to evaluation dataset JSONL file (default: eval/cases.jsonl).",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        default=DEFAULT_THRESHOLD,
        help="Minimum required average score on 1-5 scale (default: 3.0).",
    )
    parser.add_argument(
        "--min-pass-rate",
        type=float,
        default=None,
        help="Minimum required share of rows meeting or exceeding threshold (0.0 to 1.0).",
    )
    parser.add_argument(
        "--baseline",
        type=Path,
        default=None,
        help="Optional path to previous release results.json for delta comparison.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=Path("reports/results.json"),
        help="Path to output evaluation results JSON (default: reports/results.json).",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Optional limit on the number of rows to evaluate (useful for smoke tests).",
    )
    parser.add_argument(
        "--is-reasoning-model",
        action="store_true",
        help="Set when the judge model is an o-series reasoning model.",
    )
    return parser.parse_args()


def load_baseline(baseline_path: Path) -> dict[str, float]:
    """Load baseline average metrics from a previous evaluation report."""
    if not baseline_path.exists():
        log.warning("Baseline file '%s' not found. Skipping baseline delta comparison.", baseline_path)
        return {}
    try:
        with open(baseline_path, encoding="utf-8") as f:
            data = json.load(f)
        return data.get("metrics", {})
    except Exception as exc:
        log.warning("Failed to parse baseline file '%s': %s", baseline_path, exc)
        return {}


def main() -> None:
    args = parse_args()

    endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
    deployment = os.environ.get("JUDGE_DEPLOYMENT")
    api_version = os.environ.get("AZURE_OPENAI_API_VERSION")

    if not endpoint or not deployment:
        fail("Set AZURE_OPENAI_ENDPOINT and JUDGE_DEPLOYMENT environment variables before running.")

    if not args.input.exists():
        fail(f"Input cases dataset '{args.input}' not found. Generate evaluation cases first.")

    log.info("Initialising Azure OpenAI judge model (deployment: %s)...", deployment)
    model_config = AzureOpenAIModelConfiguration(
        azure_endpoint=endpoint,
        azure_deployment=deployment,
        api_version=api_version,
        credential=DefaultAzureCredential(),
    )

    evaluators = {
        "groundedness": GroundednessEvaluator(
            model_config=model_config,
            is_reasoning_model=args.is_reasoning_model,
        ),
        "relevance": RelevanceEvaluator(
            model_config=model_config,
            is_reasoning_model=args.is_reasoning_model,
        ),
        "retrieval": RetrievalEvaluator(
            model_config=model_config,
            is_reasoning_model=args.is_reasoning_model,
        ),
    }

    evaluator_config = {
        "groundedness": {
            "column_mapping": {
                "query": "${data.query}",
                "context": "${data.context}",
                "response": "${data.response}",
            }
        },
        "relevance": {
            "column_mapping": {
                "query": "${data.query}",
                "response": "${data.response}",
            }
        },
        "retrieval": {
            "column_mapping": {
                "query": "${data.query}",
                "context": "${data.context}",
            }
        },
    }

    project_scope = os.environ.get("AZURE_AI_PROJECT")
    azure_ai_project = None
    if project_scope:
        parts = [p.strip() for p in project_scope.split("/") if p.strip()]
        if len(parts) == 3:
            azure_ai_project = {
                "subscription_id": parts[0],
                "resource_group_name": parts[1],
                "project_name": parts[2],
            }
            log.info("Logging evaluation run to Azure AI Foundry project: %s", project_scope)

    input_file = args.input
    temp_sample_file: Path | None = None
    if args.limit is not None and args.limit > 0:
        lines: list[str] = []
        with open(args.input, "r", encoding="utf-8") as f:
            for idx, line in enumerate(f):
                if idx >= args.limit:
                    break
                if line.strip():
                    lines.append(line)
        temp_sample_file = args.input.parent / f".sample_{args.input.name}"
        temp_sample_file.write_text("".join(lines), encoding="utf-8")
        input_file = temp_sample_file
        log.info("Limiting evaluation to first %d records from '%s'.", len(lines), args.input)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    log.info("Evaluating dataset '%s' across Groundedness, Relevance, and Retrieval...", input_file)

    evaluate_kwargs: dict[str, Any] = {
        "data": str(input_file),
        "evaluators": evaluators,
        "evaluator_config": evaluator_config,
        "output_path": str(args.output),
    }
    if azure_ai_project:
        evaluate_kwargs["azure_ai_project"] = azure_ai_project

    try:
        result = evaluate(**evaluate_kwargs)
    finally:
        if temp_sample_file and temp_sample_file.exists():
            temp_sample_file.unlink(missing_ok=True)

    metrics = result.get("metrics", {})
    rows = result.get("rows", [])
    total_rows = len(rows)

    log.info("\n================ EVALUATION SUMMARY ================")
    log.info("Dataset: %s (%d records)", args.input, total_rows)

    baseline_metrics = load_baseline(args.baseline) if args.baseline else {}
    failures: list[str] = []
    weak_rows_by_metric: dict[str, list[dict[str, Any]]] = {m: [] for m in EVALUATOR_METRICS}

    for name, keys in EVALUATOR_METRICS.items():
        avg_score = metrics.get(keys["avg_key"])
        if avg_score is None:
            # Look up exact key name only; do not fall back to loose prefix matching
            # which could erroneously match threshold or configuration keys.
            fail(
                f"Metric '{name}' missing from evaluation results (expected key: '{keys['avg_key']}'). "
                f"Inspect SDK version or verify metric key names.",
                exit_code=2,
            )

        row_scores = [
            float(r[keys["row_key"]])
            for r in rows
            if keys["row_key"] in r and r[keys["row_key"]] is not None
        ]

        passing_rows = sum(1 for s in row_scores if s >= args.threshold)
        pass_rate = (passing_rows / len(row_scores)) if row_scores else 0.0

        for idx, row in enumerate(rows):
            score = row.get(keys["row_key"])
            if score is not None and float(score) < args.threshold:
                weak_rows_by_metric[name].append({
                    "row_index": idx + 1,
                    "query": row.get("inputs.query") or row.get("query", "N/A"),
                    "score": float(score),
                    "reason": row.get(keys["reason_key"], ""),
                })

        baseline_score = baseline_metrics.get(keys["avg_key"])
        delta_str = ""
        if baseline_score is not None:
            delta = avg_score - float(baseline_score)
            delta_str = f" [baseline: {float(baseline_score):.2f}, delta: {delta:+.2f}]"

        log.info(
            "- %s: average %.2f%s | pass rate: %.1f%% (%d/%d at or above %.1f)",
            name.capitalize(),
            avg_score,
            delta_str,
            pass_rate * 100,
            passing_rows,
            len(row_scores),
            args.threshold,
        )

        if avg_score < args.threshold:
            failures.append(f"{name} average {avg_score:.2f} is below threshold {args.threshold:.2f}")

        if args.min_pass_rate is not None and pass_rate < args.min_pass_rate:
            failures.append(
                f"{name} pass rate {pass_rate:.1%} is below minimum {args.min_pass_rate:.1%}"
            )

    has_weak_rows = any(weak_rows_by_metric.values())
    if has_weak_rows:
        log.info("\n--- Weak Rows Below Threshold (Score < %.1f) ---", args.threshold)
        for name, items in weak_rows_by_metric.items():
            if items:
                log.info("[%s: %d weak cases]", name.upper(), len(items))
                for item in items[:5]:  # Show up to 5 weak cases per metric
                    log.info(
                        "  Row #%d: score=%.1f | query='%s'",
                        item["row_index"],
                        item["score"],
                        item["query"],
                    )
                    if item["reason"]:
                        log.info("    Reason: %s", item["reason"][:120])
                if len(items) > 5:
                    log.info("  ... and %d more rows (see full results in '%s')", len(items) - 5, args.output)

    log.info("\nArtifact output written to: %s", args.output)
    studio_url = result.get("studio_url")
    if studio_url:
        log.info("Azure AI Foundry portal run: %s", studio_url)

    if failures:
        log.error("\nEVALUATION GATE FAILED:\n- %s", "\n- ".join(failures))
        sys.exit(1)

    log.info("\nEVALUATION GATE PASSED: All metrics meet required quality thresholds.")


if __name__ == "__main__":
    main()