"""Export Unlimited-OCR model forward paths to ONNX.

This exports tensor forward graphs, not the Python autoregressive generate()
loop or PIL-based image preprocessing used by model.infer().
"""

import argparse
import math
import os
import sys
import warnings
from pathlib import Path

# Colab/Jupyter can leak a notebook-only backend into the export environment.
# Transformers validates remote-code imports before loading the model and may
# import matplotlib even though export does not use it.
if os.environ.get("MPLBACKEND", "").startswith("module://matplotlib_inline"):
    os.environ["MPLBACKEND"] = "Agg"

import torch
from transformers import AutoModel

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from test_inference import validate_local_model_files  # noqa: E402


IMAGE_TOKEN_ID = 128815
BOS_TOKEN_ID = 0


class TextLogitsWrapper(torch.nn.Module):
    def __init__(self, model: torch.nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
            return_dict=False,
        )
        return outputs[0]


class TextDecodeWithCacheWrapper(torch.nn.Module):
    def __init__(self, model: torch.nn.Module, num_layers: int) -> None:
        super().__init__()
        self.model = model
        self.num_layers = num_layers

    def forward(
        self,
        input_ids: torch.LongTensor,
        position_ids: torch.LongTensor,
        *past_key_values: torch.Tensor,
    ) -> tuple[torch.Tensor, ...]:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=None,
            position_ids=position_ids,
            past_key_values=unflatten_past_key_values(
                past_key_values,
                self.num_layers,
            ),
            use_cache=True,
            return_dict=False,
        )
        return (outputs[0], *flatten_past_key_values(outputs[1]))


class ImagePrefillLogitsWrapper(torch.nn.Module):
    def __init__(self, model: torch.nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor,
        images_crop: torch.Tensor,
        images_ori: torch.Tensor,
        images_seq_mask: torch.Tensor,
        images_spatial_crop: torch.LongTensor,
    ) -> torch.Tensor:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            images=[(images_crop, images_ori)],
            images_seq_mask=images_seq_mask,
            images_spatial_crop=images_spatial_crop,
            use_cache=False,
            return_dict=False,
        )
        return outputs[0]


class ImagePrefillWithCacheWrapper(torch.nn.Module):
    def __init__(self, model: torch.nn.Module) -> None:
        super().__init__()
        self.model = model

    def forward(
        self,
        input_ids: torch.LongTensor,
        attention_mask: torch.Tensor,
        images_crop: torch.Tensor,
        images_ori: torch.Tensor,
        images_seq_mask: torch.Tensor,
        images_spatial_crop: torch.LongTensor,
    ) -> tuple[torch.Tensor, ...]:
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            images=[(images_crop, images_ori)],
            images_seq_mask=images_seq_mask,
            images_spatial_crop=images_spatial_crop,
            use_cache=True,
            return_dict=False,
        )
        return (outputs[0], *flatten_past_key_values(outputs[1]))


def flatten_past_key_values(past_key_values) -> tuple[torch.Tensor, ...]:
    if past_key_values is None:
        return ()
    if hasattr(past_key_values, "to_legacy_cache"):
        past_key_values = past_key_values.to_legacy_cache()

    values = []
    for layer_past in past_key_values:
        key, value = layer_past[:2]
        values.extend((key, value))
    return tuple(values)


def unflatten_past_key_values(
    past_key_values: tuple[torch.Tensor, ...],
    num_layers: int,
) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
    expected_tensors = num_layers * 2
    if len(past_key_values) != expected_tensors:
        raise ValueError(
            f"expected {expected_tensors} flattened past tensors, "
            f"got {len(past_key_values)}"
        )

    return tuple(
        (past_key_values[layer_idx * 2], past_key_values[layer_idx * 2 + 1])
        for layer_idx in range(num_layers)
    )


def select_export_device(requested_device: str) -> torch.device:
    if requested_device != "auto":
        return torch.device(requested_device)
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


def select_export_dtype(requested_dtype: str, device: torch.device) -> torch.dtype:
    if requested_dtype == "auto":
        return torch.float16 if device.type == "cuda" else torch.float32

    dtype_name = requested_dtype.removeprefix("torch.")
    dtype = getattr(torch, dtype_name, None)
    if not isinstance(dtype, torch.dtype):
        raise ValueError(f"Unsupported dtype: {requested_dtype}")
    if dtype is torch.bfloat16:
        raise SystemExit(
            "ONNX Runtime does not support bfloat16 Conv in this graph. "
            "Use --dtype auto, --dtype float16 on CUDA, or --dtype float32 on CPU."
        )
    return dtype


