from __future__ import annotations

import argparse
import inspect
from pathlib import Path

import torch
from torch import nn

from .models import UNetInstance, UNetSmall


class _SegOnlyWrapper(nn.Module):
    """Wrap an instance model to export only segmentation logits."""

    def __init__(self, model: UNetInstance):
        super().__init__()
        self.model = model

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        logits, _embeddings = self.model(x)
        return logits


def _looks_like_instance_state(state_dict: dict[str, torch.Tensor]) -> bool:
    keys = list(state_dict.keys())
    return any(k.startswith("emb_head.") or k.startswith("seg_head.") for k in keys)


def load_model_from_ckpt(
    ckpt_path: str | Path,
    model_type: str = "auto",
    *,
    with_embeddings: bool = False,
) -> tuple[nn.Module, bool]:
    load_kwargs = {"map_location": "cpu"}
    # torch>=2.4 supports weights_only=True and avoids unpickling arbitrary objects.
    try:
        ckpt = torch.load(str(ckpt_path), weights_only=True, **load_kwargs)
    except TypeError:
        ckpt = torch.load(str(ckpt_path), **load_kwargs)
    state_dict = ckpt.get("model_state") if isinstance(ckpt, dict) and "model_state" in ckpt else ckpt
    if not isinstance(state_dict, dict):
        raise ValueError(f"Checkpoint does not contain a valid state_dict: {ckpt_path}")

    ckpt_args = ckpt.get("args", {}) if isinstance(ckpt, dict) else {}
    if not isinstance(ckpt_args, dict):
        ckpt_args = {}

    model_type = str(model_type).strip().lower()
    if model_type not in {"auto", "seg", "instance"}:
        raise ValueError(f"Unknown --model_type={model_type!r}; expected auto|seg|instance")

    inferred_instance = _looks_like_instance_state(state_dict)
    use_instance = inferred_instance if model_type == "auto" else (model_type == "instance")

    if use_instance:
        model = UNetInstance(
            in_channels=1,
            num_classes=3,
            base_channels=int(ckpt_args.get("base_channels", 32)),
            embedding_dim=int(ckpt_args.get("embedding_dim", 8)),
            normalize_embeddings=bool(ckpt_args.get("normalize_embeddings", True)),
        )
        model.load_state_dict(state_dict)
        model.eval()
        return (model, True) if with_embeddings else (_SegOnlyWrapper(model), True)

    if with_embeddings:
        raise ValueError("--with_embeddings requires an instance checkpoint/model_type=instance")

    model = UNetSmall(
        in_channels=1,
        num_classes=3,
        base_channels=int(ckpt_args.get("base_channels", 32)),
    )
    model.load_state_dict(state_dict)
    model.eval()
    return model, False


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Export trained segmentation model to ONNX")
    ap.add_argument("--ckpt", required=True, help="Path to .pt checkpoint")
    ap.add_argument("--out", required=True, help="Output ONNX path")
    ap.add_argument(
        "--model_type",
        default="auto",
        choices=["auto", "seg", "instance"],
        help="Checkpoint model family. 'auto' infers from state_dict keys.",
    )
    ap.add_argument(
        "--with_embeddings",
        action="store_true",
        help="For instance checkpoints, export dual outputs: logits and embeddings.",
    )
    ap.add_argument("--T", type=int, default=220, help="Dummy input T dimension")
    ap.add_argument("--X", type=int, default=180, help="Dummy input X dimension")
    ap.add_argument("--opset", type=int, default=17)
    ap.add_argument(
        "--legacy",
        action="store_true",
        help="Force legacy ONNX exporter (dynamo=False).",
    )
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    model, is_instance = load_model_from_ckpt(
        args.ckpt,
        model_type=args.model_type,
        with_embeddings=bool(args.with_embeddings),
    )

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    dummy = torch.zeros((1, 1, args.T, args.X), dtype=torch.float32)
    output_names = ["logits"]
    dynamic_axes = {"input": {2: "T", 3: "X"}, "logits": {2: "T", 3: "X"}}
    if bool(args.with_embeddings):
        if not is_instance:
            raise ValueError("--with_embeddings was requested but checkpoint is not instance-aware")
        output_names = ["logits", "embeddings"]
        dynamic_axes["embeddings"] = {2: "T", 3: "X"}

    export_kwargs = dict(
        input_names=["input"],
        output_names=output_names,
        dynamic_axes=dynamic_axes,
        opset_version=args.opset,
    )
    supports_dynamo = "dynamo" in inspect.signature(torch.onnx.export).parameters

    def _export(use_dynamo: bool) -> None:
        if supports_dynamo:
            torch.onnx.export(model, dummy, str(out_path), dynamo=bool(use_dynamo), **export_kwargs)
        else:
            torch.onnx.export(model, dummy, str(out_path), **export_kwargs)

    def _raise_missing_onnx(exc: ModuleNotFoundError) -> None:
        if exc.name == "onnx":
            raise RuntimeError(
                "The 'onnx' package is required to export models. "
                "Install it in this environment with: python -m pip install onnx"
            ) from exc
        raise exc

    if args.legacy or not supports_dynamo:
        try:
            _export(use_dynamo=False)
        except ModuleNotFoundError as exc:
            _raise_missing_onnx(exc)
    else:
        try:
            _export(use_dynamo=True)
        except ModuleNotFoundError as exc:
            if exc.name == "onnxscript":
                print("onnxscript not found. Falling back to legacy ONNX exporter (dynamo=False).")
                try:
                    _export(use_dynamo=False)
                except ModuleNotFoundError as exc2:
                    _raise_missing_onnx(exc2)
            else:
                _raise_missing_onnx(exc)
        except Exception as exc:
            print(
                f"Dynamo ONNX export failed ({type(exc).__name__}). "
                "Falling back to legacy ONNX exporter (dynamo=False)."
            )
            try:
                _export(use_dynamo=False)
            except ModuleNotFoundError as exc2:
                _raise_missing_onnx(exc2)

    print(f"Exported ONNX model: {out_path}")


if __name__ == "__main__":
    main()
