"""
WeightObfuscator — Main orchestrator for full-model obfuscation.

Ties together layer graph analysis, transforms, and verification into
a single API that makes ALL model files untraceable while preserving behavior.

Obfuscates:
- Weight parameters (permutation, rotation, scaling)
- Parameter names (hashed to opaque keys)
- Knowledge files (noise-injected structural/functional data)
- Architecture spec (randomized layer names)
- Manifest metadata (stripped provenance)
- Fingerprint (recomputed from new weights)
"""

from __future__ import annotations

import copy
import hashlib
import logging
from dataclasses import dataclass
from typing import Any, Dict, Optional, Tuple

import numpy as np

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

from model_osmosis.format.osmo_format import ArchitectureSpec, OsmoModel
from model_osmosis.obfuscation.graph import LayerGraph
from model_osmosis.obfuscation.transforms import (
    obfuscate_architecture,
    obfuscate_auxiliary_files,
    obfuscate_knowledge,
    obfuscate_parameter_names,
    permute_neurons,
    randomize_metadata,
    rotate_basis,
    scale_weights,
)
from model_osmosis.obfuscation.verification import ObfuscationVerifier


@dataclass
class ObfuscationConfig:
    """Configuration for full-model obfuscation.

    Quality control:
        preserve_quality: When True, uses only exact transforms (permutation
            + scaling) and skips rotation which introduces float error.
            Knowledge obfuscation uses permutation-only (no noise).
            Default True. Set False for maximum untraceability at the
            cost of ~1e-7 per-round floating point error.

    Weight transforms (function-preserving):
        enable_permutation: Apply random neuron permutation (exact, no error).
        enable_rotation: Apply orthogonal basis rotation (~1e-7 error).
            Automatically disabled when preserve_quality=True.
        enable_scaling: Apply per-neuron random scaling (exact).
        scale_range: Range for random scale factors (min, max).
        num_rounds: Number of obfuscation rounds for extra mixing.

    Full-file obfuscation (makes ALL output files untraceable):
        obfuscate_all_files: Master switch for full-file obfuscation.
        randomize_metadata: Strip provenance info from manifest.
        obfuscate_knowledge_files: Noise-inject knowledge data.
        obfuscate_layer_names: Randomize layer names in architecture spec.
        obfuscate_param_names: Hash parameter keys to opaque names.
        obfuscate_auxiliary: Strip provenance from tokenizer/config files.
        knowledge_noise_scale: Noise magnitude for knowledge obfuscation.
            Only used when preserve_quality=False.

    Verification:
        seed: Random seed for reproducibility (None = random).
        verify: Run functional equivalence verification after obfuscation.
        verification_samples: Number of random samples for verification.
        verification_tolerance: Maximum allowed output difference.
    """

    # Quality control
    preserve_quality: bool = True

    # Weight transforms
    enable_permutation: bool = True
    enable_rotation: bool = True
    enable_scaling: bool = True
    scale_range: Tuple[float, float] = (0.5, 2.0)
    num_rounds: int = 1

    # Full-file obfuscation
    obfuscate_all_files: bool = True
    randomize_metadata: bool = True
    obfuscate_knowledge_files: bool = True
    obfuscate_layer_names: bool = True
    obfuscate_param_names: bool = True
    obfuscate_auxiliary: bool = True
    knowledge_noise_scale: float = 0.01

    # Verification
    seed: Optional[int] = None
    verify: bool = True
    verification_samples: int = 100
    verification_tolerance: float = 1e-6

    @classmethod
    def auto(cls, model: "OsmoModel") -> "ObfuscationConfig":
        """Automatically determine the best safe obfuscation settings based on architecture."""
        arch_type = model.architecture.architecture_type.lower()
        
        # Architectures known to break permutation equivalence 
        # (e.g., due to tied vocab, RoPE, Grouped-Query Attention)
        complex_transformers = ["llama", "qwen", "qwen2", "mistral", "phi"]
        
        if any(llm in arch_type for llm in complex_transformers):
            logger.info(
                "Detected complex transformer (%s). Using safe (scaling only) obfuscation.",
                arch_type,
            )
            return cls(
                preserve_quality=True,
                enable_permutation=False,
                enable_rotation=False,
                enable_scaling=True,
                verify=True,
            )
        else:
            logger.info(
                "Detected standard model (%s). Maximizing obfuscation.", arch_type
            )
            return cls(
                preserve_quality=True,
                enable_permutation=True,
                enable_rotation=False,
                enable_scaling=True,
                verify=True,
            )


