#!/usr/bin/env python3
"""Groundedness evaluation gate script for Azure AI Evaluation SDK.

Scores candidate responses against their retrieved context with GroundednessEvaluator.
Exit codes: 0 gate passed, 1 gate failed, 2 input or judge unavailable.

The gate fails closed: if the judge model cannot be reached it exits 2.
The word-overlap fallback runs only with --mock, for local wiring tests.
Never use --mock in a release pipeline.
"""

import argparse
import json
import os
import sys
from typing import Any


def evaluate_row(
    evaluator: Any,
    query: str,
    context: str,
    response: str,
) -> tuple[float, str]:
    """Score a single (query, context, response) triplet using the evaluator."""
    if evaluator is not None:
        try:
            result = evaluator(query=query, context=context, response=response)
            score = float(result.get("groundedness", 1.0))
            reason = str(result.get("groundedness_reason", ""))
            return score, reason
        except Exception as e:
            return 1.0, f"Evaluator call error: {e}"

    # Deterministic local fallback for testing without Azure OpenAI credentials
    norm_resp = " ".join(response.lower().split())
    norm_ctx = " ".join(context.lower().split())
    if not norm_resp or "could not find" in norm_resp:
        return 5.0, "Refusal or empty response is grounded"
    words = [w for w in norm_resp.split() if len(w) >= 3]
    matches = sum(1 for w in words if w in norm_ctx)
    ratio = matches / max(1, len(words))
    # Pass if at least 25% of content tokens overlap or substring match is present
    score = 5.0 if (ratio >= 0.25 or norm_resp in norm_ctx) else 1.0
    reason = f"Mock word overlap ratio: {ratio:.2f}"
    return score, reason


def run_gate(
    input_file: str,
    output_file: str,
    threshold: float,
    min_score: float,
    mock: bool = False,
) -> int:
    """Run batch evaluation over input JSONL and enforce release threshold."""
    if not os.path.exists(input_file):
        print(f"Error: input file '{input_file}' not found", file=sys.stderr)
        return 2

    evaluator = None
    endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
    deployment = os.environ.get("JUDGE_DEPLOYMENT")

    if endpoint and deployment:
        try:
            from azure.ai.evaluation import AzureOpenAIModelConfiguration, GroundednessEvaluator
            from azure.identity import DefaultAzureCredential

            model_config = AzureOpenAIModelConfiguration(
                azure_endpoint=endpoint,
                azure_deployment=deployment,
                credential=DefaultAzureCredential(),
            )
            evaluator = GroundednessEvaluator(model_config)
            print(f"Initialised GroundednessEvaluator with deployment '{deployment}'")
        except ImportError:
            print("Notice: azure-ai-evaluation is not installed.")
        except Exception as err:
            print(f"Notice: Failed to connect to Azure AI Evaluation ({err}).")

    if evaluator is None and not mock:
        print("Error: no judge model available. Set AZURE_OPENAI_ENDPOINT and JUDGE_DEPLOYMENT, "
              "or pass --mock for a local wiring test.", file=sys.stderr)
        return 2
    if evaluator is None:
        print("WARNING: --mock set. Scores come from word overlap, not a judge model.")

    total = 0
    passed = 0
    results = []

    with open(input_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            query = row.get("query", "")
            context = row.get("context", "")
            response = row.get("response", "")

            score, reason = evaluate_row(evaluator, query, context, response)
            is_pass = score >= min_score
            total += 1
            if is_pass:
                passed += 1

            row_res = dict(row)
            row_res["groundedness_score"] = score
            row_res["groundedness_reason"] = reason
            row_res["groundedness_result"] = "pass" if is_pass else "fail"
            results.append(row_res)

    pass_rate = (passed / total) if total > 0 else 0.0

    if output_file:
        os.makedirs(os.path.dirname(os.path.abspath(output_file)), exist_ok=True)
        with open(output_file, "w", encoding="utf-8") as f:
            for r in results:
                f.write(json.dumps(r) + "\n")
        print(f"Wrote {len(results)} evaluated rows to '{output_file}'")

    print(f"Evaluation complete: {passed}/{total} passed ({pass_rate:.1%}). Required: {threshold:.1%}")

    if pass_rate < threshold:
        print(f"[FAIL] Release gate rejected: pass rate {pass_rate:.1%} is below {threshold:.1%}", file=sys.stderr)
        return 1

    print(f"[PASS] Release gate accepted: pass rate {pass_rate:.1%} meets threshold {threshold:.1%}")
    return 0


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate RAG groundedness release gate.")
    parser.add_argument("--input", required=True, help="Path to input JSONL dataset")
    parser.add_argument("--output", default="eval/eval_results.jsonl", help="Path to output evaluated JSONL")
    parser.add_argument("--threshold", type=float, default=0.95, help="Minimum pass rate required (default: 0.95)")
    parser.add_argument("--min-score", type=float, default=4.0, help="Minimum groundedness score to pass a row (default: 4.0)")
    parser.add_argument("--mock", action="store_true",
                        help="Use the word-overlap fallback when no judge is set (local tests only)")

    args = parser.parse_args()
    code = run_gate(args.input, args.output, args.threshold, args.min_score, args.mock)
    sys.exit(code)


if __name__ == "__main__":
    main()