"""
OsmoFormat Serialization - Save and load .osmo models to/from disk.

Supports:
- Directory format (model.osmo/) for easy inspection
- Archive format (model.osmo.zip) for distribution
- Streaming loading for large models
"""

import json
import shutil
import tempfile
import zipfile
from pathlib import Path
from typing import Any, Dict, Iterator, Optional, Tuple, Union

import numpy as np

try:
    from safetensors.numpy import save_file as save_safetensors_numpy
    from safetensors.numpy import load_file as load_safetensors_numpy
    HAS_SAFETENSORS = True
except ImportError:
    HAS_SAFETENSORS = False

from model_osmosis.format.osmo_format import (
    OsmoModel,
    OsmoManifest,
    ArchitectureSpec,
    OSMO_FORMAT_VERSION,
    MANIFEST_FILE,
    ARCHITECTURE_FILE,
    PARAMETERS_DIR,
    KNOWLEDGE_DIR,
    AUXILIARY_DIR,
    FINGERPRINT_FILE,
)


class OsmoSerializer:
    """Serialize OsmoModel to disk."""

    def __init__(self, use_compression: bool = True):
        """
        Args:
            use_compression: If True, create .osmo.zip archives
        """
        self.use_compression = use_compression

    def save(
        self,
        model: OsmoModel,
        path: Union[str, Path],
        shard_size_mb: int = 5000,
    ) -> Path:
        """
        Save an OsmoModel to disk.

        Args:
            model: The OsmoModel to save
            path: Destination path (directory or .zip)
            shard_size_mb: Maximum shard size in MB for large models

        Returns:
            Path to saved model
        """
        path = Path(path)

        # Ensure .osmo extension
        if not path.suffix:
            path = path.with_suffix(".osmo")

        # Create directory structure
        if self.use_compression:
            # Save to temp dir, then zip
            with tempfile.TemporaryDirectory() as tmpdir:
                tmppath = Path(tmpdir) / "model"
                self._save_directory(model, tmppath, shard_size_mb)
                self._create_archive(tmppath, path)
        else:
            self._save_directory(model, path, shard_size_mb)

        return path

    def _save_directory(
        self,
        model: OsmoModel,
        path: Path,
        shard_size_mb: int,
    ) -> None:
        """Save model to directory structure."""
        path.mkdir(parents=True, exist_ok=True)

        # Save manifest
        manifest_path = path / MANIFEST_FILE
        with open(manifest_path, "w") as f:
            json.dump(model.manifest.to_dict(), f, indent=2)

        # Save architecture
        arch_path = path / ARCHITECTURE_FILE
        with open(arch_path, "w") as f:
            json.dump(model.architecture.to_dict(), f, indent=2)

        # Save parameters
        params_dir = path / PARAMETERS_DIR
        params_dir.mkdir(exist_ok=True)
        self._save_parameters(model.parameters, params_dir, shard_size_mb)

        # Save knowledge if present
        if model.knowledge:
            knowledge_dir = path / KNOWLEDGE_DIR
            knowledge_dir.mkdir(exist_ok=True)
            for name, arr in model.knowledge.items():
                np.save(knowledge_dir / f"{name}.npy", arr)

        # Save auxiliary files (tokenizer, config, preprocessor, etc.)
        if model.auxiliary_files:
            aux_dir = path / AUXILIARY_DIR
            aux_dir.mkdir(exist_ok=True)
            for name, data in model.auxiliary_files.items():
                aux_file = aux_dir / name
                aux_file.parent.mkdir(parents=True, exist_ok=True)
                if isinstance(data, (dict, list)):
                    with open(aux_file, "w") as f:
                        json.dump(data, f, indent=2, default=str)
                elif isinstance(data, bytes):
                    with open(aux_file, "wb") as f:
                        f.write(data)
                elif isinstance(data, str):
                    with open(aux_file, "w", encoding="utf-8") as f:
                        f.write(data)

        # Save fingerprint
        fingerprint_path = path / FINGERPRINT_FILE
        with open(fingerprint_path, "w") as f:
            json.dump({"fingerprint": model.fingerprint}, f)

    def _save_parameters(
        self,
        parameters: Dict[str, np.ndarray],
        params_dir: Path,
        shard_size_mb: int,
    ) -> None:
        """Save parameters, sharding if necessary."""
        if not HAS_SAFETENSORS:
            # Fallback to numpy
            for name, arr in parameters.items():
                safe_name = name.replace("/", "_").replace(".", "_")
                np.save(params_dir / f"{safe_name}.npy", arr)
            # Save name mapping
            with open(params_dir / "names.json", "w") as f:
                name_map = {
                    name.replace("/", "_").replace(".", "_"): name
                    for name in parameters.keys()
                }
                json.dump(name_map, f)
            return

        # Calculate total size
        total_bytes = sum(arr.nbytes for arr in parameters.values())
        shard_bytes = shard_size_mb * 1024 * 1024

        if total_bytes <= shard_bytes:
            # Single shard
            save_safetensors_numpy(parameters, str(params_dir / "model.safetensors"))
        else:
            # Multiple shards
            current_shard = {}
            current_size = 0
            shard_idx = 0

            for name, arr in parameters.items():
                if current_size + arr.nbytes > shard_bytes and current_shard:
                    # Save current shard
                    save_safetensors_numpy(
                        current_shard,
                        str(params_dir / f"model-{shard_idx:05d}.safetensors"),
                    )
                    shard_idx += 1
                    current_shard = {}
                    current_size = 0

                current_shard[name] = arr
                current_size += arr.nbytes

            # Save final shard
            if current_shard:
                save_safetensors_numpy(
                    current_shard,
                    str(params_dir / f"model-{shard_idx:05d}.safetensors"),
                )

            # Save index
            with open(params_dir / "model.safetensors.index.json", "w") as f:
                json.dump({"total_shards": shard_idx + 1}, f)

    def _create_archive(self, source_dir: Path, dest_path: Path) -> None:
        """Create a zip archive from directory."""
        archive_path = dest_path.with_suffix(".osmo.zip")
        with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as zf:
            for file_path in source_dir.rglob("*"):
                if file_path.is_file():
                    arcname = file_path.relative_to(source_dir)
                    zf.write(file_path, arcname)

        # Rename to final path if needed
        if archive_path != dest_path:
            shutil.move(str(archive_path), str(dest_path))


