import argparse
import json
import re
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

from benchmark.adapters import ADAPTERS, AdapterError, InvalidPrediction
from benchmark.models import PRESETS, load_dotenv
from benchmark.dataset import (
    BENCHMARK_NAME,
    BENCHMARK_VERSION,
    VERSIONS,
    display_path,
    load_cases,
    load_tasks,
    results_dir,
    sha256_file,
    version_paths,
)

OPTIONAL_CASE_FIELDS = ("domain", "group_id", "answer_position", "phenomena")
from benchmark.metrics import evaluate


def make_adapter(args):
    try:
        return ADAPTERS[args.provider].from_args(args)
    except AdapterError as exc:
        raise SystemExit(str(exc)) from exc


def slug(text):
    return re.sub(r"[^A-Za-z0-9._-]+", "_", text)


def default_output(args, model):
    name = f"{args.provider}__{slug(model)}"
    base = results_dir(args.version)
    if args.split != "public":
        base = base / args.split
    filters = [f for f in (args.category, args.difficulty) if f]
    if not filters and not args.limit:
        return base / f"{name}.json"
    if args.limit:
        filters.append(f"limit{args.limit}")
    return base / "partial" / f"{name}__{'_'.join(map(slug, filters))}.json"


def run_case(adapter, case, question, include_state=True):
    record = {
        "id": case["id"],
        "category": case["category"],
        "difficulty": case["difficulty"],
        "task_id": case["task_id"],
        "expected": case["expected"],
        "valid_answers": case["valid_answers"],
        "scored": case.get("scored", True),
    }
    if include_state:
        record["state"] = case["state"]
    for field in OPTIONAL_CASE_FIELDS:
        if field in case:
            record[field] = case[field]

    started = time.perf_counter()
    try:
        pred = adapter.predict(case["state"], question)
    except InvalidPrediction as exc:
        record["error"] = f"InvalidPrediction: {exc}"
        record["error_type"] = "invalid_prediction"
    except AdapterError as exc:
        record["error"] = str(exc)
        record["error_type"] = "api_error"
    except Exception as exc:
        record["error"] = f"{type(exc).__name__}: {exc}"
        record["error_type"] = "exception"
    else:
        predicted = pred["choice"]
        probs = pred["probabilities"]
        record.update({
            "predicted": predicted,
            "probabilities": probs,
            "predicted_probability": probs.get(predicted),
            "expected_probability": (
                probs.get(case["expected"]) if case["expected"] is not None else None
            ),
            "correct": predicted in case["valid_answers"],
        })
    record["latency_ms"] = round((time.perf_counter() - started) * 1000, 1)
    return record


def print_record(i, record):
    if record.get("error"):
        print(f"ERR  | {i:03d} | {record['id']:10} | {record['error']}")
        return
    status = "PASS" if record["correct"] else ("INFO" if not record["scored"] else "FAIL")
    p = record["predicted_probability"]
    p_text = f"{p:.4f}" if isinstance(p, (int, float)) else "-"
    print(
        f"{status:4} | {i:03d} | {record['id']:10} | "
        f"{record['category']:20} | "
        f"expected={str(record['expected']):22} | "
        f"predicted={str(record['predicted']):22} | p={p_text}"
    )


def print_summary(report):
    print("\n" + "=" * 100)
    print("SUMMARY")
    print("=" * 100)
    print(
        f"Accuracy                  : {report['accuracy']:.2%} "
        f"({report['correct']}/{report['scored_total']})"
    )
    print(f"High-confidence error rate: {report['high_confidence_error_rate']:.2%}")
    print(f"High-confidence errors    : {report['high_confidence_errors']}")
    print(f"API / invalid errors      : {report['api_errors']}")
    if report["mean_expected_probability"] is not None:
        print(f"Mean expected probability : {report['mean_expected_probability']:.4f}")
    if report.get("group_total"):
        print(
            f"Group consistency         : {report['group_consistency']:.2%} "
            f"({report['group_total']} contrast groups)"
        )

    sections = (
        ("By category", "by_category"),
        ("By difficulty", "by_difficulty"),
        ("By domain", "by_domain"),
    )
    for title, key in sections:
        if key not in report:
            continue
        print(f"\n{title}:")
        for name, stats in report[key].items():
            print(
                f"  {name:22} {stats['correct']:3}/{stats['total']:3}  "
                f"{stats['accuracy']:.2%}"
            )