class WeightObfuscator:
    """Apply function-preserving transforms to make model weights untraceable.

    Uses three complementary techniques:

    1. **Neuron Permutation**: Randomly reorder neurons within hidden layers.
       Changes weight matrices completely while preserving computation.
       Mathematically exact — zero error.

    2. **Orthogonal Basis Rotation**: Rotate the representation space using
       random orthogonal matrices. Makes statistical fingerprinting impossible.
       Tiny float error (~1e-7 for float32).

    3. **Weight Scaling**: Random per-neuron scaling with inverse compensation.
       Changes weight magnitudes while preserving computation. Exact.

    Together these transforms ensure:
    - Model outputs are identical (within float precision)
    - Weight matrices have completely different values
    - No statistical method can match weights to source models
    - Cosine similarity between original and obfuscated weights ≈ 0

    Usage::

        obfuscator = WeightObfuscator(ObfuscationConfig(seed=42))
        obfuscated_model = obfuscator.obfuscate(osmo_model, sample_data=X)
    """

    def __init__(self, config: Optional[ObfuscationConfig] = None):
        self.config = config or ObfuscationConfig()
        self._rng = np.random.default_rng(self.config.seed)
        self._graph_cache: Dict[str, list] = {}

    def obfuscate(
        self,
        model: OsmoModel,
        sample_data: Optional[np.ndarray] = None,
    ) -> OsmoModel:
        """Obfuscate an OsmoModel's weights while preserving behavior.

        Args:
            model: The OsmoModel to obfuscate.
            sample_data: Optional input data for verification. If verify=True
                and no sample_data is provided, random data is generated.

        Returns:
            A new OsmoModel with obfuscated parameters and cleaned metadata.

        Raises:
            ObfuscationError: If obfuscation fails.
            VerificationError: If obfuscated model fails equivalence check.
        """
        from model_osmosis.core.exceptions import ObfuscationError, VerificationError

        # Deep copy parameters to avoid modifying the original
        obfuscated_params = {k: v.copy() for k, v in model.parameters.items()}

        # Keep original for verification
        if self.config.verify:
            original_params = {k: v.copy() for k, v in model.parameters.items()}

        # Determine effective rotation setting
        use_rotation = self.config.enable_rotation
        if self.config.preserve_quality:
            use_rotation = False

        # Build layer graph (with caching based on parameter structure)
        cache_key = self._param_structure_hash(obfuscated_params)
        if cache_key in self._graph_cache:
            groups = self._graph_cache[cache_key]
        else:
            graph = LayerGraph(obfuscated_params, model.architecture)
            groups = graph.build_transform_groups()
            self._graph_cache[cache_key] = groups

        if not groups:
            # No transform groups found — only do metadata randomization
            logger.warning(
                "No adjacent layer pairs found. Only metadata will be randomized."
            )
        else:
            mode_label = "quality-preserving" if self.config.preserve_quality else "full"
            logger.info(
                "Found %d transform group(s) (mode: %s)", len(groups), mode_label
            )

            # Apply transforms for each round
            for round_idx in range(self.config.num_rounds):
                if self.config.num_rounds > 1:
                    logger.info(
                        "Round %d/%d", round_idx + 1, self.config.num_rounds
                    )

                for group in groups:
                    if self.config.enable_permutation:
                        permute_neurons(obfuscated_params, group, self._rng)

                    if self.config.enable_scaling:
                        scale_weights(
                            obfuscated_params,
                            group,
                            self._rng,
                            self.config.scale_range,
                        )

                    if use_rotation:
                        rotate_basis(obfuscated_params, group, self._rng)

        # --- Full-file obfuscation ---

        # Randomize metadata (manifest.json)
        new_manifest = model.manifest
        if self.config.randomize_metadata:
            new_manifest = randomize_metadata(model.manifest)
            logger.info("Metadata stripped (manifest.json)")

        # Obfuscate architecture layer names (architecture.json)
        new_architecture = copy.deepcopy(model.architecture)
        if self.config.obfuscate_all_files and self.config.obfuscate_layer_names:
            new_architecture = obfuscate_architecture(
                model.architecture, self._rng
            )
            logger.info("Layer names randomized (architecture.json)")

        # Obfuscate knowledge files (knowledge/*.npy)
        new_knowledge: Dict[str, np.ndarray] = {}
        if model.knowledge:
            if self.config.obfuscate_all_files and self.config.obfuscate_knowledge_files:
                noise_scale = (
                    0.0 if self.config.preserve_quality
                    else self.config.knowledge_noise_scale
                )
                new_knowledge = obfuscate_knowledge(
                    model.knowledge, self._rng, noise_scale
                )
                logger.info(
                    "Knowledge files obfuscated (%d arrays, noise=%s)",
                    len(new_knowledge),
                    "off" if self.config.preserve_quality else f"{noise_scale:.3f}",
                )
            else:
                new_knowledge = {k: v.copy() for k, v in model.knowledge.items()}

        # Obfuscate auxiliary files (tokenizer, config, preprocessor, etc.)
        new_auxiliary: Dict[str, Any] = {}
        if model.auxiliary_files:
            if self.config.obfuscate_all_files and self.config.obfuscate_auxiliary:
                new_auxiliary = obfuscate_auxiliary_files(
                    model.auxiliary_files, self._rng
                )
                logger.info(
                    "Auxiliary files obfuscated (%d files)", len(new_auxiliary)
                )
            else:
                new_auxiliary = copy.deepcopy(model.auxiliary_files)

        # Verify functional equivalence BEFORE renaming params
        # (verifier needs matching keys between original and obfuscated)
        if self.config.verify and groups:
            if sample_data is None:
                sample_data = self._generate_sample_data(model)

            if sample_data is not None:
                verifier = ObfuscationVerifier()
                result = verifier.verify(
                    original_params=original_params,
                    obfuscated_params=obfuscated_params,
                    architecture=model.architecture,
                    sample_data=sample_data,
                    tolerance=self.config.verification_tolerance,
                )

                logger.info("%s", result.details)

                if not result.passed:
                    raise VerificationError(
                        f"Obfuscated model failed equivalence check: "
                        f"max_error={result.max_absolute_error:.2e} > "
                        f"tolerance={self.config.verification_tolerance:.2e}"
                    )

        # Report weight dissimilarity (before name hashing)
        if groups:
            self._report_dissimilarity(model.parameters, obfuscated_params)

        # Obfuscate parameter names LAST (changes dict keys)
        if self.config.obfuscate_all_files and self.config.obfuscate_param_names:
            obfuscated_params = obfuscate_parameter_names(
                obfuscated_params, self._rng
            )
            logger.info("Parameter names hashed")

        # Build new model
        obfuscated_model = OsmoModel(
            manifest=new_manifest,
            architecture=new_architecture,
            parameters=obfuscated_params,
            knowledge=new_knowledge,
            auxiliary_files=new_auxiliary,
            fingerprint="",  # Will be recomputed by __post_init__
        )

        logger.info("Obfuscation complete")
        return obfuscated_model

    def obfuscate_state_dict(
        self,
        params: Dict[str, np.ndarray],
        architecture: Optional[ArchitectureSpec] = None,
    ) -> Dict[str, np.ndarray]:
        """Obfuscate a raw parameter dictionary.

        Lower-level API for when you have parameters but not a full OsmoModel.

        Args:
            params: Dictionary of parameter name → numpy array.
            architecture: Optional architecture spec for better graph analysis.

        Returns:
            New dictionary with obfuscated parameters.
        """
        obfuscated = {k: v.copy() for k, v in params.items()}

        graph = LayerGraph(obfuscated, architecture)
        groups = graph.build_transform_groups()

        # Determine effective rotation setting
        use_rotation = self.config.enable_rotation
        if self.config.preserve_quality:
            use_rotation = False

        for _ in range(self.config.num_rounds):
            for group in groups:
                if self.config.enable_permutation:
                    permute_neurons(obfuscated, group, self._rng)
                if self.config.enable_scaling:
                    scale_weights(
                        obfuscated, group, self._rng, self.config.scale_range
                    )
                if use_rotation:
                    rotate_basis(obfuscated, group, self._rng)

        return obfuscated

    def obfuscate_torch_model(
        self,
        model: Any,
        sample_data: Optional[np.ndarray] = None,
    ) -> Any:
        """Obfuscate a PyTorch model's weights in-place.

        Args:
            model: A torch.nn.Module.
            sample_data: Optional input data for verification.

        Returns:
            The same model with obfuscated weights.
        """
        import torch

        # Extract state dict as numpy, remembering original device/dtype
        state_dict = model.state_dict()
        device_map = {k: v.device for k, v in state_dict.items()}
        params = {k: v.cpu().numpy() for k, v in state_dict.items()}

        # Obfuscate
        obfuscated = self.obfuscate_state_dict(params)

        # Load back, restoring original device and dtype
        new_state = {
            k: torch.from_numpy(v).to(dtype=state_dict[k].dtype, device=device_map[k])
            for k, v in obfuscated.items()
        }
        model.load_state_dict(new_state)

        return model

    @staticmethod
    def _param_structure_hash(params: Dict[str, np.ndarray]) -> str:
        """Compute a stable hash from parameter names, shapes, and dtypes.

        Used for graph caching — two parameter dicts with the same structure
        will produce the same hash regardless of object identity.
        """
        h = hashlib.sha256()
        for key in sorted(params.keys()):
            arr = params[key]
            h.update(f"{key}:{arr.shape}:{arr.dtype}\n".encode())
        return h.hexdigest()

    def _generate_sample_data(self, model: OsmoModel) -> Optional[np.ndarray]:
        """Generate random sample data for verification.

        Infers input dimension from the first weight matrix.
        """
        # Find the first weight matrix to determine input dimension
        for key in sorted(model.parameters.keys()):
            arr = model.parameters[key]
            if arr.ndim >= 2:
                # PyTorch convention: [out_features, in_features]
                input_dim = arr.shape[1]
                return self._rng.standard_normal(
                    (self.config.verification_samples, input_dim)
                ).astype(np.float32)

        return None

    def _report_dissimilarity(
        self,
        original: Dict[str, np.ndarray],
        obfuscated: Dict[str, np.ndarray],
    ) -> None:
        """Report how different the obfuscated weights are from originals."""
        similarities = []
        for key in original:
            if key not in obfuscated:
                continue
            a = original[key].flatten().astype(np.float64)
            b = obfuscated[key].flatten().astype(np.float64)

            norm_a = np.linalg.norm(a)
            norm_b = np.linalg.norm(b)

            if norm_a > 0 and norm_b > 0:
                cos_sim = np.dot(a, b) / (norm_a * norm_b)
                similarities.append(cos_sim)

        if similarities:
            avg_sim = np.mean(similarities)
            logger.info(
                "Average cosine similarity to original: %.4f "
                "(0.0 = completely different, 1.0 = identical)",
                avg_sim,
            )


def obfuscate_osmo(
    model: OsmoModel,
    sample_data: Optional[np.ndarray] = None,
    config: Optional[ObfuscationConfig] = None,
) -> OsmoModel:
    """Convenience function to obfuscate an OsmoModel.

    Args:
        model: The OsmoModel to obfuscate.
        sample_data: Optional input data for verification.
        config: Obfuscation settings (uses defaults if None).

    Returns:
        A new OsmoModel with obfuscated weights and cleaned metadata.
    """
    obfuscator = WeightObfuscator(config)
    return obfuscator.obfuscate(model, sample_data)
