"""
OsmoFormat - Universal Model Interchange Format

The .osmo format is a universal intermediate representation that allows
any ML model to be converted, merged, stacked, and transferred regardless
of its original framework (PyTorch, TensorFlow, JAX, ONNX, etc.).

File Structure:
    model.osmo/
    ├── manifest.json          # Metadata, version, capabilities
    ├── architecture.json      # Layer structure, connections, shapes
    ├── parameters/
    │   ├── layer_0.safetensors
    │   ├── layer_1.safetensors
    │   └── ...
    ├── knowledge/
    │   ├── structural.npz     # Architecture DNA
    │   ├── functional.npz     # Learned representations
    │   └── meta.json          # Training info, dataset hints
    └── fingerprint.json       # Model identity hash for deduplication

This format enables:
- Framework-agnostic model storage
- Intelligent merging based on layer semantics
- Cross-architecture knowledge transfer
- Model capability composition
"""

from dataclasses import dataclass, field, asdict
from enum import Enum
from pathlib import Path
from typing import Any, Dict, List, Optional, Union
import hashlib
import json
import numpy as np


class ModelModality(Enum):
    """Primary modality of the model."""
    TEXT = "text"
    VISION = "vision"
    AUDIO = "audio"
    VIDEO = "video"
    MULTIMODAL = "multimodal"
    CODE = "code"
    EMBEDDING = "embedding"
    DIFFUSION = "diffusion"
    UNKNOWN = "unknown"


class LayerType(Enum):
    """Semantic layer types for intelligent merging."""
    EMBEDDING = "embedding"          # Token/patch embeddings
    ATTENTION = "attention"          # Self/cross attention
    FFN = "ffn"                      # Feed-forward networks
    NORMALIZATION = "normalization"  # LayerNorm, RMSNorm, etc.
    CONVOLUTION = "convolution"      # Conv layers
    POOLING = "pooling"              # Pooling layers
    PROJECTION = "projection"        # Linear projections
    HEAD = "head"                    # Output heads (LM head, classifier)
    GATE = "gate"                    # Gating mechanisms (MoE, GLU)
    POSITIONAL = "positional"        # Position encodings
    CUSTOM = "custom"                # Unknown/custom layers


class FrameworkOrigin(Enum):
    """Original framework the model came from."""
    PYTORCH = "pytorch"
    TENSORFLOW = "tensorflow"
    JAX = "jax"
    ONNX = "onnx"
    MLX = "mlx"
    KERAS = "keras"
    FLAX = "flax"
    PADDLEPADDLE = "paddlepaddle"
    MXNET = "mxnet"
    UNKNOWN = "unknown"


@dataclass
class LayerSpec:
    """Specification for a single layer in the architecture."""
    name: str                           # Unique layer identifier
    layer_type: LayerType               # Semantic type
    shape: List[int]                    # Parameter shape
    dtype: str                          # Data type (float16, float32, etc.)
    inputs: List[str] = field(default_factory=list)   # Input layer names
    outputs: List[str] = field(default_factory=list)  # Output layer names
    attributes: Dict[str, Any] = field(default_factory=dict)  # Extra attributes

    def to_dict(self) -> Dict:
        d = asdict(self)
        d["layer_type"] = self.layer_type.value
        return d

    @classmethod
    def from_dict(cls, d: Dict) -> "LayerSpec":
        d = d.copy()
        d["layer_type"] = LayerType(d["layer_type"])
        return cls(**d)