def main():
    parser = argparse.ArgumentParser(description=BENCHMARK_NAME)
    target = parser.add_mutually_exclusive_group(required=True)
    target.add_argument("--provider", choices=sorted(ADAPTERS))
    target.add_argument("--preset", choices=sorted(PRESETS), help="Known provider/model pair.")
    target.add_argument("--list-presets", action="store_true")
    parser.add_argument("--model", default=None)
    parser.add_argument("--api-key", default=None, help="Prefer the provider env variable.")
    parser.add_argument("--version", default=BENCHMARK_VERSION, choices=sorted(VERSIONS))
    parser.add_argument("--split", default="public", help="public / hidden (v0.2+)")
    parser.add_argument("--dataset", default=None, help="Override the version's split file.")
    parser.add_argument("--tasks", default=None, help="Override the version's tasks file.")
    parser.add_argument("--category", default=None)
    parser.add_argument("--difficulty", default=None)
    parser.add_argument("--limit", type=int, default=None, help="Only run the first N cases.")
    parser.add_argument(
        "--scored-only",
        action="store_true",
        help="Skip ambiguous/unscored controls (they are included by default).",
    )
    parser.add_argument("--output", default=None)
    parser.add_argument(
        "--retry-errors",
        default=None,
        metavar="RESULT",
        help="Rerun only the failed cases of this result file and merge them back in.",
    )
    parser.add_argument("--timeout", type=float, default=60)
    parser.add_argument("--max-retries", type=int, default=3)
    parser.add_argument("--delay", type=float, default=0, help="Seconds to wait between cases.")
    parser.add_argument("--laya-max-len", type=int, default=None)
    parser.add_argument("--julia-path", default=None, help="Default: models/Julia-1")
    parser.add_argument("--device", default="cpu", help="Julia device (cpu / cuda).")
    args = parser.parse_args()

    if args.list_presets:
        for name, (provider, model) in PRESETS.items():
            print(f"{name:22} --provider {provider:16} --model {model}")
        return

    if args.preset:
        args.provider, preset_model = PRESETS[args.preset]
        args.model = args.model or preset_model

    load_dotenv()

    try:
        default_dataset, default_tasks = version_paths(args.version, args.split)
    except ValueError as exc:
        raise SystemExit(str(exc)) from exc
    args.dataset = args.dataset or str(default_dataset)
    args.tasks = args.tasks or str(default_tasks)
    if not Path(args.dataset).exists():
        raise SystemExit(f"Dataset bulunamadı: {display_path(args.dataset)}")

    tasks = load_tasks(args.tasks)
    cases = load_cases(args.dataset)

    if args.category:
        cases = [c for c in cases if c["category"] == args.category]
    if args.difficulty:
        cases = [c for c in cases if c["difficulty"] == args.difficulty]
    if args.scored_only:
        cases = [c for c in cases if c.get("scored", True)]
    if args.limit:
        cases = cases[: args.limit]
    if not cases:
        raise SystemExit("Filtrelerden sonra çalıştırılacak case kalmadı.")

    prior = None
    if args.retry_errors:
        prior_path = Path(args.retry_errors)
        prior = json.loads(prior_path.read_text(encoding="utf-8"))
        if prior.get("dataset_sha256") != sha256_file(args.dataset) \
                or prior.get("tasks_sha256") != sha256_file(args.tasks):
            raise SystemExit("Sonuç dosyası farklı bir dataset veya tasks dosyasına ait.")
        failed = {r["id"] for r in prior["results"] if r.get("error")}
        if not failed:
            raise SystemExit("Yeniden denenecek hatalı case yok.")
        cases = [c for c in cases if c["id"] in failed]
        print(f"Retrying {len(cases)} failed case(s) of {display_path(prior_path)}")

    adapter = make_adapter(args)

    print(f"Provider: {adapter.provider}")
    print(f"Model   : {adapter.model}")
    print(f"Cases   : {len(cases)}")
    print("-" * 110)

    results = []
    for i, case in enumerate(cases, 1):
        if args.delay and i > 1:
            time.sleep(args.delay)
        record = run_case(
            adapter, case, tasks[case["task_id"]], include_state=args.split != "hidden"
        )
        print_record(i, record)
        results.append(record)

    if prior:
        fresh = {r["id"]: r for r in results}
        results = [fresh.get(r["id"], r) for r in prior["results"]]
        still_failing = [r["id"] for r in results if r.get("error")]
        if still_failing:
            print(f"Still failing: {', '.join(still_failing)}")

    report = evaluate(results)
    print_summary(report)

    output = Path(args.retry_errors or args.output or default_output(args, adapter.model))
    output.parent.mkdir(parents=True, exist_ok=True)

    payload = {
        "benchmark": BENCHMARK_NAME,
        "benchmark_version": args.version,
        "split": args.split,
        "dataset": display_path(args.dataset),
        "dataset_sha256": sha256_file(args.dataset),
        "tasks": display_path(args.tasks),
        "tasks_sha256": sha256_file(args.tasks),
        "provider": adapter.provider,
        "model": adapter.model,
        "adapter": adapter.config(),
        "created_at_utc": datetime.now(timezone.utc).isoformat(),
        "filters": {
            "category": args.category,
            "difficulty": args.difficulty,
            "limit": args.limit,
            "include_unscored": not args.scored_only,
        },
        "report": report,
        "results": results,
    }

    output.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"\nSaved: {display_path(output)}")

    if report["api_errors"]:
        print(
            f"WARNING: {report['api_errors']} case(s) failed and were counted as incorrect.",
            file=sys.stderr,
        )


if __name__ == "__main__":
    main()