from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Dict, List

from PIL import Image
from sklearn.metrics import accuracy_score, classification_report

from src.inference import ClipPredictor, CustomModelPredictor, OpenAIVisionPredictor


def iter_test_images(data_dir: Path, labels: List[str]) -> List[Dict[str, str]]:
    items: List[Dict[str, str]] = []
    test_root = data_dir / "test"
    for label in labels:
        class_dir = test_root / label
        if not class_dir.exists():
            continue
        for p in class_dir.iterdir():
            if p.is_file() and p.suffix.lower() in {".jpg", ".jpeg", ".png", ".webp"}:
                items.append({"path": str(p), "label": label})
    return items


def evaluate_custom(custom: CustomModelPredictor, samples: List[Dict[str, str]]) -> Dict[str, object]:
    y_true = []
    y_pred = []
    for item in samples:
        image = Image.open(item["path"])
        pred = custom.predict(image)
        y_true.append(item["label"])
        y_pred.append(pred["top_prediction"]["label"])

    acc = accuracy_score(y_true, y_pred)
    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
    return {"accuracy": float(acc), "classification_report": report}


def evaluate_clip(clip: ClipPredictor, samples: List[Dict[str, str]]) -> Dict[str, object]:
    if not clip.available():
        return {"available": False, "error": "Could not load CLIP model."}

    y_true = []
    y_pred = []
    for item in samples:
        image = Image.open(item["path"])
        pred = clip.predict(image)
        y_true.append(item["label"])
        y_pred.append(pred["top_prediction"]["label"])

    acc = accuracy_score(y_true, y_pred)
    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
    return {"available": True, "accuracy": float(acc), "classification_report": report}


def evaluate_openai(openai_model: OpenAIVisionPredictor, samples: List[Dict[str, str]], max_samples: int) -> Dict[str, object]:
    if not openai_model.available():
        return {"available": False, "error": "OPENAI_API_KEY missing."}

    subset = samples[:max_samples]
    y_true = []
    y_pred = []
    for item in subset:
        image = Image.open(item["path"])
        pred = openai_model.predict(image)
        y_true.append(item["label"])
        y_pred.append(pred["top_prediction"]["label"])

    acc = accuracy_score(y_true, y_pred)
    report = classification_report(y_true, y_pred, output_dict=True, zero_division=0)
    return {
        "available": True,
        "evaluated_samples": len(subset),
        "accuracy": float(acc),
        "classification_report": report,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Compare custom model vs CLIP vs OpenAI Vision.")
    parser.add_argument("--data-dir", default="data/pokemon")
    parser.add_argument("--model-path", default="models/custom_resnet18.pth")
    parser.add_argument("--output", default="reports/model_comparison.json")
    parser.add_argument("--openai-max-samples", type=int, default=24)
    args = parser.parse_args()

    data_dir = Path(args.data_dir)
    output_path = Path(args.output)
    output_path.parent.mkdir(parents=True, exist_ok=True)

    custom = CustomModelPredictor(args.model_path)
    if not custom.available():
        raise FileNotFoundError(
            "Custom model checkpoint not found. Train model first with src/train_custom_model.py"
        )

    labels = custom.labels
    samples = iter_test_images(data_dir, labels)

    clip = ClipPredictor(labels)
    openai_model = OpenAIVisionPredictor(labels)

    result = {
        "dataset": str(data_dir),
        "num_test_samples": len(samples),
        "labels": labels,
        "custom_model": evaluate_custom(custom, samples),
        "clip_model": evaluate_clip(clip, samples),
        "openai_model": evaluate_openai(openai_model, samples, args.openai_max_samples),
    }

    output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
    print(f"Saved comparison report to: {output_path}")


if __name__ == "__main__":
    main()