#!/usr/bin/env python3
"""Run golden evaluation queries against current and candidate Azure OpenAI deployments.

Collects inference responses, token usage, and latency to evaluate behavioural drift,
grounding regressions, and verbosity inflation prior to model cutover.
"""
from __future__ import annotations

import argparse
import json
import os
import pathlib
import re
import sys
import time
from typing import Any

DEFAULT_PROMPTY_PATH = pathlib.Path(__file__).resolve().parents[1] / "prompts" / "grounded_qa.prompty"
DEFAULT_SYSTEM_PROMPT = (
    "You are a helpful assistant. Answer truthfully using only the provided context."
)


def load_dataset(dataset_path: pathlib.Path, max_samples: int | None = None) -> list[dict[str, Any]]:
    """Load evaluation records from a JSONL file, optionally capped at max_samples."""
    if not dataset_path.exists():
        sys.stderr.write(f"Error: Dataset not found at {dataset_path}\n")
        sys.exit(1)

    records: list[dict[str, Any]] = []
    with dataset_path.open("r", encoding="utf-8") as handle:
        for line_num, raw_line in enumerate(handle, start=1):
            line = raw_line.strip()
            if not line:
                continue
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as exc:
                sys.stderr.write(f"Error: Invalid JSON on line {line_num} of {dataset_path}: {exc}\n")
                sys.exit(1)
            if max_samples is not None and len(records) >= max_samples:
                break

    if not records:
        sys.stderr.write(f"Error: Dataset at {dataset_path} contains no records.\n")
        sys.exit(1)

    return records


def extract_prompty_defaults(front_matter: str) -> dict[str, str]:
    """Extract default string input values from Prompty YAML front matter."""
    defaults = {
        "client_name": "Contoso",
        "refusal_message": "I am sorry, but I cannot find the answer to your question in the approved documentation.",
    }
    try:
        import yaml

        meta = yaml.safe_load(front_matter) or {}
        inputs = meta.get("inputs", {})
        if isinstance(inputs, dict):
            for key in ("client_name", "refusal_message"):
                entry = inputs.get(key, {})
                if isinstance(entry, dict) and entry.get("default"):
                    defaults[key] = str(entry["default"])
    except ImportError:
        pass
    return defaults


def parse_prompty_system_prompt(raw_text: str) -> str:
    """Strip Prompty YAML front matter and Jinja loop blocks to extract system instructions."""
    front_matter = ""
    body = raw_text
    fm_match = re.match(r"\A---\s*\n(.*?)\n---\s*\n(.*)\Z", raw_text, re.S)
    if fm_match:
        front_matter, body = fm_match.group(1), fm_match.group(2)

    defaults = extract_prompty_defaults(front_matter)

    # Extract system: block up to user: block if present
    sys_match = re.search(r"^system:\s*\n(.*?)(?=^user:\s*$|\Z)", body, re.S | re.M)
    system_text = sys_match.group(1) if sys_match else body

    # Remove Jinja <sources>...</sources> loop block (row context is injected in the user message)
    system_text = re.sub(r"<sources>.*?</sources>", "", system_text, flags=re.S)
    system_text = system_text.replace("{{client_name}}", defaults["client_name"])
    system_text = system_text.replace("{{ client_name }}", defaults["client_name"])
    system_text = system_text.replace("{{refusal_message}}", defaults["refusal_message"])
    system_text = system_text.replace("{{ refusal_message }}", defaults["refusal_message"])
    return system_text.strip()


def resolve_system_prompt(prompt_path: pathlib.Path | None) -> str:
    """Load a plain-text or .prompty prompt file, failing fast if an explicit path is missing."""
    target_path = prompt_path
    if target_path is not None:
        if not target_path.exists():
            sys.stderr.write(f"Error: System prompt file not found at {target_path}\n")
            sys.exit(1)
    elif DEFAULT_PROMPTY_PATH.exists():
        target_path = DEFAULT_PROMPTY_PATH
    else:
        sys.stderr.write(
            f"Notice: {DEFAULT_PROMPTY_PATH} not found and --system-prompt omitted; "
            "using built-in default system prompt.\n"
        )
        return DEFAULT_SYSTEM_PROMPT

    raw_text = target_path.read_text(encoding="utf-8")
    if target_path.suffix.lower() == ".prompty":
        return parse_prompty_system_prompt(raw_text)
    return raw_text.strip()


