"""
Framework Converters - Convert any ML model to/from OsmoFormat.

Supported frameworks:
- PyTorch (torch.nn.Module, state_dict)
- TensorFlow/Keras (tf.keras.Model, SavedModel)
- JAX/Flax (params dict)
- ONNX (onnx.ModelProto)
- HuggingFace (transformers models)

Usage:
    # Convert PyTorch to OsmoFormat
    osmo_model = pytorch_to_osmo(my_torch_model)

    # Convert OsmoFormat back to PyTorch
    torch_model = osmo_to_pytorch(osmo_model)

    # Or use the universal converter
    converter = UniversalConverter()
    osmo = converter.to_osmo(any_model)
    result = converter.from_osmo(osmo, target_framework="tensorflow")
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime
import hashlib
import json
import logging
from typing import Any, Dict, List, Optional, Tuple, Type, Union
import re

logger = logging.getLogger("model_osmosis.converters")

import numpy as np

from model_osmosis.format.osmo_format import (
    OsmoModel,
    OsmoManifest,
    ArchitectureSpec,
    LayerSpec,
    LayerType,
    ModelModality,
    FrameworkOrigin,
    ModelCapabilities,
    TrainingInfo,
    OSMO_FORMAT_VERSION,
)


TENSOR_FIDELITY_AUX_KEY = "tensor_fidelity.json"
TENSOR_FIDELITY_VERSION = "1"


@dataclass
class ConversionReport:
    """Conversion diagnostics exposed after to_osmo/from_osmo calls."""

    missing_keys: List[str] = field(default_factory=list)
    unexpected_keys: List[str] = field(default_factory=list)
    lossy_casts: List[str] = field(default_factory=list)
    dtype_transports: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "missing_keys": list(self.missing_keys),
            "unexpected_keys": list(self.unexpected_keys),
            "lossy_casts": list(self.lossy_casts),
            "dtype_transports": list(self.dtype_transports),
        }


def _empty_report() -> ConversionReport:
    return ConversionReport()


def _sha256_bytes(raw: bytes) -> str:
    return hashlib.sha256(raw).hexdigest()


def _array_order(arr: np.ndarray) -> str:
    if arr.flags.f_contiguous and not arr.flags.c_contiguous:
        return "F"
    return "C"


def _as_dict_copy(value: Any) -> Dict[str, Any]:
    if isinstance(value, dict):
        return dict(value)
    return {}


def _serialize_tensor_fidelity(fidelity: Dict[str, Any]) -> Dict[str, Any]:
    return {
        "version": fidelity.get("version", TENSOR_FIDELITY_VERSION),
        "tensors": _as_dict_copy(fidelity.get("tensors")),
    }


def _extract_tensor_fidelity(osmo: OsmoModel) -> Dict[str, Any]:
    raw = osmo.auxiliary_files.get(TENSOR_FIDELITY_AUX_KEY)
    if raw is None:
        return {"version": TENSOR_FIDELITY_VERSION, "tensors": {}}
    if isinstance(raw, str):
        try:
            raw = json.loads(raw)
        except json.JSONDecodeError:
            return {"version": TENSOR_FIDELITY_VERSION, "tensors": {}}
    if not isinstance(raw, dict):
        return {"version": TENSOR_FIDELITY_VERSION, "tensors": {}}
    tensors = raw.get("tensors", {})
    if not isinstance(tensors, dict):
        tensors = {}
    return {
        "version": str(raw.get("version", TENSOR_FIDELITY_VERSION)),
        "tensors": tensors,
    }


def _format_strict_load_error(
    framework: str,
    missing_keys: List[str],
    unexpected_keys: List[str],
) -> str:
    missing_preview = ", ".join(missing_keys[:5])
    unexpected_preview = ", ".join(unexpected_keys[:5])
    return (
        f"Strict {framework} load failed: missing_keys={len(missing_keys)} "
        f"(examples: [{missing_preview}]), unexpected_keys={len(unexpected_keys)} "
        f"(examples: [{unexpected_preview}])."
    )


def _normalize_task_set(tasks: Optional[List[str]]) -> set[str]:
    return {
        str(task).strip().lower()
        for task in (tasks or [])
        if str(task).strip()
    }


def _select_hf_model_family(tasks: Optional[List[str]]) -> str:
    normalized = _normalize_task_set(tasks)
    if "text-generation" in normalized:
        return "causal_lm"
    if "text2text-generation" in normalized:
        return "seq2seq_lm"
    if "fill-mask" in normalized:
        return "masked_lm"
    return "base"


def _hf_model_from_config(auto_cls: Any, config: Any) -> Any:
    try:
        return auto_cls.from_config(config, trust_remote_code=True)
    except TypeError:
        return auto_cls.from_config(config)


def _build_torch_state_dict_from_osmo(
    osmo: OsmoModel,
    *,
    allow_lossy_cast: bool,
    report: ConversionReport,
) -> Dict[str, Any]:
    import torch

    fidelity = _extract_tensor_fidelity(osmo)
    fidelity_tensors = _as_dict_copy(fidelity.get("tensors"))
    state_dict: Dict[str, Any] = {}

    for name in sorted(osmo.parameters.keys()):
        arr = osmo.parameters[name]
        meta = _as_dict_copy(fidelity_tensors.get(name))
        encoding_mode = str(meta.get("encoding_mode", "numpy"))

        if encoding_mode == "raw_bytes":
            original_dtype = str(meta.get("original_dtype", ""))
            if original_dtype != "bfloat16":
                raise ValueError(
                    f"Unsupported raw-bytes tensor transport dtype '{original_dtype}' for {name}"
                )
            raw = np.asarray(arr, dtype=np.uint8).tobytes(order="C")
            expected_raw_len = int(meta.get("raw_len", len(raw)))
            if expected_raw_len != len(raw):
                raise ValueError(
                    f"Tensor fidelity length mismatch for {name}: "
                    f"expected={expected_raw_len}, got={len(raw)}"
                )
            checksum = str(meta.get("checksum", ""))
            actual_checksum = _sha256_bytes(raw)
            if checksum and checksum != actual_checksum:
                raise ValueError(
                    f"Tensor fidelity checksum mismatch for {name}: "
                    f"expected={checksum}, got={actual_checksum}"
                )

            shape = tuple(meta.get("shape", []))
            order = str(meta.get("order", "C"))
            bits = np.frombuffer(raw, dtype=np.uint16).copy()
            if shape:
                bits = bits.reshape(shape, order=order)
            bit_tensor = torch.from_numpy(bits.view(np.int16))
            try:
                state_dict[name] = bit_tensor.view(torch.bfloat16)
            except Exception as exc:
                raise ValueError(
                    f"Failed to reconstruct bfloat16 tensor for {name}: {exc}"
                ) from exc
            report.dtype_transports.append(name)
            continue

        if encoding_mode == "lossy_cast":
            if not allow_lossy_cast:
                raise ValueError(
                    f"Lossy cast tensor '{name}' is blocked by default. "
                    "Set allow_lossy_cast=True to permit degraded restoration."
                )
            report.lossy_casts.append(name)

        order = str(meta.get("order", _array_order(arr)))
        raw = arr.tobytes(order=order)
        checksum = str(meta.get("checksum", ""))
        actual_checksum = _sha256_bytes(raw)
        if checksum and checksum != actual_checksum:
            raise ValueError(
                f"Tensor fidelity checksum mismatch for {name}: "
                f"expected={checksum}, got={actual_checksum}"
            )
        state_dict[name] = torch.from_numpy(arr)

    return state_dict


def _infer_layer_type(name: str, shape: Tuple[int, ...]) -> LayerType:
    """Infer semantic layer type from parameter name and shape."""
    name_lower = name.lower()

    # Embedding layers
    if any(x in name_lower for x in ["embed", "wte", "wpe", "token"]):
        return LayerType.EMBEDDING

    # Attention layers
    if any(x in name_lower for x in ["attn", "attention", "q_proj", "k_proj", "v_proj", "query", "key", "value"]):
        return LayerType.ATTENTION

    # Feed-forward / MLP
    if any(x in name_lower for x in ["mlp", "ffn", "feed_forward", "fc1", "fc2", "dense", "intermediate"]):
        return LayerType.FFN

    # Normalization
    if any(x in name_lower for x in ["norm", "ln_", "layer_norm", "layernorm", "rmsnorm"]):
        return LayerType.NORMALIZATION

    # Convolution
    if any(x in name_lower for x in ["conv", "dwconv"]):
        return LayerType.CONVOLUTION

    # Pooling
    if "pool" in name_lower:
        return LayerType.POOLING

    # Output heads
    if any(x in name_lower for x in ["lm_head", "classifier", "head", "output"]):
        return LayerType.HEAD

    # Gating (MoE, GLU)
    if any(x in name_lower for x in ["gate", "router", "expert"]):
        return LayerType.GATE

    # Positional
    if any(x in name_lower for x in ["position", "pos_"]):
        return LayerType.POSITIONAL

    # Projection (generic linear)
    if len(shape) == 2:
        return LayerType.PROJECTION

    return LayerType.CUSTOM


def _infer_modality(config: Dict[str, Any], layer_names: List[str]) -> ModelModality:
    """Infer model modality from config and layer names."""
    names_str = " ".join(layer_names).lower()
    config_str = str(config).lower()

    # Vision indicators
    if any(x in names_str or x in config_str for x in ["patch_embed", "pixel", "image", "vit", "resnet", "conv2d"]):
        return ModelModality.VISION

    # Audio indicators
    if any(x in names_str or x in config_str for x in ["whisper", "wav2vec", "hubert", "audio", "mel"]):
        return ModelModality.AUDIO

    # Multimodal indicators
    if any(x in names_str or x in config_str for x in ["clip", "llava", "blip", "multimodal"]):
        return ModelModality.MULTIMODAL

    # Diffusion indicators
    if any(x in names_str or x in config_str for x in ["unet", "vae", "diffusion", "denoise"]):
        return ModelModality.DIFFUSION

    # Code indicators
    if any(x in config_str for x in ["code", "starcoder", "codegen", "codellama"]):
        return ModelModality.CODE

    # Embedding models
    if any(x in config_str for x in ["embedding", "sentence-transformer", "e5", "bge"]):
        return ModelModality.EMBEDDING

    # Default to text for transformer-like models
    if any(x in names_str for x in ["lm_head", "embed_tokens", "wte"]):
        return ModelModality.TEXT

    return ModelModality.UNKNOWN


class BaseConverter(ABC):
    """Base class for framework converters."""

    def __init__(self) -> None:
        self.last_report: ConversionReport = _empty_report()

    @property
    @abstractmethod
    def framework(self) -> FrameworkOrigin:
        """Return the framework this converter handles."""
        pass

    @abstractmethod
    def to_osmo(self, model: Any, name: str = "", **kwargs) -> OsmoModel:
        """Convert a model to OsmoFormat."""
        pass

    @abstractmethod
    def from_osmo(self, osmo: OsmoModel, **kwargs) -> Any:
        """Convert an OsmoModel back to the native framework."""
        pass

    @abstractmethod
    def can_convert(self, model: Any) -> bool:
        """Check if this converter can handle the given model."""
        pass


class PyTorchConverter(BaseConverter):
    """Convert PyTorch models to/from OsmoFormat."""

    @property
    def framework(self) -> FrameworkOrigin:
        return FrameworkOrigin.PYTORCH

    def can_convert(self, model: Any) -> bool:
        try:
            import torch
            return isinstance(model, (torch.nn.Module, dict))
        except ImportError:
            return False

    def to_osmo(
        self,
        model: Any,
        name: str = "",
        config: Optional[Dict] = None,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> OsmoModel:
        """
        Convert PyTorch model to OsmoFormat.

        Args:
            model: torch.nn.Module or state_dict
            name: Model name
            config: Optional model config dict

        Returns:
            OsmoModel
        """
        import torch

        # Get state dict
        if isinstance(model, torch.nn.Module):
            state_dict = model.state_dict()
            model_name = name or model.__class__.__name__
        else:
            state_dict = model
            model_name = name or "pytorch_model"

        # Convert to numpy with tensor-fidelity metadata.
        parameters: Dict[str, np.ndarray] = {}
        layers = []
        hidden_size = 0
        fidelity_tensors: Dict[str, Dict[str, Any]] = {}
        report = _empty_report()

        for param_name in sorted(state_dict.keys()):
            tensor = state_dict[param_name]
            if hasattr(tensor, "cpu"):
                tensor_cpu = tensor.detach().cpu() if hasattr(tensor, "detach") else tensor.cpu()
                torch_dtype = str(getattr(tensor_cpu, "dtype", "unknown"))
            else:
                tensor_cpu = np.asarray(tensor)
                torch_dtype = str(getattr(tensor_cpu, "dtype", "unknown"))
            encoding_mode = "numpy"

            try:
                arr = tensor_cpu.numpy() if hasattr(tensor_cpu, "numpy") else np.asarray(tensor_cpu)
            except TypeError:
                # Preserve exact tensor bytes by default (e.g. torch.bfloat16).
                if torch_dtype == "torch.bfloat16":
                    if allow_lossy_cast:
                        arr = tensor_cpu.to(dtype=torch.float32).numpy()
                        encoding_mode = "lossy_cast"
                        report.lossy_casts.append(param_name)
                    else:
                        raw_int16 = tensor_cpu.contiguous().view(torch.int16).numpy()
                        raw_bytes = raw_int16.view(np.uint16).tobytes(order="C")
                        arr = np.frombuffer(raw_bytes, dtype=np.uint8).copy()
                        encoding_mode = "raw_bytes"
                        report.dtype_transports.append(param_name)
                else:
                    if not allow_lossy_cast:
                        raise TypeError(
                            "Unsupported torch dtype conversion for "
                            f"{param_name} ({torch_dtype}). "
                            "Re-run with allow_lossy_cast=True to permit numeric casting."
                        )
                    if tensor_cpu.is_floating_point():
                        arr = tensor_cpu.to(dtype=torch.float32).numpy()
                    else:
                        arr = tensor_cpu.to(dtype=torch.int64).numpy()
                    encoding_mode = "lossy_cast"
                    report.lossy_casts.append(param_name)

            if not isinstance(arr, np.ndarray):
                arr = np.asarray(arr)

            parameters[param_name] = arr

            fidelity_shape = list(tensor_cpu.shape) if hasattr(tensor_cpu, "shape") else list(arr.shape)
            fidelity_order = "C"
            fidelity_raw = b""
            if encoding_mode == "raw_bytes":
                fidelity_raw = np.asarray(arr, dtype=np.uint8).tobytes(order="C")
            else:
                fidelity_order = _array_order(arr)
                fidelity_raw = arr.tobytes(order=fidelity_order)

            fidelity_tensors[param_name] = {
                "original_dtype": torch_dtype.replace("torch.", ""),
                "storage_dtype": str(arr.dtype),
                "shape": fidelity_shape,
                "storage_shape": list(arr.shape),
                "order": fidelity_order,
                "encoding_mode": encoding_mode,
                "raw_len": len(fidelity_raw),
                "checksum": _sha256_bytes(fidelity_raw),
            }

            # Create layer spec
            layer_type = _infer_layer_type(param_name, tuple(fidelity_shape))
            layers.append(
                LayerSpec(
                    name=param_name,
                    layer_type=layer_type,
                    shape=fidelity_shape,
                    dtype=fidelity_tensors[param_name]["original_dtype"],
                )
            )

            # Infer hidden size from embedding or attention layers
            if layer_type in (LayerType.EMBEDDING, LayerType.ATTENTION) and len(fidelity_shape) >= 2:
                hidden_size = max(hidden_size, fidelity_shape[-1])

        # Build config from actual model if available
        actual_config = config or {}
        if hasattr(model, "config"):
            try:
                actual_config = model.config.to_dict()
            except (AttributeError, TypeError):
                pass

        # Infer architecture details — use config first, then heuristic.
        # The heuristic (count // 4) assumes q,k,v,o per layer but fails for
        # GQA/MQA architectures where k,v may be shared/fewer.
        attn_count = len([l for l in layers if l.layer_type == LayerType.ATTENTION])
        num_layers = actual_config.get("num_hidden_layers", 0)
        if not num_layers and attn_count > 0:
            num_layers = max(attn_count // 4, 1)

        architecture = ArchitectureSpec(
            layers=layers,
            hidden_size=hidden_size or actual_config.get("hidden_size", 768),
            num_layers=num_layers or actual_config.get("num_hidden_layers", 12),
            num_attention_heads=actual_config.get("num_attention_heads"),
            vocab_size=actual_config.get("vocab_size"),
            max_position=actual_config.get("max_position_embeddings"),
            intermediate_size=actual_config.get("intermediate_size"),
            architecture_type=actual_config.get("model_type", "unknown"),
        )

        # Infer modality
        modality = _infer_modality(actual_config, list(parameters.keys()))

        # Build manifest
        manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name=model_name,
            modality=modality,
            framework_origin=FrameworkOrigin.PYTORCH,
            parameter_count=sum(p.size for p in parameters.values()),
            created_at=datetime.now().isoformat(),
        )

        osmo = OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=parameters,
        )
        osmo.auxiliary_files[TENSOR_FIDELITY_AUX_KEY] = _serialize_tensor_fidelity(
            {"version": TENSOR_FIDELITY_VERSION, "tensors": fidelity_tensors}
        )
        self.last_report = report
        return osmo

    def from_osmo(
        self,
        osmo: OsmoModel,
        model_class: Optional[Type] = None,
        config: Optional[Any] = None,
        strict_load: bool = True,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> Any:
        """
        Convert OsmoModel back to PyTorch.

        Args:
            osmo: The OsmoModel to convert
            model_class: Optional torch.nn.Module class to instantiate
            config: Optional config for model instantiation

        Returns:
            PyTorch state_dict or model instance
        """
        import torch

        report = _empty_report()
        state_dict = _build_torch_state_dict_from_osmo(
            osmo, allow_lossy_cast=allow_lossy_cast, report=report
        )

        # If model class provided, instantiate and load
        if model_class is not None:
            if config is not None:
                model = model_class(config)
            else:
                model = model_class()
            load_result = model.load_state_dict(state_dict, strict=False)
            missing = list(load_result.missing_keys)
            unexpected = list(load_result.unexpected_keys)
            report.missing_keys = missing
            report.unexpected_keys = unexpected
            if strict_load and (missing or unexpected):
                raise ValueError(_format_strict_load_error("PyTorch", missing, unexpected))
            self.last_report = report
            return model

        self.last_report = report
        return state_dict


class TensorFlowConverter(BaseConverter):
    """Convert TensorFlow/Keras models to/from OsmoFormat."""

    @property
    def framework(self) -> FrameworkOrigin:
        return FrameworkOrigin.TENSORFLOW

    def can_convert(self, model: Any) -> bool:
        try:
            import tensorflow as tf
            return isinstance(model, (tf.keras.Model, tf.Module))
        except ImportError:
            return False

    def to_osmo(
        self,
        model: Any,
        name: str = "",
        **kwargs,
    ) -> OsmoModel:
        """Convert TensorFlow model to OsmoFormat."""
        import tensorflow as tf

        model_name = name or model.name or "tf_model"

        # Get weights
        parameters: Dict[str, np.ndarray] = {}
        layers = []
        hidden_size = 0
        fidelity_tensors: Dict[str, Dict[str, Any]] = {}

        vars_sorted = sorted(model.trainable_variables, key=lambda v: v.name)
        for var in vars_sorted:
            param_name = var.name
            arr = var.numpy()
            parameters[param_name] = arr

            order = _array_order(arr)
            raw = arr.tobytes(order=order)
            fidelity_tensors[param_name] = {
                "original_dtype": str(arr.dtype),
                "storage_dtype": str(arr.dtype),
                "shape": list(arr.shape),
                "storage_shape": list(arr.shape),
                "order": order,
                "encoding_mode": "numpy",
                "raw_len": len(raw),
                "checksum": _sha256_bytes(raw),
            }

            layer_type = _infer_layer_type(param_name, arr.shape)
            layers.append(
                LayerSpec(
                    name=param_name,
                    layer_type=layer_type,
                    shape=list(arr.shape),
                    dtype=str(arr.dtype),
                )
            )

            if layer_type in (LayerType.EMBEDDING, LayerType.ATTENTION) and len(arr.shape) >= 2:
                hidden_size = max(hidden_size, arr.shape[-1])

        architecture = ArchitectureSpec(
            layers=layers,
            hidden_size=hidden_size or 768,
            num_layers=len([l for l in layers if l.layer_type == LayerType.ATTENTION]) // 4 or 1,
        )

        modality = _infer_modality({}, list(parameters.keys()))

        manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name=model_name,
            modality=modality,
            framework_origin=FrameworkOrigin.TENSORFLOW,
            parameter_count=sum(p.size for p in parameters.values()),
            created_at=datetime.now().isoformat(),
        )

        osmo = OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=parameters,
        )
        osmo.auxiliary_files[TENSOR_FIDELITY_AUX_KEY] = _serialize_tensor_fidelity(
            {"version": TENSOR_FIDELITY_VERSION, "tensors": fidelity_tensors}
        )
        self.last_report = _empty_report()
        return osmo

    def from_osmo(
        self,
        osmo: OsmoModel,
        model: Optional[Any] = None,
        strict_load: bool = True,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> Dict[str, np.ndarray]:
        """
        Convert OsmoModel to TensorFlow weights.

        Args:
            osmo: The OsmoModel
            model: Optional tf.keras.Model to load weights into

        Returns:
            Weight dict or model with loaded weights
        """
        report = _empty_report()
        fidelity = _extract_tensor_fidelity(osmo)
        fidelity_tensors = _as_dict_copy(fidelity.get("tensors"))

        if model is not None:
            target_names = [var.name for var in model.trainable_variables]
            source_names = sorted(osmo.parameters.keys())
            missing = [name for name in target_names if name not in osmo.parameters]
            unexpected = [name for name in source_names if name not in target_names]
            report.missing_keys = missing
            report.unexpected_keys = unexpected
            if strict_load and (missing or unexpected):
                raise ValueError(_format_strict_load_error("TensorFlow", missing, unexpected))

            for var in model.trainable_variables:
                if var.name in osmo.parameters:
                    arr = osmo.parameters[var.name]
                    meta = _as_dict_copy(fidelity_tensors.get(var.name))
                    if str(meta.get("encoding_mode", "numpy")) == "lossy_cast" and not allow_lossy_cast:
                        raise ValueError(
                            f"Lossy cast tensor '{var.name}' is blocked by default. "
                            "Set allow_lossy_cast=True to permit degraded restoration."
                        )
                    var.assign(arr)
            self.last_report = report
            return model

        self.last_report = report
        return osmo.parameters


class JAXConverter(BaseConverter):
    """Convert JAX/Flax models to/from OsmoFormat."""

    @property
    def framework(self) -> FrameworkOrigin:
        return FrameworkOrigin.JAX

    def can_convert(self, model: Any) -> bool:
        # JAX params are typically nested dicts with jax arrays
        if not isinstance(model, dict):
            return False
        try:
            import jax.numpy as jnp
            # Check if any leaf is a jax array
            def _has_jax_arrays(d):
                for v in d.values():
                    if isinstance(v, dict):
                        if _has_jax_arrays(v):
                            return True
                    elif hasattr(v, "device"):  # JAX arrays have device attribute
                        return True
                return False
            return _has_jax_arrays(model)
        except ImportError:
            return False

    def to_osmo(
        self,
        model: Dict,
        name: str = "",
        **kwargs,
    ) -> OsmoModel:
        """Convert JAX params dict to OsmoFormat."""
        import jax

        model_name = name or "jax_model"

        # Flatten nested params dict
        flat_params = {}

        def _flatten(d, prefix=""):
            for k in sorted(d.keys(), key=lambda item: str(item)):
                v = d[k]
                key = f"{prefix}/{k}" if prefix else k
                if isinstance(v, dict):
                    _flatten(v, key)
                else:
                    flat_params[key] = np.array(v)

        _flatten(model)

        # Build layers
        ordered_params = {k: flat_params[k] for k in sorted(flat_params.keys())}
        layers = []
        hidden_size = 0
        fidelity_tensors: Dict[str, Dict[str, Any]] = {}

        for param_name, arr in ordered_params.items():
            layer_type = _infer_layer_type(param_name, arr.shape)
            layers.append(
                LayerSpec(
                    name=param_name,
                    layer_type=layer_type,
                    shape=list(arr.shape),
                    dtype=str(arr.dtype),
                )
            )
            if layer_type in (LayerType.EMBEDDING, LayerType.ATTENTION) and len(arr.shape) >= 2:
                hidden_size = max(hidden_size, arr.shape[-1])

            order = _array_order(arr)
            raw = arr.tobytes(order=order)
            fidelity_tensors[param_name] = {
                "original_dtype": str(arr.dtype),
                "storage_dtype": str(arr.dtype),
                "shape": list(arr.shape),
                "storage_shape": list(arr.shape),
                "order": order,
                "encoding_mode": "numpy",
                "raw_len": len(raw),
                "checksum": _sha256_bytes(raw),
            }

        architecture = ArchitectureSpec(
            layers=layers,
            hidden_size=hidden_size or 768,
            num_layers=len([l for l in layers if l.layer_type == LayerType.ATTENTION]) // 4 or 1,
        )

        modality = _infer_modality({}, list(flat_params.keys()))

        manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name=model_name,
            modality=modality,
            framework_origin=FrameworkOrigin.JAX,
            parameter_count=sum(p.size for p in ordered_params.values()),
            created_at=datetime.now().isoformat(),
        )

        osmo = OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=ordered_params,
        )
        osmo.auxiliary_files[TENSOR_FIDELITY_AUX_KEY] = _serialize_tensor_fidelity(
            {"version": TENSOR_FIDELITY_VERSION, "tensors": fidelity_tensors}
        )
        self.last_report = _empty_report()
        return osmo

    def from_osmo(
        self,
        osmo: OsmoModel,
        strict_load: bool = True,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> Dict:
        """
        Convert OsmoModel to JAX params dict.

        Returns nested dict structure from flattened params.
        """
        import jax.numpy as jnp

        report = _empty_report()
        fidelity = _extract_tensor_fidelity(osmo)
        fidelity_tensors = _as_dict_copy(fidelity.get("tensors"))

        # Unflatten to nested dict
        result = {}
        for name in sorted(osmo.parameters.keys()):
            arr = osmo.parameters[name]
            meta = _as_dict_copy(fidelity_tensors.get(name))
            if str(meta.get("encoding_mode", "numpy")) == "lossy_cast" and not allow_lossy_cast:
                raise ValueError(
                    f"Lossy cast tensor '{name}' is blocked by default. "
                    "Set allow_lossy_cast=True to permit degraded restoration."
                )
            parts = name.split("/")
            current = result
            for part in parts[:-1]:
                if part not in current:
                    current[part] = {}
                current = current[part]
            current[parts[-1]] = jnp.array(arr)

        target_params = kwargs.get("target_params")
        if target_params is not None:
            expected: Dict[str, Any] = {}

            def _flatten_target(d, prefix=""):
                for key in sorted(d.keys(), key=lambda item: str(item)):
                    value = d[key]
                    full = f"{prefix}/{key}" if prefix else str(key)
                    if isinstance(value, dict):
                        _flatten_target(value, full)
                    else:
                        expected[full] = value

            if isinstance(target_params, dict):
                _flatten_target(target_params)
                missing = [k for k in sorted(expected.keys()) if k not in osmo.parameters]
                unexpected = [k for k in sorted(osmo.parameters.keys()) if k not in expected]
                report.missing_keys = missing
                report.unexpected_keys = unexpected
                if strict_load and (missing or unexpected):
                    raise ValueError(_format_strict_load_error("JAX", missing, unexpected))

        self.last_report = report
        return result


class ONNXConverter(BaseConverter):
    """Convert ONNX models to/from OsmoFormat."""

    @property
    def framework(self) -> FrameworkOrigin:
        return FrameworkOrigin.ONNX

    def can_convert(self, model: Any) -> bool:
        try:
            import onnx
            return isinstance(model, onnx.ModelProto)
        except ImportError:
            return False

    def to_osmo(
        self,
        model: Any,
        name: str = "",
        **kwargs,
    ) -> OsmoModel:
        """Convert ONNX model to OsmoFormat."""
        import onnx
        from onnx import numpy_helper

        model_name = name or model.graph.name or "onnx_model"

        # Extract initializers (weights)
        pairs = []
        for initializer in model.graph.initializer:
            arr = numpy_helper.to_array(initializer)
            pairs.append((initializer.name, arr))
        pairs.sort(key=lambda x: x[0])

        parameters: Dict[str, np.ndarray] = {}
        layers = []
        hidden_size = 0
        fidelity_tensors: Dict[str, Dict[str, Any]] = {}

        for param_name, arr in pairs:
            parameters[param_name] = arr

            layer_type = _infer_layer_type(param_name, arr.shape)
            layers.append(
                LayerSpec(
                    name=param_name,
                    layer_type=layer_type,
                    shape=list(arr.shape),
                    dtype=str(arr.dtype),
                )
            )

            if layer_type in (LayerType.EMBEDDING, LayerType.ATTENTION) and len(arr.shape) >= 2:
                hidden_size = max(hidden_size, arr.shape[-1])

            order = _array_order(arr)
            raw = arr.tobytes(order=order)
            fidelity_tensors[param_name] = {
                "original_dtype": str(arr.dtype),
                "storage_dtype": str(arr.dtype),
                "shape": list(arr.shape),
                "storage_shape": list(arr.shape),
                "order": order,
                "encoding_mode": "numpy",
                "raw_len": len(raw),
                "checksum": _sha256_bytes(raw),
            }

        architecture = ArchitectureSpec(
            layers=layers,
            hidden_size=hidden_size or 768,
            num_layers=len([l for l in layers if l.layer_type == LayerType.ATTENTION]) // 4 or 1,
        )

        modality = _infer_modality({}, list(parameters.keys()))

        manifest = OsmoManifest(
            version=OSMO_FORMAT_VERSION,
            name=model_name,
            modality=modality,
            framework_origin=FrameworkOrigin.ONNX,
            parameter_count=sum(p.size for p in parameters.values()),
            created_at=datetime.now().isoformat(),
        )

        osmo = OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=parameters,
        )
        osmo.auxiliary_files[TENSOR_FIDELITY_AUX_KEY] = _serialize_tensor_fidelity(
            {"version": TENSOR_FIDELITY_VERSION, "tensors": fidelity_tensors}
        )
        self.last_report = _empty_report()
        return osmo

    def from_osmo(
        self,
        osmo: OsmoModel,
        onnx_model: Optional[Any] = None,
        strict_load: bool = True,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> Any:
        """
        Update ONNX model weights from OsmoModel.

        Args:
            osmo: The OsmoModel
            onnx_model: Existing ONNX model to update weights in

        Returns:
            Updated ONNX model or weight dict
        """
        report = _empty_report()
        fidelity = _extract_tensor_fidelity(osmo)
        fidelity_tensors = _as_dict_copy(fidelity.get("tensors"))
        if onnx_model is None:
            self.last_report = report
            return osmo.parameters

        import onnx
        from onnx import numpy_helper

        expected_names = [initializer.name for initializer in onnx_model.graph.initializer]
        source_names = sorted(osmo.parameters.keys())
        missing = [name for name in expected_names if name not in osmo.parameters]
        unexpected = [name for name in source_names if name not in expected_names]
        report.missing_keys = missing
        report.unexpected_keys = unexpected
        if strict_load and (missing or unexpected):
            raise ValueError(_format_strict_load_error("ONNX", missing, unexpected))

        # Update initializers
        for i, initializer in enumerate(onnx_model.graph.initializer):
            if initializer.name in osmo.parameters:
                meta = _as_dict_copy(fidelity_tensors.get(initializer.name))
                if str(meta.get("encoding_mode", "numpy")) == "lossy_cast" and not allow_lossy_cast:
                    raise ValueError(
                        f"Lossy cast tensor '{initializer.name}' is blocked by default. "
                        "Set allow_lossy_cast=True to permit degraded restoration."
                    )
                new_tensor = numpy_helper.from_array(
                    osmo.parameters[initializer.name],
                    name=initializer.name,
                )
                onnx_model.graph.initializer[i].CopyFrom(new_tensor)

        self.last_report = report
        return onnx_model


class HuggingFaceConverter(BaseConverter):
    """Convert HuggingFace transformers models to/from OsmoFormat."""

    @property
    def framework(self) -> FrameworkOrigin:
        return FrameworkOrigin.PYTORCH  # HF uses PyTorch under the hood

    def can_convert(self, model: Any) -> bool:
        try:
            from transformers import PreTrainedModel
            return isinstance(model, PreTrainedModel)
        except ImportError:
            return False

    def to_osmo(
        self,
        model: Any,
        name: str = "",
        include_tokenizer: bool = True,
        include_config: bool = True,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> OsmoModel:
        """Convert HuggingFace model to OsmoFormat.

        Args:
            model: HuggingFace PreTrainedModel.
            name: Optional model name.
            include_tokenizer: Capture tokenizer files in auxiliary_files.
            include_config: Capture config/preprocessor files in auxiliary_files.
        """
        from transformers import PreTrainedModel

        # Use PyTorch converter but add HF-specific metadata
        pytorch_converter = PyTorchConverter()

        config_dict = {}
        if hasattr(model, "config"):
            try:
                config_dict = model.config.to_dict()
            except (AttributeError, TypeError):
                pass

        model_name = name or config_dict.get("_name_or_path", model.__class__.__name__)

        osmo = pytorch_converter.to_osmo(
            model,
            name=model_name,
            config=config_dict,
            allow_lossy_cast=allow_lossy_cast,
        )

        # Enhance manifest with HF-specific info
        osmo.manifest.original_repo = config_dict.get("_name_or_path")
        osmo.manifest.capabilities.tasks = self._infer_tasks(config_dict, model=model)
        osmo.manifest.capabilities.context_length = config_dict.get("max_position_embeddings")

        # Capture auxiliary files (tokenizer, config, preprocessor)
        auxiliary = {}
        repo_id = config_dict.get("_name_or_path", "")

        if include_config and config_dict:
            auxiliary["config.json"] = config_dict

            # Capture generation config if available
            if hasattr(model, "generation_config"):
                try:
                    auxiliary["generation_config.json"] = model.generation_config.to_dict()
                except Exception as e:
                    logger.debug("Could not capture generation_config: %s", e)

        if include_tokenizer and repo_id:
            auxiliary.update(self._capture_tokenizer_files(repo_id))

        merged_aux = dict(osmo.auxiliary_files or {})
        merged_aux.update(auxiliary)
        osmo.auxiliary_files = merged_aux
        self.last_report = _empty_report()
        return osmo

    def _capture_tokenizer_files(self, repo_id: str) -> Dict[str, Any]:
        """Try to load tokenizer files from a HuggingFace repo.

        Returns a dict of filename -> content for each captured file.
        """
        import json as _json

        captured: Dict[str, Any] = {}

        try:
            from transformers import AutoTokenizer
            tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)

            # Get the tokenizer config
            try:
                tok_config = _json.loads(tokenizer.backend_tokenizer.to_str())
                captured["tokenizer.json"] = tok_config
            except Exception as e:
                logger.debug("Could not capture tokenizer.json: %s", e)

            # tokenizer_config.json
            try:
                tok_init = tokenizer.init_kwargs.copy() if hasattr(tokenizer, "init_kwargs") else {}
                tok_init["tokenizer_class"] = tokenizer.__class__.__name__
                captured["tokenizer_config.json"] = tok_init
            except Exception as e:
                logger.debug("Could not capture tokenizer_config.json: %s", e)

            # special_tokens_map
            try:
                special = tokenizer.special_tokens_map
                captured["special_tokens_map.json"] = special
            except Exception as e:
                logger.debug("Could not capture special_tokens_map.json: %s", e)

        except Exception as e:
            logger.debug("Could not load tokenizer from %s: %s", repo_id, e)

        return captured

    def _infer_tasks(self, config: Dict, model: Optional[Any] = None) -> List[str]:
        """Infer tasks from model config."""
        tasks = []
        model_type = str(config.get("model_type", "") or "").lower()
        config_blob = str(config)
        architectures = config.get("architectures", [])
        if isinstance(architectures, list):
            arch_blob = " ".join(str(item) for item in architectures).lower()
        else:
            arch_blob = str(architectures).lower()
        model_class = model.__class__.__name__.lower() if model is not None else ""

        # Treat generation tasks conservatively: only when LM-specific hints exist.
        # This avoids misclassifying base backbones (e.g. GPT2Model) as causal LM.
        causal_hints = ("forcausallm", "lmheadmodel", "causallm")
        if (
            "lm_head" in config_blob
            or any(hint in arch_blob for hint in causal_hints)
            or any(hint in model_class for hint in causal_hints)
        ):
            tasks.append("text-generation")
        if "classifier" in config_blob:
            tasks.append("classification")
        if model_type in ["bert", "roberta"]:
            tasks.extend(["fill-mask", "classification"])
        if model_type in ["t5", "bart"]:
            tasks.append("text2text-generation")
        if model_type in ["whisper"]:
            tasks.append("speech-recognition")
        if model_type in ["clip", "blip"]:
            tasks.extend(["image-classification", "image-text-matching"])

        return tasks

    def from_osmo(
        self,
        osmo: OsmoModel,
        model_id: Optional[str] = None,
        strict_load: bool = True,
        allow_lossy_cast: bool = False,
        **kwargs,
    ) -> Any:
        """
        Convert OsmoModel to HuggingFace model.

        Args:
            osmo: The OsmoModel
            model_id: HuggingFace model ID to use as base architecture

        Returns:
            HuggingFace PreTrainedModel with loaded weights
        """
        from transformers import AutoModel, AutoConfig

        if model_id is None:
            model_id = osmo.manifest.original_repo

        if model_id is None:
            raise ValueError("Need model_id to reconstruct HuggingFace model")

        # Load config and create model
        config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
        model_family = _select_hf_model_family(osmo.manifest.capabilities.tasks)
        model = None

        if model_family == "causal_lm":
            try:
                from transformers import AutoModelForCausalLM
                model = _hf_model_from_config(AutoModelForCausalLM, config)
            except Exception:
                logger.debug(
                    "Could not initialize AutoModelForCausalLM for %s; falling back to AutoModel",
                    model_id,
                    exc_info=True,
                )
        elif model_family == "seq2seq_lm":
            try:
                from transformers import AutoModelForSeq2SeqLM
                model = _hf_model_from_config(AutoModelForSeq2SeqLM, config)
            except Exception:
                logger.debug(
                    "Could not initialize AutoModelForSeq2SeqLM for %s; falling back to AutoModel",
                    model_id,
                    exc_info=True,
                )
        elif model_family == "masked_lm":
            try:
                from transformers import AutoModelForMaskedLM
                model = _hf_model_from_config(AutoModelForMaskedLM, config)
            except Exception:
                logger.debug(
                    "Could not initialize AutoModelForMaskedLM for %s; falling back to AutoModel",
                    model_id,
                    exc_info=True,
                )

        if model is None:
            model = _hf_model_from_config(AutoModel, config)

        report = _empty_report()
        state_dict = _build_torch_state_dict_from_osmo(
            osmo, allow_lossy_cast=allow_lossy_cast, report=report
        )
        load_result = model.load_state_dict(state_dict, strict=False)
        missing = list(load_result.missing_keys)
        unexpected = list(load_result.unexpected_keys)
        report.missing_keys = missing
        report.unexpected_keys = unexpected
        if strict_load and (missing or unexpected):
            raise ValueError(_format_strict_load_error("HuggingFace", missing, unexpected))
        self.last_report = report
        return model


class UniversalConverter:
    """
    Universal converter that auto-detects framework and converts to/from OsmoFormat.

    This is the main entry point for converting any model.
    """

    def __init__(self):
        self.last_report: Dict[str, Any] = _empty_report().to_dict()
        self.converters: List[BaseConverter] = [
            HuggingFaceConverter(),  # Check HF first (most specific)
            PyTorchConverter(),
            TensorFlowConverter(),
            JAXConverter(),
            ONNXConverter(),
        ]

    def to_osmo(self, model: Any, name: str = "", **kwargs) -> OsmoModel:
        """
        Convert any supported model to OsmoFormat.

        Args:
            model: Any supported ML model
            name: Optional model name

        Returns:
            OsmoModel

        Raises:
            ValueError: If no converter supports this model type
        """
        for converter in self.converters:
            if converter.can_convert(model):
                osmo = converter.to_osmo(model, name=name, **kwargs)
                self.last_report = converter.last_report.to_dict()
                return osmo

        raise ValueError(
            f"No converter found for model type: {type(model)}. "
            f"Supported frameworks: PyTorch, TensorFlow, JAX, ONNX, HuggingFace"
        )

    def from_osmo(
        self,
        osmo: OsmoModel,
        target_framework: str = "pytorch",
        **kwargs,
    ) -> Any:
        """
        Convert OsmoModel to target framework.

        Args:
            osmo: The OsmoModel to convert
            target_framework: One of "pytorch", "tensorflow", "jax", "onnx", "huggingface"

        Returns:
            Model in target framework format
        """
        framework_map = {
            "pytorch": PyTorchConverter,
            "torch": PyTorchConverter,
            "tensorflow": TensorFlowConverter,
            "tf": TensorFlowConverter,
            "keras": TensorFlowConverter,
            "jax": JAXConverter,
            "flax": JAXConverter,
            "onnx": ONNXConverter,
            "huggingface": HuggingFaceConverter,
            "hf": HuggingFaceConverter,
        }

        target = target_framework.lower()
        if target not in framework_map:
            raise ValueError(
                f"Unknown target framework: {target_framework}. "
                f"Supported: {list(framework_map.keys())}"
            )

        converter = framework_map[target]()
        result = converter.from_osmo(osmo, **kwargs)
        self.last_report = converter.last_report.to_dict()
        return result

    def convert(
        self,
        model: Any,
        target_framework: str,
        name: str = "",
        **kwargs,
    ) -> Any:
        """
        Direct conversion between frameworks via OsmoFormat.

        Args:
            model: Source model (any supported framework)
            target_framework: Target framework name
            name: Optional model name

        Returns:
            Model in target framework format
        """
        osmo = self.to_osmo(model, name=name, **kwargs)
        return self.from_osmo(osmo, target_framework=target_framework, **kwargs)


# Convenience functions
def to_osmo(model: Any, name: str = "", **kwargs) -> OsmoModel:
    """Convert any model to OsmoFormat."""
    return UniversalConverter().to_osmo(model, name=name, **kwargs)


def from_osmo(osmo: OsmoModel, target_framework: str = "pytorch", **kwargs) -> Any:
    """Convert OsmoModel to target framework."""
    return UniversalConverter().from_osmo(osmo, target_framework=target_framework, **kwargs)
