"""Judge the candidate answers with the Azure AI Evaluation SDK (judge route B).

Reads runs/candidate_cases.jsonl, which export_run.py writes from the application's own search
and answer code. Each line has "query", "context" and "response". The script scores groundedness
on a 1 to 5 scale and fails the gate (exit code 1) if the average is below the threshold.

Threshold
---------
    Default 3.0, the SDK's default pass threshold on the 1 to 5 scale. Check the current default
    in the SDK documentation for the version you pin.

Usage
-----
    python eval/eval_azure_ai.py \\
        [--input runs/candidate_cases.jsonl] \\
        [--threshold 3.0] \\
        [--min-pass-rate 0.9] \\
        [--output reports/azure-eval-report.json]

Environment
-----------
    AZURE_OPENAI_ENDPOINT      endpoint of the judge resource
    JUDGE_DEPLOYMENT           name of the judge deployment
    AZURE_OPENAI_API_VERSION   optional
Sign-in uses DefaultAzureCredential (the Azure CLI session in CI). No API key is read.
"""

from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path
from typing import NoReturn

from azure.ai.evaluation import AzureOpenAIModelConfiguration, GroundednessEvaluator, evaluate
from azure.identity import DefaultAzureCredential

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

DEFAULT_THRESHOLD = 3.0
ROW_SCORE = "outputs.groundedness.groundedness"
AVERAGE_SCORE = "groundedness.groundedness"


def fail(message: str) -> NoReturn:
    log.error("SETUP ERROR: %s", message)
    sys.exit(2)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Groundedness gate with the Azure AI Evaluation SDK.")
    parser.add_argument("--input", type=Path, default=Path("runs/candidate_cases.jsonl"))
    parser.add_argument("--threshold", type=float, default=DEFAULT_THRESHOLD)
    parser.add_argument(
        "--min-pass-rate",
        type=float,
        help="also fail if the share of rows at or above the threshold is below this (0 to 1)",
    )
    parser.add_argument("--output", type=Path, default=Path("reports/azure-eval-report.json"))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
    deployment = os.environ.get("JUDGE_DEPLOYMENT")
    if not endpoint or not deployment:
        fail("Set AZURE_OPENAI_ENDPOINT and JUDGE_DEPLOYMENT in the pipeline variables.")
    if not args.input.exists():
        fail(f"'{args.input}' not found. Run eval/export_run.py with --answer first.")

    model_config = AzureOpenAIModelConfiguration(
        azure_endpoint=endpoint,
        azure_deployment=deployment,
        api_version=os.environ.get("AZURE_OPENAI_API_VERSION"),
        credential=DefaultAzureCredential(),
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    result = evaluate(
        data=str(args.input),
        evaluators={"groundedness": GroundednessEvaluator(model_config)},
        evaluator_config={
            "groundedness": {
                "column_mapping": {
                    "query": "${data.query}",
                    "context": "${data.context}",
                    "response": "${data.response}",
                }
            }
        },
        output_path=str(args.output),
    )

    average = result.get("metrics", {}).get(AVERAGE_SCORE)
    if average is None:
        fail(f"The result has no '{AVERAGE_SCORE}' metric. Check the SDK version and the input rows.")
    scores = [row[ROW_SCORE] for row in result.get("rows", []) if ROW_SCORE in row]
    below = sum(1 for score in scores if score < args.threshold)
    log.info("Average groundedness: %.2f (threshold %.2f). Rows below it: %d of %d.",
             average, args.threshold, below, len(scores))

    failures = []
    if average < args.threshold:
        failures.append(f"average {average:.2f} is below {args.threshold:.2f}")
    if args.min_pass_rate is not None and scores:
        pass_rate = (len(scores) - below) / len(scores)
        if pass_rate < args.min_pass_rate:
            failures.append(f"pass rate {pass_rate:.0%} is below {args.min_pass_rate:.0%}")
    if failures:
        log.error("GATE FAILED: %s. Row-level results are in '%s'.", "; ".join(failures), args.output)
        sys.exit(1)
    log.info("GATE PASSED. Row-level results are in '%s'.", args.output)


if __name__ == "__main__":
    main()