def build_azure_client(endpoint: str | None, api_version: str) -> Any:
    """Initialise an AzureOpenAI client using Microsoft Entra ID credentials."""
    if not endpoint:
        sys.stderr.write(
            "Error: Azure OpenAI endpoint is required for live runs. "
            "Set AZURE_OPENAI_ENDPOINT, pass --endpoint, or pass --dry-run for offline testing.\n"
        )
        sys.exit(1)

    try:
        from azure.identity import DefaultAzureCredential, get_bearer_token_provider
        from openai import AzureOpenAI
    except ImportError as exc:
        sys.stderr.write(
            f"Error: Missing required package ({exc}). "
            "Run 'pip install openai azure-identity' or pass --dry-run for offline testing.\n"
        )
        sys.exit(1)

    token_provider = get_bearer_token_provider(
        DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
    )
    return AzureOpenAI(
        azure_endpoint=endpoint,
        azure_ad_token_provider=token_provider,
        api_version=api_version,
    )


def evaluate_deployment(
    client: Any,
    deployment_name: str,
    records: list[dict[str, Any]],
    system_prompt: str,
    dry_run: bool = False,
) -> list[dict[str, Any]]:
    """Query a deployment across all records and record outputs and token metrics."""
    results: list[dict[str, Any]] = []

    for idx, row in enumerate(records):
        query = row.get("query", row.get("question", row.get("question_text", "")))
        context = row.get("context", row.get("content", row.get("source_passage", "")))

        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": f"<sources>\n{context}\n</sources>\n\nQuestion: {query}"},
        ]

        if dry_run:
            response_text = f"[Simulated answer from {deployment_name}] based on context."
            completion_tokens = 45 + (idx % 10) * 5
            prompt_tokens = 120
            latency_ms = 420.0
        else:
            start_time = time.perf_counter()
            completion = client.chat.completions.create(
                model=deployment_name,
                messages=messages,
                temperature=0.0,
            )
            latency_ms = round((time.perf_counter() - start_time) * 1000.0, 2)
            choice = completion.choices[0]
            response_text = choice.message.content or ""
            usage = completion.usage
            completion_tokens = usage.completion_tokens if usage else 0
            prompt_tokens = usage.prompt_tokens if usage else 0

        res_record = dict(row)
        # Ensure standard evaluator fields exist regardless of input key names
        res_record["query"] = query
        res_record["context"] = context
        res_record["deployment"] = deployment_name
        res_record["response"] = response_text
        res_record["completion_tokens"] = completion_tokens
        res_record["prompt_tokens"] = prompt_tokens
        res_record["latency_ms"] = latency_ms
        results.append(res_record)

    return results


def save_results(output_file: pathlib.Path, records: list[dict[str, Any]]) -> None:
    """Save evaluation results to JSONL."""
    output_file.parent.mkdir(parents=True, exist_ok=True)
    with output_file.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def compute_metrics(
    current_results: list[dict[str, Any]],
    candidate_results: list[dict[str, Any]],
) -> dict[str, Any]:
    """Compute comparative token and latency metrics between current and candidate."""
    curr_tokens = [r["completion_tokens"] for r in current_results]
    cand_tokens = [r["completion_tokens"] for r in candidate_results]
    curr_lat = [r["latency_ms"] for r in current_results]
    cand_lat = [r["latency_ms"] for r in candidate_results]

    avg_curr_tokens = sum(curr_tokens) / len(curr_tokens) if curr_tokens else 0.0
    avg_cand_tokens = sum(cand_tokens) / len(cand_tokens) if cand_tokens else 0.0
    avg_curr_lat = sum(curr_lat) / len(curr_lat) if curr_lat else 0.0
    avg_cand_lat = sum(cand_lat) / len(cand_lat) if cand_lat else 0.0

    token_ratio = round(avg_cand_tokens / avg_curr_tokens, 3) if avg_curr_tokens else 1.0
    latency_ratio = round(avg_cand_lat / avg_curr_lat, 3) if avg_curr_lat else 1.0

    return {
        "sample_count": len(current_results),
        "current_avg_completion_tokens": round(avg_curr_tokens, 1),
        "candidate_avg_completion_tokens": round(avg_cand_tokens, 1),
        "token_ratio_candidate_to_current": token_ratio,
        "current_avg_latency_ms": round(avg_curr_lat, 1),
        "candidate_avg_latency_ms": round(avg_cand_lat, 1),
        "latency_ratio_candidate_to_current": latency_ratio,
        "token_inflation_warning": token_ratio > 1.25,
    }