class OsmoLoader:
    """Load OsmoModel from disk."""

    _TEXT_AUX_EXTENSIONS = {
        ".txt",
        ".md",
        ".yaml",
        ".yml",
        ".ini",
        ".cfg",
        ".conf",
        ".py",
        ".csv",
        ".tsv",
        ".merges",
        ".modelcard",
        ".special_tokens_map",
        ".template",
    }

    @classmethod
    def _load_auxiliary_file(cls, aux_file: Path) -> Any:
        """Load auxiliary files with binary-safe behavior."""
        suffix = aux_file.suffix.lower()
        if suffix == ".json":
            with open(aux_file, encoding="utf-8") as f:
                return json.load(f)
        if suffix in cls._TEXT_AUX_EXTENSIONS:
            with open(aux_file, "r", encoding="utf-8") as f:
                return f.read()
        with open(aux_file, "rb") as f:
            return f.read()

    def load(
        self,
        path: Union[str, Path],
        load_parameters: bool = True,
        load_knowledge: bool = True,
    ) -> OsmoModel:
        """
        Load an OsmoModel from disk.

        Args:
            path: Path to .osmo directory or .zip archive
            load_parameters: Whether to load parameters (set False for metadata only)
            load_knowledge: Whether to load extracted knowledge

        Returns:
            Loaded OsmoModel
        """
        path = Path(path)

        # Handle archive
        if path.suffix == ".zip" or (path.is_file() and zipfile.is_zipfile(path)):
            return self._load_from_archive(path, load_parameters, load_knowledge)

        return self._load_from_directory(path, load_parameters, load_knowledge)

    def _load_from_directory(
        self,
        path: Path,
        load_parameters: bool,
        load_knowledge: bool,
    ) -> OsmoModel:
        """Load from directory structure."""
        # Load manifest
        with open(path / MANIFEST_FILE) as f:
            manifest = OsmoManifest.from_dict(json.load(f))

        # Load architecture
        with open(path / ARCHITECTURE_FILE) as f:
            architecture = ArchitectureSpec.from_dict(json.load(f))

        # Load parameters
        parameters = {}
        if load_parameters:
            parameters = self._load_parameters(path / PARAMETERS_DIR)

        # Load knowledge
        knowledge = {}
        if load_knowledge:
            knowledge_dir = path / KNOWLEDGE_DIR
            if knowledge_dir.exists():
                for npy_file in knowledge_dir.glob("*.npy"):
                    name = npy_file.stem
                    knowledge[name] = np.load(npy_file)

        # Load auxiliary files (tokenizer, config, preprocessor, etc.)
        auxiliary_files = {}
        aux_dir = path / AUXILIARY_DIR
        if aux_dir.exists():
            for aux_file in aux_dir.rglob("*"):
                if aux_file.is_file():
                    rel_name = aux_file.relative_to(aux_dir).as_posix()
                    auxiliary_files[rel_name] = self._load_auxiliary_file(aux_file)

        # Load fingerprint
        fingerprint = ""
        fingerprint_path = path / FINGERPRINT_FILE
        if fingerprint_path.exists():
            with open(fingerprint_path) as f:
                fingerprint = json.load(f).get("fingerprint", "")

        return OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=parameters,
            knowledge=knowledge,
            auxiliary_files=auxiliary_files,
            fingerprint=fingerprint,
        )

    def _load_parameters(self, params_dir: Path) -> Dict[str, np.ndarray]:
        """Load parameters from directory."""
        if not params_dir.exists():
            return {}

        parameters = {}

        # Check for safetensors
        safetensor_files = list(params_dir.glob("*.safetensors"))
        if safetensor_files and HAS_SAFETENSORS:
            for sf_file in safetensor_files:
                params = load_safetensors_numpy(str(sf_file))
                parameters.update(params)
            return parameters

        # Fallback to numpy
        name_map = {}
        names_file = params_dir / "names.json"
        if names_file.exists():
            with open(names_file) as f:
                name_map = json.load(f)

        for npy_file in params_dir.glob("*.npy"):
            safe_name = npy_file.stem
            original_name = name_map.get(safe_name, safe_name)
            parameters[original_name] = np.load(npy_file)

        return parameters

    def _load_from_archive(
        self,
        archive_path: Path,
        load_parameters: bool,
        load_knowledge: bool,
    ) -> OsmoModel:
        """Load from zip archive."""
        with tempfile.TemporaryDirectory() as tmpdir:
            tmppath = Path(tmpdir)

            with zipfile.ZipFile(archive_path, "r") as zf:
                zf.extractall(tmppath)

            return self._load_from_directory(tmppath, load_parameters, load_knowledge)

    def stream_parameters(
        self,
        path: Union[str, Path],
    ) -> Iterator[Tuple[str, np.ndarray]]:
        """
        Stream parameters one at a time for memory-efficient loading.

        Yields:
            Tuples of (parameter_name, parameter_array)
        """
        path = Path(path)
        params_dir = path / PARAMETERS_DIR

        if not params_dir.exists():
            return

        # Stream from safetensors
        safetensor_files = sorted(params_dir.glob("*.safetensors"))
        if safetensor_files and HAS_SAFETENSORS:
            for sf_file in safetensor_files:
                params = load_safetensors_numpy(str(sf_file))
                for name, arr in params.items():
                    yield name, arr
            return

        # Stream from numpy
        name_map = {}
        names_file = params_dir / "names.json"
        if names_file.exists():
            with open(names_file) as f:
                name_map = json.load(f)

        for npy_file in sorted(params_dir.glob("*.npy")):
            safe_name = npy_file.stem
            original_name = name_map.get(safe_name, safe_name)
            yield original_name, np.load(npy_file)


