"""
ONNX export functionality for BADAS models.

Provides functions to export BADAS models to ONNX format with
embedded weights and validation.
"""

import torch
from pathlib import Path
from typing import Optional, Tuple, Union
from dataclasses import dataclass

from badas.export.wrappers import create_export_wrapper
from badas.export.utils import patch_vjepa2_rope, embed_onnx_weights, validate_onnx


@dataclass
class ONNXExportConfig:
    """Configuration for ONNX export.

    Attributes:
        opset_version: ONNX opset version (default: 17).
        embed_weights: Embed weights in single file (default: True).
        include_predictor: Include future prediction module (default: False).
        validate: Validate exported model (default: True).
        dynamic_batch: Enable dynamic batch size (default: False).
        verbose: Print detailed export info (default: False).
    """
    opset_version: int = 17
    embed_weights: bool = True
    include_predictor: bool = False
    validate: bool = True
    dynamic_batch: bool = False
    verbose: bool = False


@dataclass
class ONNXExportResult:
    """Result of ONNX export operation.

    Attributes:
        success: Whether export succeeded.
        onnx_path: Path to exported ONNX file.
        file_size_mb: Size of exported file in MB.
        validation: Validation results (if performed).
        error: Error message (if failed).
    """
    success: bool
    onnx_path: Optional[str] = None
    file_size_mb: Optional[float] = None
    validation: Optional[dict] = None
    error: Optional[str] = None


def export_to_onnx(
    checkpoint_path: str,
    output_path: str,
    config: Optional[ONNXExportConfig] = None,
    model_name_override: Optional[str] = None,
) -> ONNXExportResult:
    """
    Export a BADAS checkpoint to ONNX format.

    Args:
        checkpoint_path: Path to .ckpt checkpoint file.
        output_path: Path for output ONNX file.
        config: Export configuration (uses defaults if None).
        model_name_override: HuggingFace model ID to use instead of
            the path stored in checkpoint. Required when deploying
            to machines without access to original training paths.

    Returns:
        ONNXExportResult with export status and details.

    Example:
        >>> from badas.export import export_to_onnx, ONNXExportConfig
        >>>
        >>> # Simple export
        >>> result = export_to_onnx("model.ckpt", "model.onnx")
        >>> print(f"Exported: {result.onnx_path} ({result.file_size_mb:.1f} MB)")
        >>>
        >>> # With custom config
        >>> config = ONNXExportConfig(opset_version=18, verbose=True)
        >>> result = export_to_onnx("model.ckpt", "model.onnx", config=config)
        >>>
        >>> # With model override for deployment
        >>> result = export_to_onnx(
        ...     "flash.ckpt",
        ...     "flash.onnx",
        ...     model_name_override="nexar-ai/vjepa2-vitb-fpc16-256-nexar",
        ... )
    """
    config = config or ONNXExportConfig()

    try:
        # Patch V-JEPA2 for ONNX compatibility
        if patch_vjepa2_rope() and config.verbose:
            print("   Patched V-JEPA2 RoPE for ONNX compatibility")

        # Load model
        from badas.inference.checkpoint import load_model_from_checkpoint

        model, model_config = load_model_from_checkpoint(
            checkpoint_path,
            map_location='cpu',
            model_name_override=model_name_override,
        )

        # Create export wrapper
        wrapper_kwargs = {}
        if config.include_predictor:
            wrapper_kwargs.update({
                'include_predictor': True,
                'frame_count': model_config.data.frame_count,
                'future_seconds': model_config.model.future_prediction_seconds,
                'original_fps': model_config.model.original_fps,
            })

        export_model = create_export_wrapper(model, **wrapper_kwargs)
        export_model = export_model.float().eval()

        # Create dummy input
        dummy_input = torch.randn(
            1,
            model_config.data.frame_count,
            3,
            model_config.data.img_size,
            model_config.data.img_size,
            dtype=torch.float32,
        )

        # Verify forward pass
        with torch.no_grad():
            expected_output = export_model(dummy_input)

        if config.verbose:
            print(f"   Forward pass OK, output shape: {expected_output.shape}")

        # Setup output path
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Configure dynamic axes
        dynamic_axes = None
        if config.dynamic_batch:
            dynamic_axes = {
                'pixel_values_videos': {0: 'batch_size'},
                'logits': {0: 'batch_size'},
            }

        # Export to ONNX
        torch.onnx.export(
            export_model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=config.opset_version,
            do_constant_folding=True,
            input_names=['pixel_values_videos'],
            output_names=['logits'],
            dynamic_axes=dynamic_axes,
            verbose=False,
        )

        # Embed weights if requested
        if config.embed_weights:
            embed_onnx_weights(str(output_path))
            if config.verbose:
                print("   Weights embedded in single file")

        file_size_mb = output_path.stat().st_size / 1024 / 1024

        # Validate if requested
        validation = None
        if config.validate:
            validation = validate_onnx(
                str(output_path),
                input_data=dummy_input,
                expected_output=expected_output,
            )
            if config.verbose:
                status = "PASSED" if validation['structure_valid'] else "FAILED"
                print(f"   Validation: {status}")
                if validation.get('max_diff') is not None:
                    print(f"   Max numerical diff: {validation['max_diff']:.6f}")

        return ONNXExportResult(
            success=True,
            onnx_path=str(output_path),
            file_size_mb=file_size_mb,
            validation=validation,
        )

    except Exception as e:
        return ONNXExportResult(
            success=False,
            error=str(e),
        )


