"""Export trained Essence Wars agents to ONNX format.

This module converts PyTorch checkpoints (.pt) to ONNX models (.onnx)
that can be loaded by the Rust game engine for use in the desktop app.

The ONNX model is self-contained: observation normalization is baked into
the graph, so no sidecar metadata files are needed. Just drop the .onnx
file into your agents directory and it works.

Usage:
    # From Python
    from essence_wars.agents.export_onnx import export_checkpoint

    export_checkpoint("best_model.pt", "my_agent.onnx")

    # From CLI
    python -m essence_wars.agents.export_onnx best_model.pt my_agent.onnx
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
import torch
import torch.nn as nn

from essence_wars._core import ACTION_SPACE_SIZE, STATE_TENSOR_SIZE


class _NormalizedInferenceWrapper(nn.Module):
    """Wraps a network with baked-in observation normalization.

    This produces a self-contained ONNX model: obs in -> (policy_logits, value) out.
    The normalizer mean/var are frozen as buffers in the graph.
    """

    def __init__(
        self,
        network: nn.Module,
        obs_mean: torch.Tensor | None = None,
        obs_var: torch.Tensor | None = None,
        epsilon: float = 1e-8,
    ) -> None:
        super().__init__()
        self.network = network

        if obs_mean is not None and obs_var is not None:
            self.register_buffer("obs_mean", obs_mean)
            self.register_buffer("obs_var", obs_var)
            self.epsilon = epsilon
            self.has_normalizer = True
        else:
            self.has_normalizer = False

    def forward(
        self, obs: torch.Tensor, action_mask: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if self.has_normalizer:
            obs = (obs - self.obs_mean) / torch.sqrt(self.obs_var + self.epsilon)

        logits, value = self.network(obs, action_mask)
        return logits, value


def _load_checkpoint(checkpoint_path: str | Path) -> dict[str, Any]:
    """Load a PyTorch checkpoint and return the dict."""
    path = Path(checkpoint_path)
    if not path.exists():
        raise FileNotFoundError(f"Checkpoint not found: {path}")

    checkpoint: dict[str, Any] = torch.load(
        path, map_location="cpu", weights_only=False
    )

    if "network_state_dict" not in checkpoint:
        raise ValueError(
            f"Invalid checkpoint: missing 'network_state_dict'. "
            f"Keys found: {list(checkpoint.keys())}"
        )

    return checkpoint


def _reconstruct_network(checkpoint: dict[str, Any]) -> nn.Module:
    """Reconstruct the network architecture from checkpoint config."""
    from essence_wars.agents.embeddings import create_network

    config = checkpoint.get("config")
    if config is None:
        # Fallback: assume flat PPO with default settings
        from essence_wars.agents.networks import EssenceWarsNetwork

        network = EssenceWarsNetwork()
    else:
        observation_mode = getattr(config, "observation_mode", "flat")
        hidden_dim = getattr(config, "hidden_dim", 256)
        embed_dim = getattr(config, "embed_dim", 64)
        include_embed_section = getattr(config, "include_embed_section", False)

        network = create_network(
            observation_mode=observation_mode,
            network_type="ppo",
            embed_dim=embed_dim,
            hidden_dim=hidden_dim,
            include_embed_section=include_embed_section,
        )

    network.load_state_dict(checkpoint["network_state_dict"])
    network.eval()
    return network


def export_checkpoint(
    checkpoint_path: str | Path,
    output_path: str | Path,
    *,
    verbose: bool = True,
) -> Path:
    """Export a PyTorch checkpoint to ONNX format.

    The observation normalizer (if present) is baked into the ONNX graph,
    making the model fully self-contained.

    Args:
        checkpoint_path: Path to .pt checkpoint file
        output_path: Path for output .onnx file
        verbose: Print progress messages

    Returns:
        Path to the exported .onnx file
    """
    output_path = Path(output_path)

    if verbose:
        print(f"Loading checkpoint: {checkpoint_path}")

    checkpoint = _load_checkpoint(checkpoint_path)
    network = _reconstruct_network(checkpoint)

    # Extract normalizer state if present
    obs_mean = None
    obs_var = None
    if "obs_normalizer" in checkpoint:
        norm_state = checkpoint["obs_normalizer"]
        obs_mean = torch.tensor(norm_state["mean"], dtype=torch.float32)
        obs_var = torch.tensor(norm_state["var"], dtype=torch.float32)
        if verbose:
            print("  Baking observation normalizer into ONNX graph")

    # Wrap network with normalizer
    wrapper = _NormalizedInferenceWrapper(network, obs_mean, obs_var)
    wrapper.eval()

    # Create dummy inputs
    dummy_obs = torch.randn(1, STATE_TENSOR_SIZE)
    dummy_mask = torch.ones(1, ACTION_SPACE_SIZE, dtype=torch.bool)

    # Export to ONNX
    if verbose:
        print(f"Exporting to ONNX: {output_path}")

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

    torch.onnx.export(
        wrapper,
        (dummy_obs, dummy_mask),
        str(output_path),
        input_names=["obs", "action_mask"],
        output_names=["policy_logits", "value"],
        dynamic_shapes={
            "obs": {0: (batch_dim := torch.export.Dim("batch"))},
            "action_mask": {0: batch_dim},
        },
        opset_version=17,
        dynamo=True,
        external_data=False,  # Embed all weights in a single file
    )

    # Verify the export
    file_size = output_path.stat().st_size
    if verbose:
        config = checkpoint.get("config")
        obs_mode = getattr(config, "observation_mode", "flat") if config else "flat"
        hidden_dim = getattr(config, "hidden_dim", 256) if config else 256
        has_norm = "obs_normalizer" in checkpoint
        step = checkpoint.get("global_step", "unknown")

        print(f"  Architecture: {obs_mode}, hidden_dim={hidden_dim}")
        print(f"  Normalizer baked in: {has_norm}")
        print(f"  Training steps: {step}")
        print(f"  File size: {file_size / 1024:.1f} KB")
        print(f"Export complete: {output_path}")

    return output_path


def verify_onnx(
    onnx_path: str | Path,
    checkpoint_path: str | Path | None = None,
) -> bool:
    """Verify an ONNX model produces correct outputs.

    Optionally compares against the original PyTorch checkpoint.

    Args:
        onnx_path: Path to .onnx file
        checkpoint_path: Optional path to original .pt for comparison

    Returns:
        True if verification passes
    """
    try:
        import onnxruntime as ort
    except ImportError:
        print("onnxruntime not installed, skipping verification")
        print("  Install with: pip install onnxruntime")
        return True

    onnx_path = Path(onnx_path)
    print(f"Verifying ONNX model: {onnx_path}")

    # Load and run ONNX model
    session = ort.InferenceSession(str(onnx_path))

    # Check input/output names
    inputs = {inp.name: inp.shape for inp in session.get_inputs()}
    outputs = {out.name: out.shape for out in session.get_outputs()}
    print(f"  Inputs: {inputs}")
    print(f"  Outputs: {outputs}")

    assert "obs" in inputs, "Missing 'obs' input"
    assert "action_mask" in inputs, "Missing 'action_mask' input"
    assert "policy_logits" in outputs, "Missing 'policy_logits' output"
    assert "value" in outputs, "Missing 'value' output"

    # Run inference with random input
    obs = np.random.randn(1, STATE_TENSOR_SIZE).astype(np.float32)
    mask = np.ones((1, ACTION_SPACE_SIZE), dtype=np.bool_)

    onnx_outputs = session.run(None, {"obs": obs, "action_mask": mask})
    policy_logits = onnx_outputs[0]
    value = onnx_outputs[1]

    print(f"  Policy logits shape: {policy_logits.shape}")
    print(f"  Value shape: {value.shape}")
    assert policy_logits.shape == (1, ACTION_SPACE_SIZE)
    assert value.shape == (1,)

    # Compare with PyTorch if checkpoint provided
    if checkpoint_path is not None:
        checkpoint = _load_checkpoint(checkpoint_path)
        network = _reconstruct_network(checkpoint)

        # Apply normalizer manually
        obs_t = torch.tensor(obs, dtype=torch.float32)
        if "obs_normalizer" in checkpoint:
            norm = checkpoint["obs_normalizer"]
            mean = torch.tensor(norm["mean"], dtype=torch.float32)
            var = torch.tensor(norm["var"], dtype=torch.float32)
            obs_t = (obs_t - mean) / torch.sqrt(var + 1e-8)

        mask_t = torch.tensor(mask, dtype=torch.bool)
        with torch.no_grad():
            pt_logits, pt_value = network(obs_t, mask_t)

        # Check numerical closeness
        logit_diff = np.abs(policy_logits - pt_logits.numpy()).max()
        value_diff = abs(float(value[0]) - float(pt_value[0]))
        print(f"  Max logit difference: {logit_diff:.6f}")
        print(f"  Value difference: {value_diff:.6f}")

        if logit_diff > 1e-4 or value_diff > 1e-4:
            print("  WARNING: Outputs differ significantly!")
            return False

        print("  PyTorch/ONNX outputs match!")

    print("Verification passed!")
    return True


def main() -> None:
    """CLI entry point for ONNX export."""
    import argparse

    parser = argparse.ArgumentParser(
        description="Export trained Essence Wars agent to ONNX format",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Export a trained PPO checkpoint
  python -m essence_wars.agents.export_onnx best_model.pt my_agent.onnx

  # Export and verify against original
  python -m essence_wars.agents.export_onnx best_model.pt my_agent.onnx --verify
        """,
    )
    parser.add_argument("checkpoint", help="Path to .pt checkpoint file")
    parser.add_argument("output", help="Path for output .onnx file")
    parser.add_argument(
        "--verify",
        action="store_true",
        help="Verify ONNX output matches PyTorch (requires onnxruntime)",
    )

    args = parser.parse_args()

    output = export_checkpoint(args.checkpoint, args.output)

    if args.verify:
        verify_onnx(output, args.checkpoint)


if __name__ == "__main__":
    main()
