"""ONNX export script for SAM2 image encoder and decoder.

This module provides functionality to export SAM2 models to ONNX format,
splitting the model into separate encoder and decoder components for efficient
inference without PyTorch dependencies.

The export uses the legacy ONNX exporter (dynamo=False) as the dynamo exporter
hangs on SAM2's complex operations (particularly RoPE).
"""

from __future__ import annotations

import argparse
import json
import logging
from pathlib import Path
from typing import TYPE_CHECKING

import torch
import torch.nn as nn

from sam2_onnx.metadata import ModelMetadata

if TYPE_CHECKING:
    from sam2.modeling.sam2_base import SAM2Base

logger = logging.getLogger(__name__)


class EncoderWrapper(nn.Module):
    """Wraps SAM2 image encoder for ONNX export.

    Combines the image encoder, high-res feature projections (conv_s0, conv_s1),
    and no_mem_embed logic to produce encoder outputs ready for the decoder.

    Forward output:
        - image_embed: [B, C, H, W] main image embedding
        - high_res_feat_0: [B, C//8, 4*H, 4*W] high-res feature level 0
        - high_res_feat_1: [B, C//4, 2*H, 2*W] high-res feature level 1
    """

    def __init__(self, model: SAM2Base):
        super().__init__()
        self.image_encoder = model.image_encoder
        self.conv_s0 = model.sam_mask_decoder.conv_s0
        self.conv_s1 = model.sam_mask_decoder.conv_s1
        self.no_mem_embed = model.no_mem_embed
        self.directly_add_no_mem_embed = model.directly_add_no_mem_embed
        self.hidden_dim = model.hidden_dim
        self.num_feature_levels = model.num_feature_levels
        self.use_high_res_features = model.use_high_res_features_in_sam

    def forward(self, image: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward pass matching model.forward_image + _prepare_backbone_features logic.

        Args:
            image: [B, 3, H, W] input image

        Returns:
            Tuple of (image_embed, high_res_feat_0, high_res_feat_1)
        """
        # Run backbone
        backbone_out = self.image_encoder(image)

        # Apply high-res projections if enabled (matches forward_image)
        if self.use_high_res_features:
            backbone_out["backbone_fpn"][0] = self.conv_s0(backbone_out["backbone_fpn"][0])
            backbone_out["backbone_fpn"][1] = self.conv_s1(backbone_out["backbone_fpn"][1])

        # Extract features (matches _prepare_backbone_features)
        feature_maps = backbone_out["backbone_fpn"][-self.num_feature_levels:]

        # Get top-level vision features (lowest resolution)
        vision_feats = feature_maps[-1]  # [B, C, H, W]

        # Add no_mem_embed if directly_add_no_mem_embed is True (matches set_image)
        if self.directly_add_no_mem_embed:
            # no_mem_embed is [1, 1, C], need to reshape to [B, C, H, W]
            B, C, H, W = vision_feats.shape
            no_mem_expanded = self.no_mem_embed.view(1, C, 1, 1).expand(B, C, H, W)
            vision_feats = vision_feats + no_mem_expanded

        image_embed = vision_feats

        # Extract high-res features if available
        if self.use_high_res_features and len(feature_maps) >= 3:
            high_res_feat_0 = feature_maps[0]  # highest resolution
            high_res_feat_1 = feature_maps[1]  # mid resolution
        else:
            # Return dummy features with correct shapes if not using high-res
            B, C, H, W = image_embed.shape
            high_res_feat_0 = torch.zeros(B, C // 8, H * 4, W * 4, device=image_embed.device)  # pragma: no mutate
            high_res_feat_1 = torch.zeros(B, C // 4, H * 2, W * 2, device=image_embed.device)

        return image_embed, high_res_feat_0, high_res_feat_1


class DecoderWrapper(nn.Module):
    """Wraps SAM2 PromptEncoder + MaskDecoder for ONNX export.

    Inlines PromptEncoder logic to avoid ONNX tracing issues with None branches.
    Always runs mask embedding path and uses has_mask_input flag to select output.

    Forward inputs:
        - image_embed: [B, C, H, W]
        - high_res_feat_0: [B, C//8, 4*H, 4*W]
        - high_res_feat_1: [B, C//4, 2*H, 2*W]
        - point_coords: [B, N, 2] in image coordinates
        - point_labels: [B, N] with values 0/1/-1
        - mask_input: [B, 1, 4*H, 4*W] low-res mask (always provided, may be zeros)
        - has_mask_input: [B, 1] binary flag (0 = no mask, 1 = has mask)

    Forward output:
        - low_res_masks: [B, 3, 4*H, 4*W] (multimask_output=True always)
        - iou_predictions: [B, 3]
    """

    def __init__(self, model: SAM2Base):
        super().__init__()
        self.prompt_encoder = model.sam_prompt_encoder
        self.mask_decoder = model.sam_mask_decoder

    def forward(
        self,
        image_embed: torch.Tensor,
        high_res_feat_0: torch.Tensor,
        high_res_feat_1: torch.Tensor,
        point_coords: torch.Tensor,
        point_labels: torch.Tensor,
        mask_input: torch.Tensor,
        has_mask_input: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Forward pass with inlined PromptEncoder logic.

        All inputs must be provided (no None allowed for ONNX).
        """
        # Inline PromptEncoder._embed_points with pad=True
        sparse = self.prompt_encoder._embed_points(point_coords, point_labels, pad=True)

        # Inline PromptEncoder mask embedding logic
        # Always run mask_downscaling (even if no mask)
        mask_embed = self.prompt_encoder.mask_downscaling(mask_input)

        # Get no_mask_embed and expand
        no_mask = self.prompt_encoder.no_mask_embed.weight.reshape(1, -1, 1, 1)
        B = mask_embed.shape[0]
        C, H, W = mask_embed.shape[1], mask_embed.shape[2], mask_embed.shape[3]
        no_mask = no_mask.expand(B, C, H, W)

        # Use has_mask_input to select between mask_embed and no_mask
        # has_mask_input: [B, 1] -> [B, 1, 1, 1] for broadcasting
        has_mask_flag = has_mask_input.view(B, 1, 1, 1)
        dense = torch.where(has_mask_flag > 0, mask_embed, no_mask)

        # Get dense positional encoding
        image_pe = self.prompt_encoder.get_dense_pe()

        # Prepare high_res_features list
        high_res_features = [high_res_feat_0, high_res_feat_1]

        # Call mask decoder with multimask_output=True (always)
        low_res_masks, iou_predictions, _, _ = self.mask_decoder(
            image_embeddings=image_embed,
            image_pe=image_pe,
            sparse_prompt_embeddings=sparse,
            dense_prompt_embeddings=dense,
            multimask_output=True,
            repeat_image=False,
            high_res_features=high_res_features,
        )

        return low_res_masks, iou_predictions


def export_encoder(
    model: SAM2Base,
    output_path: Path,
    image_size: int,
    opset: int = 17,
) -> None:
    """Export encoder to ONNX using legacy exporter.

    Args:
        model: SAM2Base model
        output_path: Path to save encoder.onnx
        image_size: Image size (e.g., 1024 or 256)
        opset: ONNX opset version
    """
    encoder = EncoderWrapper(model)
    encoder.eval()

    # Deterministic dummy input [1, 3, H, W] — seeded randn, not zeros
    # (zeros cause degenerate attention/LayerNorm traces in Hiera backbone)
    gen = torch.Generator().manual_seed(0)
    dummy_input = torch.randn(1, 3, image_size, image_size, generator=gen)

    torch.onnx.export(
        encoder,
        (dummy_input,),
        str(output_path),
        input_names=["image"],
        output_names=["image_embed", "high_res_feat_0", "high_res_feat_1"],
        opset_version=opset,
        dynamo=False,  # CRITICAL: use legacy exporter
    )

    logger.info(f"Encoder exported to {output_path}")


def export_decoder(
    model: SAM2Base,
    output_path: Path,
    image_size: int,
    opset: int = 17,
) -> None:
    """Export decoder to ONNX using legacy exporter.

    Args:
        model: SAM2Base model
        output_path: Path to save decoder.onnx
        image_size: Image size (used to compute embedding dims)
        opset: ONNX opset version
    """
    decoder = DecoderWrapper(model)
    decoder.eval()

    # Compute embedding dimensions
    backbone_stride = model.backbone_stride
    embed_size = image_size // backbone_stride
    embed_dim = model.hidden_dim

    # Deterministic dummy inputs — seeded randn for non-degenerate traces
    gen = torch.Generator().manual_seed(0)
    dummy_image_embed = torch.randn(1, embed_dim, embed_size, embed_size, generator=gen)
    dummy_high_res_feat_0 = torch.randn(1, embed_dim // 8, embed_size * 4, embed_size * 4, generator=gen)
    dummy_high_res_feat_1 = torch.randn(1, embed_dim // 4, embed_size * 2, embed_size * 2, generator=gen)
    dummy_point_coords = torch.randn(1, 3, 2, generator=gen)  # 3 points
    dummy_point_labels = torch.ones(1, 3, dtype=torch.int64)
    dummy_mask_input = torch.randn(1, 1, embed_size * 4, embed_size * 4, generator=gen)
    dummy_has_mask_input = torch.zeros(1, 1)

    # Dynamic axes for point inputs
    dynamic_axes = {
        "point_coords": {1: "num_points"},
        "point_labels": {1: "num_points"},
    }

    torch.onnx.export(
        decoder,
        (
            dummy_image_embed,
            dummy_high_res_feat_0,
            dummy_high_res_feat_1,
            dummy_point_coords,
            dummy_point_labels,
            dummy_mask_input,
            dummy_has_mask_input,
        ),
        str(output_path),
        input_names=[
            "image_embed",
            "high_res_feat_0",
            "high_res_feat_1",
            "point_coords",
            "point_labels",
            "mask_input",
            "has_mask_input",
        ],
        output_names=["low_res_masks", "iou_predictions"],
        dynamic_axes=dynamic_axes,
        opset_version=opset,
        dynamo=False,  # CRITICAL: use legacy exporter
    )

    logger.info(f"Decoder exported to {output_path}")


def validate_onnx(
    encoder_path: Path,
    decoder_path: Path,
    model: SAM2Base,
    image_size: int,
) -> bool:
    """Validate ONNX exports against PyTorch model.

    Args:
        encoder_path: Path to encoder.onnx
        decoder_path: Path to decoder.onnx
        model: Original SAM2Base model
        image_size: Image size

    Returns:
        True if all outputs match within tolerance, False otherwise.
    """
    import onnxruntime as ort
    import numpy as np

    logger.info("Validating ONNX exports...")

    # Load ONNX models
    encoder_session = ort.InferenceSession(str(encoder_path))
    decoder_session = ort.InferenceSession(str(decoder_path))

    # Deterministic test inputs — seeded for reproducibility
    gen = torch.Generator().manual_seed(0)
    model.eval()
    with torch.no_grad():
        test_image = torch.randn(1, 3, image_size, image_size, generator=gen)

        # PyTorch encoder
        encoder_wrapper = EncoderWrapper(model)
        encoder_wrapper.eval()
        torch_image_embed, torch_hr0, torch_hr1 = encoder_wrapper(test_image)

        # ONNX encoder
        onnx_outputs = encoder_session.run(
            None,
            {"image": test_image.numpy()}
        )
        onnx_image_embed = onnx_outputs[0]
        onnx_hr0 = onnx_outputs[1]
        onnx_hr1 = onnx_outputs[2]

        # Compare encoder outputs
        encoder_diff = np.abs(torch_image_embed.numpy() - onnx_image_embed).max()
        hr0_diff = np.abs(torch_hr0.numpy() - onnx_hr0).max()
        hr1_diff = np.abs(torch_hr1.numpy() - onnx_hr1).max()

        logger.info(f"Encoder validation: image_embed={encoder_diff:.2e}, high_res_feat_0={hr0_diff:.2e}, high_res_feat_1={hr1_diff:.2e}")

        # Test decoder
        embed_size = image_size // model.backbone_stride
        test_points = torch.tensor([[[100.0, 100.0], [200.0, 200.0]]])
        test_labels = torch.tensor([[1, 0]])
        test_mask = torch.zeros(1, 1, embed_size * 4, embed_size * 4)
        test_has_mask = torch.zeros(1, 1)

        # PyTorch decoder
        decoder_wrapper = DecoderWrapper(model)
        decoder_wrapper.eval()
        torch_masks, torch_ious = decoder_wrapper(
            torch_image_embed,
            torch_hr0,
            torch_hr1,
            test_points,
            test_labels,
            test_mask,
            test_has_mask,
        )

        # ONNX decoder
        onnx_decoder_outputs = decoder_session.run(
            None,
            {
                "image_embed": onnx_image_embed,
                "high_res_feat_0": onnx_hr0,
                "high_res_feat_1": onnx_hr1,
                "point_coords": test_points.numpy(),
                "point_labels": test_labels.numpy(),
                "mask_input": test_mask.numpy(),
                "has_mask_input": test_has_mask.numpy(),
            }
        )
        onnx_masks = onnx_decoder_outputs[0]
        onnx_ious = onnx_decoder_outputs[1]

        # Compare decoder outputs
        mask_diff = np.abs(torch_masks.numpy() - onnx_masks).max()
        iou_diff = np.abs(torch_ious.numpy() - onnx_ious).max()

        logger.info(f"Decoder validation: low_res_masks={mask_diff:.2e}, iou_predictions={iou_diff:.2e}")

        atol = 1e-5
        passed = encoder_diff < atol and hr0_diff < atol and hr1_diff < atol and mask_diff < atol and iou_diff < atol  # pragma: no mutate
        if passed:
            logger.info("Validation passed (atol=1e-5)")
        else:
            logger.warning("Validation failed (atol=1e-5)")
        return passed


def extract_metadata(model: SAM2Base, quantization: str | None = None) -> ModelMetadata:
    """Extract metadata from SAM2 model.

    Args:
        model: SAM2Base model
        quantization: Optional quantization mode ("uint8" or "int8")

    Returns:
        ModelMetadata instance
    """
    return ModelMetadata(
        image_size=model.image_size,
        embed_dim=model.hidden_dim,
        backbone_stride=model.backbone_stride,
        num_multimask_outputs=3,
        mask_threshold=0.0,
        quantization=quantization,
    )


def quantize_model(
    input_path: Path,
    output_path: Path,
    mode: str = "uint8",
) -> None:
    """Quantize ONNX model using dynamic quantization.

    Args:
        input_path: Path to input ONNX model
        output_path: Path to save quantized model
        mode: Quantization mode ("uint8" or "int8")
    """
    from onnxruntime.quantization import quantize_dynamic, QuantType

    quant_types = {"uint8": QuantType.QUInt8, "int8": QuantType.QInt8}
    if mode not in quant_types:
        raise ValueError(f"quantization mode must be 'uint8' or 'int8', got {mode!r}")
    quant_type = quant_types[mode]

    quantize_dynamic(
        str(input_path),
        str(output_path),
        weight_type=quant_type,
    )

    logger.info(f"Quantized model saved to {output_path}")


def main():
    """CLI entry point for ONNX export."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")

    parser = argparse.ArgumentParser(description="Export SAM2 to ONNX format")
    parser.add_argument(
        "--config",
        type=str,
        required=True,
        help="Path to SAM2 config (e.g., configs/sam2.1/sam2.1_hiera_t.yaml)",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        default=None,
        help="Path to SAM2 checkpoint (optional for random weights)",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="Output directory for ONNX files",
    )
    parser.add_argument(
        "--quantize",
        type=str,
        choices=["uint8", "int8"],
        default=None,
        help="Quantization mode (optional)",
    )
    parser.add_argument(
        "--opset",
        type=int,
        default=17,
        help="ONNX opset version (default: 17)",
    )
    parser.add_argument(
        "--no-validate",
        action="store_true",
        help="Skip validation step",
    )

    args = parser.parse_args()

    # Import here to allow module-level imports without torch
    from sam2.build_sam import build_sam2

    # Build model
    print(f"Building SAM2 model from {args.config}...")
    model = build_sam2(
        args.config,
        ckpt_path=args.checkpoint,
        device="cpu",
        mode="eval",
    )
    model.eval()

    image_size = model.image_size
    print(f"Model image size: {image_size}")

    # Create output directory
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Export encoder
    encoder_path = output_dir / "encoder.onnx"
    print(f"\nExporting encoder to {encoder_path}...")
    export_encoder(model, encoder_path, image_size, args.opset)

    # Export decoder
    decoder_path = output_dir / "decoder.onnx"
    print(f"\nExporting decoder to {decoder_path}...")
    export_decoder(model, decoder_path, image_size, args.opset)

    # Quantize if requested
    if args.quantize:
        print(f"\nQuantizing models with {args.quantize}...")
        encoder_quant_path = output_dir / f"encoder_{args.quantize}.onnx"
        decoder_quant_path = output_dir / f"decoder_{args.quantize}.onnx"

        quantize_model(encoder_path, encoder_quant_path, args.quantize)
        quantize_model(decoder_path, decoder_quant_path, args.quantize)

        # Update paths for validation
        encoder_path = encoder_quant_path
        decoder_path = decoder_quant_path

    # Extract and save metadata
    metadata = extract_metadata(model, quantization=args.quantize)
    metadata_path = output_dir / "metadata.json"
    with open(metadata_path, "w") as f:
        json.dump(metadata.to_dict(), f, indent=2)
    print(f"\nMetadata saved to {metadata_path}")

    # Validate
    if not args.no_validate:
        validate_onnx(encoder_path, decoder_path, model, image_size)

    print(f"\n✓ Export complete. Files saved to {output_dir}")


if __name__ == "__main__":
    main()