@dataclass
class ArchitectureSpec:
    """Complete architecture specification."""
    layers: List[LayerSpec]
    hidden_size: int
    num_layers: int
    num_attention_heads: Optional[int] = None
    vocab_size: Optional[int] = None
    max_position: Optional[int] = None
    intermediate_size: Optional[int] = None
    architecture_type: str = "unknown"  # "transformer", "cnn", "rnn", etc.
    extra: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict:
        return {
            "layers": [l.to_dict() for l in self.layers],
            "hidden_size": self.hidden_size,
            "num_layers": self.num_layers,
            "num_attention_heads": self.num_attention_heads,
            "vocab_size": self.vocab_size,
            "max_position": self.max_position,
            "intermediate_size": self.intermediate_size,
            "architecture_type": self.architecture_type,
            "extra": self.extra,
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "ArchitectureSpec":
        d = d.copy()
        d["layers"] = [LayerSpec.from_dict(l) for l in d["layers"]]
        return cls(**d)


@dataclass
class ModelCapabilities:
    """What the model can do."""
    tasks: List[str] = field(default_factory=list)  # ["text-generation", "classification"]
    languages: List[str] = field(default_factory=list)  # ["en", "fr", "multi"]
    domains: List[str] = field(default_factory=list)  # ["code", "medical", "legal"]
    context_length: Optional[int] = None
    output_types: List[str] = field(default_factory=list)  # ["text", "embeddings"]


@dataclass
class TrainingInfo:
    """Information about how the model was trained."""
    dataset_hints: List[str] = field(default_factory=list)  # Dataset names/descriptions
    training_steps: Optional[int] = None
    learning_rate: Optional[float] = None
    batch_size: Optional[int] = None
    precision: Optional[str] = None  # "fp16", "bf16", "fp32"
    hardware: Optional[str] = None  # "A100", "TPU-v4", etc.
    base_model: Optional[str] = None  # If fine-tuned from another model


@dataclass
class OsmoManifest:
    """Top-level manifest for an .osmo model."""
    version: str = "1.0.0"
    name: str = ""
    description: str = ""
    modality: ModelModality = ModelModality.UNKNOWN
    framework_origin: FrameworkOrigin = FrameworkOrigin.UNKNOWN
    original_repo: Optional[str] = None  # HuggingFace repo, etc.
    capabilities: ModelCapabilities = field(default_factory=ModelCapabilities)
    training_info: TrainingInfo = field(default_factory=TrainingInfo)
    parameter_count: int = 0
    created_at: str = ""
    license: str = "unknown"
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict:
        return {
            "version": self.version,
            "name": self.name,
            "description": self.description,
            "modality": self.modality.value,
            "framework_origin": self.framework_origin.value,
            "original_repo": self.original_repo,
            "capabilities": asdict(self.capabilities),
            "training_info": asdict(self.training_info),
            "parameter_count": self.parameter_count,
            "created_at": self.created_at,
            "license": self.license,
            "tags": self.tags,
        }

    @classmethod
    def from_dict(cls, d: Dict) -> "OsmoManifest":
        d = d.copy()
        d["modality"] = ModelModality(d.get("modality", "unknown"))
        d["framework_origin"] = FrameworkOrigin(d.get("framework_origin", "unknown"))
        d["capabilities"] = ModelCapabilities(**d.get("capabilities", {}))
        d["training_info"] = TrainingInfo(**d.get("training_info", {}))
        return cls(**d)


@dataclass
class OsmoModel:
    """
    The complete OsmoFormat model representation.

    This is the universal intermediate format that connects all frameworks.
    Any model can be converted TO this format, and FROM this format to any target.
    """
    manifest: OsmoManifest
    architecture: ArchitectureSpec
    parameters: Dict[str, np.ndarray]  # Layer name -> numpy array
    knowledge: Dict[str, np.ndarray] = field(default_factory=dict)  # Extracted knowledge
    auxiliary_files: Dict[str, Any] = field(default_factory=dict)  # Tokenizer, configs, etc.
    fingerprint: str = ""  # SHA256 of parameters for deduplication

    def __post_init__(self):
        if not self.fingerprint:
            self.fingerprint = self._compute_fingerprint()

    def _compute_fingerprint(self) -> str:
        """Compute a unique fingerprint for this model."""
        hasher = hashlib.sha256()
        for name in sorted(self.parameters.keys()):
            hasher.update(name.encode())
            hasher.update(self.parameters[name].tobytes())
        return hasher.hexdigest()[:16]

    @property
    def num_parameters(self) -> int:
        """Total number of parameters."""
        return sum(p.size for p in self.parameters.values())

    def get_layer_parameters(self, layer_type: LayerType) -> Dict[str, np.ndarray]:
        """Get all parameters for a specific layer type."""
        result = {}
        for layer in self.architecture.layers:
            if layer.layer_type == layer_type:
                if layer.name in self.parameters:
                    result[layer.name] = self.parameters[layer.name]
        return result

    def get_embedding_dim(self) -> int:
        """Get the embedding/hidden dimension."""
        return self.architecture.hidden_size

    def compatible_with(self, other: "OsmoModel") -> bool:
        """Check if this model is compatible for merging with another."""
        # Same modality or both multimodal
        modality_ok = (
            self.manifest.modality == other.manifest.modality
            or self.manifest.modality == ModelModality.MULTIMODAL
            or other.manifest.modality == ModelModality.MULTIMODAL
        )
        # Same hidden size
        hidden_ok = self.architecture.hidden_size == other.architecture.hidden_size
        return modality_ok and hidden_ok

    def summary(self) -> str:
        """Human-readable summary."""
        return (
            f"OsmoModel: {self.manifest.name}\n"
            f"  Modality: {self.manifest.modality.value}\n"
            f"  Origin: {self.manifest.framework_origin.value}\n"
            f"  Parameters: {self.num_parameters:,}\n"
            f"  Hidden Size: {self.architecture.hidden_size}\n"
            f"  Layers: {self.architecture.num_layers}\n"
            f"  Fingerprint: {self.fingerprint}\n"
        )


# Constants
OSMO_FORMAT_VERSION = "1.0.0"
OSMO_EXTENSION = ".osmo"
MANIFEST_FILE = "manifest.json"
ARCHITECTURE_FILE = "architecture.json"
PARAMETERS_DIR = "parameters"
KNOWLEDGE_DIR = "knowledge"
AUXILIARY_DIR = "auxiliary"
FINGERPRINT_FILE = "fingerprint.json"