def image_token_count(image_size: int) -> int:
    patch_size = 16
    downsample_ratio = 4
    num_queries = math.ceil((image_size // patch_size) / downsample_ratio)
    return (num_queries + 1) * num_queries + 1


def build_text_inputs(
    sequence_length: int,
    device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor]:
    input_ids = torch.full(
        (1, sequence_length),
        fill_value=IMAGE_TOKEN_ID,
        dtype=torch.long,
        device=device,
    )
    input_ids[:, 0] = BOS_TOKEN_ID
    attention_mask = torch.ones_like(input_ids)
    return input_ids, attention_mask


def cache_head_dim(model: torch.nn.Module) -> int:
    if hasattr(model.config, "v_head_dim") and model.config.v_head_dim:
        return int(model.config.v_head_dim)
    return int(model.config.hidden_size // model.config.num_attention_heads)


def cache_num_heads(model: torch.nn.Module) -> int:
    return int(
        getattr(
            model.config,
            "num_key_value_heads",
            model.config.num_attention_heads,
        )
    )


def build_text_decode_cache_inputs(
    model: torch.nn.Module,
    sequence_length: int,
    past_sequence_length: int,
    device: torch.device,
    dtype: torch.dtype,
) -> tuple[torch.Tensor, ...]:
    if getattr(model.config, "use_mla", False):
        raise ValueError(
            "--kv-cache text decode export currently supports MHA cache shapes; "
            "this model config uses MLA cache shapes"
        )
    if past_sequence_length < 1:
        raise ValueError("--past-sequence-length must be at least 1 for cache decode")

    input_ids = torch.zeros((1, sequence_length), dtype=torch.long, device=device)
    position_ids = torch.arange(
        past_sequence_length,
        past_sequence_length + sequence_length,
        dtype=torch.long,
        device=device,
    ).unsqueeze(0)

    cache_shape = (
        1,
        cache_num_heads(model),
        past_sequence_length,
        cache_head_dim(model),
    )
    past_key_values = tuple(
        torch.zeros(cache_shape, dtype=dtype, device=device)
        for _ in range(model.config.num_hidden_layers * 2)
    )

    return (input_ids, position_ids, *past_key_values)


def build_image_prefill_inputs(
    image_size: int,
    sequence_length: int,
    device: torch.device,
    dtype: torch.dtype,
) -> tuple[torch.Tensor, ...]:
    token_count = image_token_count(image_size)
    min_sequence_length = token_count + 1
    if sequence_length < min_sequence_length:
        raise ValueError(
            f"--image-sequence-length must be at least {min_sequence_length} "
            f"for image_size={image_size}"
        )

    input_ids = torch.zeros((1, sequence_length), dtype=torch.long, device=device)
    input_ids[:, 0] = BOS_TOKEN_ID
    input_ids[:, 1 : token_count + 1] = IMAGE_TOKEN_ID
    attention_mask = torch.zeros_like(input_ids)
    attention_mask[:, :min_sequence_length] = 1

    images_seq_mask = torch.zeros_like(input_ids, dtype=torch.bool)
    images_seq_mask[:, 1 : token_count + 1] = True

    images_crop = torch.zeros(
        (1, 3, image_size, image_size),
        dtype=dtype,
        device=device,
    )
    images_ori = torch.ones(
        (1, 3, image_size, image_size),
        dtype=dtype,
        device=device,
    )
    images_spatial_crop = torch.tensor([[1, 1]], dtype=torch.long, device=device)

    return (
        input_ids,
        attention_mask,
        images_crop,
        images_ori,
        images_seq_mask,
        images_spatial_crop,
    )


def cache_input_names(num_layers: int) -> list[str]:
    names = []
    for layer_idx in range(num_layers):
        names.extend(
            [
                f"past_key_values.{layer_idx}.key",
                f"past_key_values.{layer_idx}.value",
            ]
        )
    return names


def cache_output_names(num_layers: int) -> list[str]:
    names = []
    for layer_idx in range(num_layers):
        names.extend(
            [
                f"present.{layer_idx}.key",
                f"present.{layer_idx}.value",
            ]
        )
    return names


def cache_dynamic_axes(
    input_names: list[str],
    output_names: list[str],
) -> dict[str, dict[int, str]]:
    axes = {
        "input_ids": {0: "batch", 1: "sequence"},
        "position_ids": {0: "batch", 1: "sequence"},
        "logits": {0: "batch", 1: "sequence"},
    }
    for name in input_names:
        if name.startswith("past_key_values."):
            axes[name] = {0: "batch", 2: "past_sequence"}
    for name in output_names:
        if name.startswith("present."):
            axes[name] = {0: "batch", 2: "total_sequence"}
    return axes


def export_onnx(
    wrapper: torch.nn.Module,
    example_inputs: tuple[torch.Tensor, ...],
    output_path: Path,
    input_names: list[str],
    output_names: list[str],
    opset: int,
    dynamo: bool,
    dynamic_axes: dict[str, dict[int, str]] | None,
) -> None:
    def run_export() -> None:
        torch.onnx.export(
            wrapper,
            example_inputs,
            str(output_path),
            input_names=input_names,
            output_names=output_names,
            opset_version=opset,
            dynamo=dynamo,
            external_data=True,
            dynamic_axes=dynamic_axes,
            export_params=True,
            do_constant_folding=True,
        )

    if dynamo or dynamic_axes is not None:
        run_export()
        return

    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="You are using the legacy TorchScript-based ONNX export.*",
            category=DeprecationWarning,
        )
        warnings.filterwarnings("ignore", category=torch.jit.TracerWarning)
        run_export()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--model",
        default=".",
        help="Model path or Hugging Face model id. Defaults to current directory.",
    )
    parser.add_argument(
        "--output",
        default="onnx/unlimited_ocr.onnx",
        help="Output ONNX path. Large weights are written as external data.",
    )
    parser.add_argument(
        "--target",
        choices=("image-prefill", "text"),
        default="image-prefill",
        help=(
            "Export image-prefill for OCR vision+LM logits, or text for the "
            "language-model logits path only."
        ),
    )
    parser.add_argument(
        "--device",
        default="auto",
        help="Export device. auto prefers CUDA and otherwise uses CPU.",
    )
    parser.add_argument(
        "--dtype",
        default="auto",
        help=(
            "Model dtype for export. auto uses float16 on CUDA and float32 "
            "on CPU. bfloat16 is not supported by ONNX Runtime for this graph."
        ),
    )
    parser.add_argument(
        "--opset",
        type=int,
        default=18,
        help="ONNX opset version.",
    )
    parser.add_argument(
        "--sequence-length",
        type=int,
        default=16,
        help=(
            "Dummy sequence length for --target text. Cache decode usually "
            "uses 1."
        ),
    )
    parser.add_argument(
        "--past-sequence-length",
        type=int,
        default=1,
        help="Dummy past KV length for --target text --kv-cache.",
    )
    parser.add_argument(
        "--image-size",
        type=int,
        default=1024,
        help="Dummy square image size for --target image-prefill.",
    )
    parser.add_argument(
        "--image-sequence-length",
        type=int,
        default=512,
        help=(
            "Fixed sequence length for --target image-prefill. Increase this "
            "if prompt tokens plus generated tokens exceed the default."
        ),
    )
    parser.add_argument(
        "--dynamic-text",
        action="store_true",
        help="Mark batch and sequence axes dynamic for --target text.",
    )
    parser.add_argument(
        "--dynamic-image",
        action="store_true",
        help=(
            "Deprecated for this MoE model. Image-prefill export uses a fixed "
            "sequence length; use --image-sequence-length to set capacity."
        ),
    )
    parser.add_argument(
        "--dynamo",
        action="store_true",
        help="Use the torch.export-based ONNX exporter. Legacy tracing is default.",
    )
    parser.add_argument(
        "--kv-cache",
        action="store_true",
        help=(
            "Export cache-aware graph outputs. For image-prefill this returns "
            "present.* tensors. For text this exports a decode graph that "
            "accepts flattened past_key_values.* tensors and returns updated "
            "present.* tensors."
        ),
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.dynamic_text and args.target != "text":
        raise SystemExit("--dynamic-text is only supported with --target text")
    if args.dynamic_image and args.target != "image-prefill":
        raise SystemExit(
            "--dynamic-image is only supported with --target image-prefill"
        )

    device = select_export_device(args.device)
    dtype = select_export_dtype(args.dtype, device)

    print(f"Validating model files in {args.model!r}")
    validate_local_model_files(args.model)

    print(f"Loading model from {args.model!r} on {device} with {dtype}")
    model = AutoModel.from_pretrained(
        args.model,
        trust_remote_code=True,
        use_safetensors=True,
        dtype=dtype,
    )
    model.config.device = str(device)
    model.config.inference_dtype = dtype
    model.config.use_cache = args.kv_cache
    model.config.output_attentions = False
    model.config.output_hidden_states = False
    if hasattr(model.config, "sliding_window"):
        model.config.sliding_window = None
    model = model.eval().to(device)

    output_names = ["logits"]
    if args.target == "text" and args.kv_cache:
        wrapper = TextDecodeWithCacheWrapper(
            model,
            num_layers=model.config.num_hidden_layers,
        ).eval()
        example_inputs = build_text_decode_cache_inputs(
            model,
            args.sequence_length,
            args.past_sequence_length,
            device,
            dtype,
        )
        input_names = [
            "input_ids",
            "position_ids",
            *cache_input_names(model.config.num_hidden_layers),
        ]
        output_names = [
            "logits",
            *cache_output_names(model.config.num_hidden_layers),
        ]
        dynamic_axes = cache_dynamic_axes(input_names, output_names)
    elif args.target == "text":
        wrapper = TextLogitsWrapper(model).eval()
        example_inputs = build_text_inputs(args.sequence_length, device)
        input_names = ["input_ids", "attention_mask"]
        dynamic_axes = (
            {
                "input_ids": {0: "batch", 1: "sequence"},
                "attention_mask": {0: "batch", 1: "sequence"},
                "logits": {0: "batch", 1: "sequence"},
            }
            if args.dynamic_text
            else None
        )
    elif args.kv_cache:
        wrapper = ImagePrefillWithCacheWrapper(model).eval()
        example_inputs = build_image_prefill_inputs(
            args.image_size,
            args.image_sequence_length,
            device,
            dtype,
        )
        input_names = [
            "input_ids",
            "attention_mask",
            "images_crop",
            "images_ori",
            "images_seq_mask",
            "images_spatial_crop",
        ]
        output_names = [
            "logits",
            *cache_output_names(model.config.num_hidden_layers),
        ]
        dynamic_axes = None
        if args.dynamic_image:
            print(
                "--dynamic-image is ignored for this MoE model; exporting fixed "
                f"sequence length {args.image_sequence_length}."
            )
    else:
        wrapper = ImagePrefillLogitsWrapper(model).eval()
        example_inputs = build_image_prefill_inputs(
            args.image_size,
            args.image_sequence_length,
            device,
            dtype,
        )
        input_names = [
            "input_ids",
            "attention_mask",
            "images_crop",
            "images_ori",
            "images_seq_mask",
            "images_spatial_crop",
        ]
        dynamic_axes = None
        if args.dynamic_image:
            print(
                "--dynamic-image is ignored for this MoE model; exporting fixed "
                f"sequence length {args.image_sequence_length}."
            )

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

    print(f"Exporting {args.target!r} graph to {output_path}")
    print(
        "Large checkpoints are stored with ONNX external data next to the .onnx file."
    )
    try:
        with torch.no_grad():
            export_onnx(
                wrapper,
                example_inputs,
                output_path,
                input_names,
                output_names,
                args.opset,
                args.dynamo,
                dynamic_axes=dynamic_axes,
            )
    except ModuleNotFoundError as error:
        if error.name in {"onnx", "onnxscript"}:
            raise SystemExit(
                "Missing ONNX export dependency. Install it with:\n"
                "  uv add --group export onnx onnxscript\n"
                "or run once with:\n"
                "  uv run --with onnx --with onnxscript python scripts/export_onnx.py"
            ) from error
        raise

    print(f"Export complete: {output_path}")


if __name__ == "__main__":
    main()