from __future__ import annotations

import argparse
import sys
from pathlib import Path

from alphabet.utils import setup_logging


def main() -> None:
    parser = argparse.ArgumentParser(description="Alphabet (A-Z) classifier CLI")
    parser.add_argument("--log-level", default="INFO")
    subparsers = parser.add_subparsers(dest="command", required=True)

    p_train = subparsers.add_parser("train", help="Download EMNIST and train the model")
    p_train.add_argument("--config", default="config/config.yaml")
    p_train.add_argument("--device", default=None, help="cuda | cpu (auto-detects if omitted)")

    p_export = subparsers.add_parser("export", help="Export a checkpoint to ONNX and TorchScript")
    p_export.add_argument("--config", default="config/config.yaml")
    p_export.add_argument("--checkpoint", required=True, help="Path to best.pt")
    p_export.add_argument("--device", default="cpu")

    p_infer = subparsers.add_parser("infer", help="Run inference on a single image (for testing)")
    p_infer.add_argument("--model", required=True, help="Path to .onnx model")
    p_infer.add_argument("--image", required=True, help="Path to a letter crop image")
    p_infer.add_argument("--config", default="config/config.yaml")

    args = parser.parse_args()
    setup_logging(args.log_level)

    if args.command == "train":
        from alphabet.train import train_from_config
        summary = train_from_config(args.config, device=args.device)
        print(f"\nBest val accuracy: {summary['best_val_accuracy']:.4f}")
        return

    if args.command == "export":
        from alphabet.export import export_from_config
        exported = export_from_config(args.config, args.checkpoint, device=args.device)
        print(exported)
        return

    if args.command == "infer":
        import cv2
        import numpy as np
        from alphabet.dataset import get_class_list
        from alphabet.infer import OnnxBackend, predict_crops
        from alphabet.utils import load_yaml

        cfg = load_yaml(args.config)
        infer_cfg = cfg.get("infer", {})
        model_cfg = cfg["model"]
        class_list = get_class_list(cfg)

        crop = cv2.imread(args.image, cv2.IMREAD_COLOR)
        if crop is None:
            print(f"Error: could not read image at {args.image}", file=sys.stderr)
            sys.exit(1)

        backend = OnnxBackend(args.model)
        results = predict_crops(
            backend,
            [crop],
            img_size=int(model_cfg.get("img_size", 64)),
            mean=float(model_cfg.get("mean", 0.5)),
            std=float(model_cfg.get("std", 0.5)),
            min_confidence=float(infer_cfg.get("min_confidence", 0.60)),
            class_list=class_list,
        )
        r = results[0]
        flag_str = " [FLAGGED for review]" if r["flag"] else ""
        print(f"Prediction : {r['character']}{flag_str}")
        print(f"Confidence : {r['confidence']:.4f}")
        top3 = sorted(enumerate(r["probabilities"]), key=lambda x: -x[1])[:3]
        for idx, prob in top3:
            print(f"  {class_list[idx]}: {prob:.4f}")
        return


if __name__ == "__main__":
    main()