"""
TensorFlowAdapter — Load and convert TensorFlow/Keras models for ModelOsmosis.

Lazy-imports TensorFlow so it's not a required dependency.

Usage::

    from model_osmosis.adapters import TensorFlowAdapter

    adapter = TensorFlowAdapter()
    expert = adapter.load("path/to/saved_model")
    osmo = adapter.to_osmo(keras_model)
    keras_model = adapter.from_osmo(osmo_model)
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional, Tuple

import numpy as np


class TensorFlowExpert:
    """Wrapper around a TensorFlow/Keras model for use in ModelOsmosis pipelines."""

    def __init__(self, model: Any, name: Optional[str] = None):
        self.model = model
        self.name = name or getattr(model, "name", "tf_model")
        self._weights_cache: Optional[Dict[str, np.ndarray]] = None

    @property
    def weights(self) -> Dict[str, np.ndarray]:
        """Extract weights as a dict of name → numpy array."""
        if self._weights_cache is None:
            self._weights_cache = {}
            for layer in self.model.layers:
                for weight in layer.weights:
                    key = weight.name
                    self._weights_cache[key] = weight.numpy()
        return self._weights_cache

    @property
    def input_shape(self) -> Optional[Tuple]:
        """Get the model's input shape."""
        try:
            return tuple(self.model.input_shape)
        except (AttributeError, RuntimeError):
            return None

    @property
    def output_shape(self) -> Optional[Tuple]:
        """Get the model's output shape."""
        try:
            return tuple(self.model.output_shape)
        except (AttributeError, RuntimeError):
            return None

    def predict(self, data: np.ndarray) -> np.ndarray:
        """Run inference on numpy data."""
        return self.model.predict(data, verbose=0)


class TensorFlowAdapter:
    """Adapter for TensorFlow/Keras models.

    Supports:
    - Loading Keras .h5 and SavedModel formats
    - Extracting weights as numpy arrays
    - Converting to/from OsmoFormat
    - Loading weights back into a Keras model

    Requires: tensorflow >= 2.0
    """

    def __init__(self):
        self._tf = None
        self._keras = None

    def _lazy_import(self):
        """Lazy-import TensorFlow to avoid hard dependency."""
        if self._tf is None:
            try:
                import tensorflow as tf

                self._tf = tf
                self._keras = tf.keras
            except ImportError:
                raise ImportError(
                    "TensorFlow is required for TensorFlowAdapter. "
                    "Install it with: pip install tensorflow"
                )

    def load(self, path: str) -> TensorFlowExpert:
        """Load a TensorFlow/Keras model from disk.

        Args:
            path: Path to a SavedModel directory or .h5 file.

        Returns:
            TensorFlowExpert wrapping the loaded model.
        """
        self._lazy_import()
        model = self._keras.models.load_model(path)
        return TensorFlowExpert(model, name=path.split("/")[-1])

    def extract_weights(self, model: Any) -> Dict[str, np.ndarray]:
        """Extract all weights from a TF/Keras model as numpy arrays.

        Args:
            model: A tf.keras.Model or TensorFlowExpert.

        Returns:
            Dict of weight name → numpy array.
        """
        if isinstance(model, TensorFlowExpert):
            return model.weights

        self._lazy_import()
        weights = {}
        for layer in model.layers:
            for weight in layer.weights:
                weights[weight.name] = weight.numpy()
        return weights

    def load_weights(self, model: Any, weights: Dict[str, np.ndarray]) -> Any:
        """Load numpy weights back into a TF/Keras model.

        Args:
            model: The target tf.keras.Model.
            weights: Dict of weight name → numpy array.

        Returns:
            The model with updated weights.
        """
        self._lazy_import()
        for layer in model.layers:
            layer_weights = []
            for weight in layer.weights:
                if weight.name in weights:
                    layer_weights.append(weights[weight.name])
                else:
                    layer_weights.append(weight.numpy())
            layer.set_weights(layer_weights)
        return model

    def to_osmo(self, model: Any) -> Any:
        """Convert a TF/Keras model to OsmoFormat.

        Args:
            model: A tf.keras.Model or TensorFlowExpert.

        Returns:
            OsmoModel with the model's weights and architecture.
        """
        from model_osmosis.format.osmo_format import (
            ArchitectureSpec,
            LayerSpec,
            LayerType,
            OsmoManifest,
            OsmoModel,
        )

        weights = self.extract_weights(model)

        # Build architecture spec from model layers
        keras_model = model.model if isinstance(model, TensorFlowExpert) else model
        layers = []
        for layer in keras_model.layers:
            layer_type = LayerType.LINEAR
            layer_class = layer.__class__.__name__.lower()
            if "conv" in layer_class:
                layer_type = LayerType.CONV
            elif "norm" in layer_class:
                layer_type = LayerType.NORM
            elif "attention" in layer_class:
                layer_type = LayerType.ATTENTION
            elif "embedding" in layer_class:
                layer_type = LayerType.EMBEDDING

            output_shape = getattr(layer, "output_shape", None)
            shape = list(output_shape) if output_shape else []

            layers.append(
                LayerSpec(
                    name=layer.name,
                    layer_type=layer_type,
                    shape=shape,
                    dtype="float32",
                )
            )

        # Determine hidden size from first dense layer
        hidden_size = 0
        for layer in keras_model.layers:
            if hasattr(layer, "units"):
                hidden_size = layer.units
                break

        architecture = ArchitectureSpec(
            layers=layers,
            hidden_size=hidden_size,
            num_layers=len(layers),
            extra={"framework": "tensorflow"},
        )

        manifest = OsmoManifest(
            model_name=getattr(keras_model, "name", "tf_model"),
            description=f"TensorFlow/Keras model converted via TensorFlowAdapter",
            tags=["tensorflow", "keras"],
        )

        return OsmoModel(
            manifest=manifest,
            architecture=architecture,
            parameters=weights,
        )

    def from_osmo(self, osmo_model: Any, target_model: Any = None) -> Any:
        """Convert an OsmoModel back to TF/Keras format.

        Args:
            osmo_model: The OsmoModel to convert.
            target_model: An existing Keras model to load weights into.
                If None, returns the raw weight dictionary.

        Returns:
            The Keras model with loaded weights, or weight dict if no target.
        """
        if target_model is None:
            return osmo_model.parameters

        return self.load_weights(target_model, osmo_model.parameters)