def export_model_to_onnx(
    model: torch.nn.Module,
    output_path: str,
    input_shape: Tuple[int, ...],
    config: Optional[ONNXExportConfig] = None,
) -> ONNXExportResult:
    """
    Export an already-loaded model to ONNX.

    Use this when you already have a model instance and don't need
    to load from checkpoint.

    Args:
        model: PyTorch model (should be an ExportWrapper or compatible).
        output_path: Path for output ONNX file.
        input_shape: Input tensor shape (B, T, C, H, W).
        config: Export configuration.

    Returns:
        ONNXExportResult with export status.

    Example:
        >>> from badas.export import export_model_to_onnx
        >>> from badas.export.wrappers import ExportWrapper
        >>>
        >>> wrapper = ExportWrapper(model)
        >>> result = export_model_to_onnx(
        ...     wrapper,
        ...     "model.onnx",
        ...     input_shape=(1, 16, 3, 256, 256),
        ... )
    """
    config = config or ONNXExportConfig()

    try:
        model = model.float().eval()
        dummy_input = torch.randn(*input_shape, dtype=torch.float32)

        with torch.no_grad():
            expected_output = model(dummy_input)

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

        dynamic_axes = None
        if config.dynamic_batch:
            dynamic_axes = {
                'pixel_values_videos': {0: 'batch_size'},
                'logits': {0: 'batch_size'},
            }

        torch.onnx.export(
            model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=config.opset_version,
            do_constant_folding=True,
            input_names=['pixel_values_videos'],
            output_names=['logits'],
            dynamic_axes=dynamic_axes,
            verbose=False,
        )

        if config.embed_weights:
            embed_onnx_weights(str(output_path))

        file_size_mb = output_path.stat().st_size / 1024 / 1024

        validation = None
        if config.validate:
            validation = validate_onnx(
                str(output_path),
                input_data=dummy_input,
                expected_output=expected_output,
            )

        return ONNXExportResult(
            success=True,
            onnx_path=str(output_path),
            file_size_mb=file_size_mb,
            validation=validation,
        )

    except Exception as e:
        return ONNXExportResult(
            success=False,
            error=str(e),
        )


__all__ = [
    'ONNXExportConfig',
    'ONNXExportResult',
    'export_to_onnx',
    'export_model_to_onnx',
]
