"""Compare one Sherlock model and prompt on a fixed set of labeled messages."""

import argparse
import hashlib
import json
import os
from pathlib import Path
import time

import mlflow
import requests

from get_token import get_token


DEFAULT_API_BASE = "https://api-sherlock.cloudferro.com/openai/v1"
EXAMPLE_DIR = Path(__file__).resolve().parent
LABELS = {"storage", "access", "compute"}


def required_environment(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        raise SystemExit(f"Set {name} in this terminal first.")
    return value


def complete(session, api_base, model, prompt, message):
    """Send only the instructions and input text to Sherlock."""
    started = time.perf_counter()
    response = session.post(
        f"{api_base}/chat/completions",
        json={
            "model": model,
            "messages": [
                {"role": "system", "content": prompt},
                {"role": "user", "content": message},
            ],
            "temperature": 0,
            "max_tokens": 64,
        },
        timeout=(10, 90),
    )
    response.raise_for_status()
    body = response.json()
    choice = body["choices"][0]
    answer = choice["message"].get("content")
    if not isinstance(answer, str):
        raise ValueError("Sherlock returned no text answer.")
    return {
        "response": answer,
        "latency_seconds": time.perf_counter() - started,
        "reported_model": body.get("model"),
        "finish_reason": choice.get("finish_reason"),
        "usage": body.get("usage"),
    }


def load_inputs(dataset_path, prompts_path, prompt_name):
    dataset_bytes = dataset_path.read_bytes()
    cases = json.loads(dataset_bytes)
    prompts = json.loads(prompts_path.read_text(encoding="utf-8"))
    prompt = prompts[prompt_name]
    if not isinstance(prompt, str) or not prompt.strip():
        raise ValueError("The selected prompt must be nonempty text.")
    if not isinstance(cases, list) or not cases:
        raise ValueError("The dataset must contain at least one case.")
    ids = set()
    for case in cases:
        if (
            not isinstance(case, dict)
            or not isinstance(case.get("id"), str)
            or not case["id"].strip()
            or case["id"] in ids
            or not isinstance(case.get("input"), str)
            or not case["input"].strip()
            or case.get("expected") not in LABELS
        ):
            raise ValueError("Each case needs a unique id, input text, and a valid expected label.")
        ids.add(case["id"])
    return cases, prompt, hashlib.sha256(dataset_bytes).hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--list-models", action="store_true")
    parser.add_argument("--model", help="Exact chat model ID returned by Sherlock.")
    parser.add_argument("--prompt", choices=["baseline", "guided"], default="baseline")
    parser.add_argument("--run-name")
    parser.add_argument("--experiment", default="sherlock-model-prompt-comparison")
    parser.add_argument("--dataset", type=Path, default=EXAMPLE_DIR / "support_messages.json")
    parser.add_argument("--prompts", type=Path, default=EXAMPLE_DIR / "prompts.json")
    args = parser.parse_args()
    if not args.list_models and (not args.model or not args.model.strip()):
        parser.error("Supply --model with an exact Sherlock model ID, or use --list-models.")

    sherlock_key = required_environment("SHERLOCK_API_KEY")
    api_base = (os.environ.get("SHERLOCK_BASE_URL", "").strip() or DEFAULT_API_BASE).rstrip("/")
    with requests.Session() as session:
        # This key belongs to Sherlock. It is never used as the MLflow token.
        session.headers["Authorization"] = f"Bearer {sherlock_key}"
        if args.list_models:
            response = session.get(f"{api_base}/models", timeout=(10, 30))
            response.raise_for_status()
            for model in response.json()["data"]:
                print(model["id"])
            return

        cases, prompt, dataset_hash = load_inputs(args.dataset, args.prompts, args.prompt)
        mlflow.set_tracking_uri(required_environment("MLFLOW_TRACKING_URI"))
        os.environ["MLFLOW_TRACKING_TOKEN"] = get_token()
        mlflow.set_experiment(args.experiment)
        run_name = args.run_name or f"{args.model.rsplit('/', 1)[-1]}-{args.prompt}"

        with mlflow.start_run(run_name=run_name) as run:
            mlflow.log_params({
                "provider": "Sherlock",
                "model_id": args.model,
                "api_base": api_base,
                "prompt_variant": args.prompt,
                "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(),
                "dataset_sha256": dataset_hash,
                "case_count": len(cases),
                "temperature": 0,
                "max_tokens": 64,
            })
            mlflow.log_text(prompt, "prompt.txt")
            mlflow.log_artifact(str(args.dataset), artifact_path="inputs")
            mlflow.log_artifact(str(args.prompts), artifact_path="inputs")
            mlflow.log_artifact(str(Path(__file__)), artifact_path="source")

            results = []
            try:
                for case in cases:
                    result = complete(session, api_base, args.model, prompt, case["input"])
                    prediction = result["response"].strip().lower()
                    result.update({
                        "id": case["id"],
                        "input": case["input"],
                        "expected": case["expected"],
                        "prediction": prediction,
                        "valid_format": prediction in LABELS,
                        "correct": prediction == case["expected"],
                    })
                    results.append(result)
            except Exception as error:
                # Record the failure type, without logging credentials or request headers.
                mlflow.log_text(
                    json.dumps({"case_id": case["id"], "error_type": type(error).__name__}),
                    "error.json",
                )
                raise
            finally:
                mlflow.log_text(json.dumps(results, indent=2, ensure_ascii=False), "results.json")
                mlflow.log_metric("cases_completed", len(results))

            correct_count = sum(result["correct"] for result in results)
            mlflow.log_metrics({
                "accuracy": correct_count / len(cases),
                "valid_label_rate": sum(result["valid_format"] for result in results) / len(cases),
                "mean_latency_seconds": sum(result["latency_seconds"] for result in results) / len(cases),
            })

    print(f"Correct labels: {correct_count}/{len(cases)}")
    print(f"Accuracy: {correct_count / len(cases):.3f}")
    print(f"Run ID: {run.info.run_id}")


if __name__ == "__main__":
    main()