"""
Streaming parameter loader — RAM-efficient layer-by-layer parameter loading.

Instead of loading an entire model into memory, this streams parameters
one layer at a time, yielding (name, tensor) pairs.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Dict, Generator, List, Optional, Tuple, Union

import numpy as np


class StreamingParameterLoader:
    """Load model parameters incrementally to minimize RAM usage.

    Supports:
    - Live PyTorch nn.Module objects (extracts params without copying all at once)
    - PyTorch checkpoint files (.pt, .pth) via lazy loading
    - State dict files
    """

    def stream_parameters(
        self,
        source: Any,
    ) -> Generator[Tuple[str, np.ndarray], None, None]:
        """Yield (parameter_name, weight_array) pairs one at a time.

        The caller can process and discard each parameter before the next
        is loaded, keeping total RAM usage bounded.

        Args:
            source: One of:
                - A live PyTorch nn.Module
                - A file path (str or Path) to a saved model/state_dict
                - A dict (state_dict already loaded)

        Yields:
            (name, array) tuples where name is the parameter key and
            array is a numpy ndarray of the weights.
        """
        if isinstance(source, dict):
            yield from self._stream_from_dict(source)
        elif isinstance(source, (str, Path)):
            yield from self._stream_from_file(source)
        else:
            yield from self._stream_from_module(source)

    def _stream_from_module(
        self, model: Any
    ) -> Generator[Tuple[str, np.ndarray], None, None]:
        """Stream parameters from a live PyTorch module."""
        import torch.nn as nn

        if not isinstance(model, nn.Module):
            raise TypeError(
                f"Expected nn.Module, got {type(model).__name__}"
            )

        for name, param in model.named_parameters():
            # Detach and convert to numpy one at a time
            arr = param.detach().cpu().numpy()
            yield (name, arr)

    def _stream_from_dict(
        self, state_dict: dict
    ) -> Generator[Tuple[str, np.ndarray], None, None]:
        """Stream parameters from a state dict."""
        for name, tensor in state_dict.items():
            if hasattr(tensor, "numpy"):
                arr = tensor.cpu().numpy() if hasattr(tensor, "cpu") else tensor.numpy()
            else:
                arr = np.asarray(tensor)
            yield (name, arr)

    def _stream_from_file(
        self, path: Union[str, Path]
    ) -> Generator[Tuple[str, np.ndarray], None, None]:
        """Stream parameters from a saved PyTorch file.

        Uses mmap_mode when possible to avoid loading the entire file.
        """
        import torch

        path = Path(path)

        if path.suffix in (".pt", ".pth", ".bin"):
            # Load with weights_only for security, mmap if available
            try:
                state_dict = torch.load(
                    path, map_location="cpu", weights_only=True
                )
            except TypeError:
                # Older PyTorch without weights_only
                state_dict = torch.load(path, map_location="cpu")

            # If it's a full checkpoint (has 'state_dict' key), unwrap
            if isinstance(state_dict, dict) and "state_dict" in state_dict:
                state_dict = state_dict["state_dict"]
            elif isinstance(state_dict, dict) and "model_state_dict" in state_dict:
                state_dict = state_dict["model_state_dict"]

            yield from self._stream_from_dict(state_dict)

        elif path.suffix == ".safetensors":
            # HuggingFace safetensors format — inherently supports lazy loading
            try:
                from safetensors import safe_open

                with safe_open(str(path), framework="numpy") as f:
                    for name in f.keys():
                        yield (name, f.get_tensor(name))
            except ImportError:
                raise ImportError(
                    "Install safetensors to load .safetensors files: "
                    "pip install safetensors"
                )
        else:
            raise ValueError(f"Unsupported file format: {path.suffix}")

    def get_all_param_names(self, source: Any) -> List[str]:
        """Get all parameter names without loading weights.

        Useful for planning merge operations.
        """
        if isinstance(source, dict):
            return list(source.keys())

        try:
            import torch.nn as nn

            if isinstance(source, nn.Module):
                return [n for n, _ in source.named_parameters()]
        except ImportError:
            pass

        if isinstance(source, (str, Path)):
            # Must load to get names (but don't keep weights)
            names = []
            for name, _ in self.stream_parameters(source):
                names.append(name)
            return names

        return []

    def save_to_safetensors(
        self,
        state_dict: Dict[str, Any],
        path: Union[str, Path],
        metadata: Optional[Dict[str, str]] = None,
    ) -> Path:
        """Save a state dict to safetensors format.

        Args:
            state_dict: Dict mapping parameter names to tensors (torch or numpy).
            path: Output file path.
            metadata: Optional metadata to include in the file header.

        Returns:
            The path to the saved file.

        Example::

            loader = StreamingParameterLoader()
            loader.save_to_safetensors(
                model.state_dict(),
                "model.safetensors",
                metadata={"format": "pt", "library": "model-osmosis"}
            )
        """
        import torch

        try:
            from safetensors.torch import save_file
        except ImportError:
            raise ImportError(
                "Install safetensors to save .safetensors files: "
                "pip install safetensors"
            )

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

        # Convert numpy arrays to torch tensors and ensure contiguous
        torch_dict = {}
        for name, tensor in state_dict.items():
            if isinstance(tensor, np.ndarray):
                torch_dict[name] = torch.from_numpy(tensor).contiguous()
            elif hasattr(tensor, "contiguous"):
                torch_dict[name] = tensor.contiguous()
            else:
                torch_dict[name] = torch.tensor(tensor).contiguous()

        save_file(torch_dict, str(path), metadata=metadata or {})
        return path

    def save_to_pytorch(
        self,
        state_dict: Dict[str, Any],
        path: Union[str, Path],
        include_metadata: bool = False,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> Path:
        """Save a state dict to PyTorch format (.pt/.pth).

        Args:
            state_dict: Dict mapping parameter names to tensors.
            path: Output file path.
            include_metadata: If True, wrap state_dict with metadata.
            metadata: Additional metadata to include (requires include_metadata=True).

        Returns:
            The path to the saved file.
        """
        import torch

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

        # Convert numpy arrays to torch tensors
        torch_dict = {}
        for name, tensor in state_dict.items():
            if isinstance(tensor, np.ndarray):
                torch_dict[name] = torch.from_numpy(tensor)
            elif not isinstance(tensor, torch.Tensor):
                torch_dict[name] = torch.tensor(tensor)
            else:
                torch_dict[name] = tensor

        if include_metadata:
            save_obj = {
                "state_dict": torch_dict,
                "metadata": metadata or {},
            }
        else:
            save_obj = torch_dict

        torch.save(save_obj, path)
        return path