def save_osmo(
    model: OsmoModel,
    path: Union[str, Path],
    compress: bool = True,
    shard_size_mb: int = 5000,
    obfuscate: bool = False,
    obfuscation_config: Optional[object] = None,
) -> Path:
    """
    Convenience function to save an OsmoModel.

    Args:
        model: The model to save
        path: Destination path
        compress: Whether to create a zip archive
        shard_size_mb: Maximum shard size for large models
        obfuscate: Whether to obfuscate weights before saving
        obfuscation_config: ObfuscationConfig instance (uses defaults if None)

    Returns:
        Path to saved model
    """
    if obfuscate:
        from model_osmosis.obfuscation import WeightObfuscator

        obfuscator = WeightObfuscator(obfuscation_config)
        model = obfuscator.obfuscate(model)

    serializer = OsmoSerializer(use_compression=compress)
    return serializer.save(model, path, shard_size_mb)


def load_osmo(
    path: Union[str, Path],
    load_parameters: bool = True,
    load_knowledge: bool = True,
    target_framework: Optional[str] = None,
) -> Union[OsmoModel, Any]:
    """
    Convenience function to load an OsmoModel.

    Args:
        path: Path to .osmo model
        load_parameters: Whether to load parameters
        load_knowledge: Whether to load knowledge
        target_framework: If set, convert to this framework on load.
            One of 'pytorch', 'tensorflow', 'jax', 'onnx', 'huggingface'.
            If None, returns an OsmoModel.

    Returns:
        OsmoModel (if no target_framework) or model in the target framework.
    """
    loader = OsmoLoader()
    osmo = loader.load(path, load_parameters, load_knowledge)

    if target_framework:
        from model_osmosis.format.converters import from_osmo
        return from_osmo(osmo, target_framework=target_framework)

    return osmo