def parse_args() -> argparse.Namespace:
    """Parse command-line options for deployment comparison."""
    parser = argparse.ArgumentParser(
        description="Compare inference outputs and token usage across two Azure OpenAI deployments."
    )
    parser.add_argument(
        "--dataset",
        "-d",
        type=pathlib.Path,
        required=True,
        help="Path to the golden evaluation set JSONL file (e.g. data/golden/A-policies/eval_set.jsonl)",
    )
    parser.add_argument(
        "--current-deployment",
        type=str,
        default=os.environ.get("CURRENT_DEPLOYMENT", "gpt-4o-prod-v1"),
        help="Current model deployment name (default: env CURRENT_DEPLOYMENT or gpt-4o-prod-v1)",
    )
    parser.add_argument(
        "--candidate-deployment",
        type=str,
        default=os.environ.get("CANDIDATE_DEPLOYMENT", "gpt-4o-candidate-v2"),
        help="Candidate model deployment name (default: env CANDIDATE_DEPLOYMENT or gpt-4o-candidate-v2)",
    )
    parser.add_argument(
        "--system-prompt",
        "-p",
        type=pathlib.Path,
        default=None,
        help="Optional path to .prompty or text prompt file (default: snippets/prompts/grounded_qa.prompty)",
    )
    parser.add_argument(
        "--output-dir",
        "-o",
        type=pathlib.Path,
        default=pathlib.Path("eval"),
        help="Directory to save evaluation answer files and comparison summary (default: eval)",
    )
    parser.add_argument(
        "--endpoint",
        type=str,
        default=os.environ.get("AZURE_OPENAI_ENDPOINT"),
        help="Azure OpenAI endpoint URL (default: env AZURE_OPENAI_ENDPOINT)",
    )
    # verify: confirm current Azure OpenAI data-plane inference API version in Microsoft Learn
    parser.add_argument(
        "--api-version",
        type=str,
        default=os.environ.get("AZURE_OPENAI_API_VERSION", "2024-10-21"),
        help="Azure OpenAI inference API version (default: env AZURE_OPENAI_API_VERSION or 2024-10-21)",
    )
    parser.add_argument(
        "--max-samples",
        "-n",
        type=int,
        default=None,
        help="Optional maximum number of dataset rows to run (useful for quick smoke tests)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Run in offline simulation mode without calling Azure OpenAI endpoints.",
    )
    return parser.parse_args()


def main() -> int:
    """Execute comparative evaluation across two deployments."""
    args = parse_args()
    if args.max_samples is not None and args.max_samples <= 0:
        sys.stderr.write("Error: --max-samples must be a positive integer.\n")
        return 1

    system_prompt = resolve_system_prompt(args.system_prompt)
    records = load_dataset(args.dataset, max_samples=args.max_samples)
    print(f"Loaded {len(records)} test queries from {args.dataset}")

    client = None if args.dry_run else build_azure_client(args.endpoint, args.api_version)

    print(f"Querying current deployment: {args.current_deployment}")
    curr_results = evaluate_deployment(client, args.current_deployment, records, system_prompt, args.dry_run)
    curr_output = args.output_dir / f"answers_{args.current_deployment}.jsonl"
    save_results(curr_output, curr_results)
    print(f"Saved {len(curr_results)} answers to {curr_output}")

    print(f"Querying candidate deployment: {args.candidate_deployment}")
    cand_results = evaluate_deployment(client, args.candidate_deployment, records, system_prompt, args.dry_run)
    cand_output = args.output_dir / f"answers_{args.candidate_deployment}.jsonl"
    save_results(cand_output, cand_results)
    print(f"Saved {len(cand_results)} answers to {cand_output}")

    args.output_dir.mkdir(parents=True, exist_ok=True)
    metrics = compute_metrics(curr_results, cand_results)
    metrics_path = args.output_dir / "deployment_comparison.json"
    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")

    print("\n--- Comparative Metrics Summary ---")
    print(f"Samples: {metrics['sample_count']}")
    print(f"Current Avg Completion Tokens: {metrics['current_avg_completion_tokens']}")
    print(
        f"Candidate Avg Completion Tokens: {metrics['candidate_avg_completion_tokens']} "
        f"(Ratio: {metrics['token_ratio_candidate_to_current']})"
    )
    print(f"Current Avg Latency: {metrics['current_avg_latency_ms']} ms")
    print(
        f"Candidate Avg Latency: {metrics['candidate_avg_latency_ms']} ms "
        f"(Ratio: {metrics['latency_ratio_candidate_to_current']})"
    )

    if metrics["token_inflation_warning"]:
        print("WARNING: Candidate model generates > 1.25x completion tokens. Refine prompt length instructions.")

    return 0


if __name__ == "__main__":
    sys.exit(main())