from __future__ import annotations

import json
import os
import zipfile
from dataclasses import asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Union

from ..utils.exceptions import ToxoError


class PackageWriterMixin:
    def save(self, path: Union[str, Path], format: str = "toxo") -> None:
        """
        Save the Toxo layer to disk.

        Args:
            path: Save path
            format: Save format ('toxo', 'json', 'pickle')
        """
        path = Path(path)

        try:
            if format == "toxo":
                self._save_toxo_format(path)
            elif format == "json":
                self._save_json_format(path)
            elif format == "pickle":
                self._save_pickle_format(path)
            else:
                raise ValueError(f"Unsupported format: {format}")

            self.logger.info(f"Toxo layer saved to {path}")

        except Exception as e:
            self.logger.error(f"Failed to save layer: {str(e)}")
            raise ToxoError(f"Save operation failed: {str(e)}")

    def _save_json_format(self, path: Path) -> None:
        """Save minimal JSON snapshot (best-effort)."""
        data = {
            "name": getattr(self, "name", None),
            "domain": getattr(self, "domain", None),
            "task_type": getattr(self, "task_type", None),
            "description": getattr(self, "description", None),
            "is_trained": bool(getattr(self, "is_trained", False)),
            "metadata": asdict(getattr(self, "metadata", {})) if hasattr(self, "metadata") else {},
        }
        if path.suffix != ".json":
            path = path.with_suffix(".json")
        with open(path, "w") as f:
            json.dump(data, f, indent=2, default=str)

    def _save_pickle_format(self, path: Path) -> None:
        """
        Save a best-effort pickle snapshot.

        NOTE: Pickling the live `ToxoLayer` instance is brittle because it can
        contain non-pickleable objects (thread locks, open clients, loggers).
        We therefore pickle a lightweight, serializable snapshot instead.
        """
        import pickle

        if path.suffix != ".pkl":
            path = path.with_suffix(".pkl")
        with open(path, "wb") as f:
            snapshot = {
                "name": getattr(self, "name", None),
                "domain": getattr(self, "domain", None),
                "task_type": getattr(self, "task_type", None),
                "description": getattr(self, "description", None),
                "is_trained": bool(getattr(self, "is_trained", False)),
                "metadata": asdict(getattr(self, "metadata", {})) if hasattr(self, "metadata") else {},
                "training_history": getattr(self, "training_history", []),
                "feedback_history": getattr(self, "feedback_history", []),
                "saved_at": datetime.now().isoformat(),
                "format": "toxo-pickle-snapshot-v1",
            }
            pickle.dump(snapshot, f, protocol=pickle.HIGHEST_PROTOCOL)

    def _save_toxo_format(self, path: Path):
        """Save in Toxo package format (.toxo) with ALL available components."""
        if path.suffix != ".toxo":
            path = path.with_suffix(".toxo")

        temp_dir = path.parent / f"{path.stem}_temp"
        temp_dir.mkdir(exist_ok=True)

        try:
            components_info: Dict[str, Any] = {}

            components_dir = temp_dir / "components"
            components_dir.mkdir(exist_ok=True)

            self._save_training_data(temp_dir, components_info)
            self._save_memory_systems(temp_dir, components_info)
            self._save_cognitive_systems(temp_dir, components_info)
            self._save_agent_framework(temp_dir, components_info)
            self._save_processing_pipeline(temp_dir, components_info)
            self._save_enterprise_features(temp_dir, components_info)

            # Save soft prompt if available
            if getattr(self, "soft_prompt_engine", None) and hasattr(self.soft_prompt_engine, "get_current_prompt"):
                try:
                    current_prompt = self.soft_prompt_engine.get_current_prompt()
                    if current_prompt is not None:
                        try:
                            import numpy as np  # type: ignore
                        except Exception as e:
                            raise RuntimeError(f"numpy is required to save soft prompts: {e}")

                        soft_prompt_path = components_dir / "soft_prompt.npy"
                        soft_prompt_meta_path = components_dir / "soft_prompt_meta.json"

                        np.save(soft_prompt_path, current_prompt)

                        training_stats = self.soft_prompt_engine.get_training_stats()
                        soft_prompt_metadata = {
                            "shape": list(getattr(current_prompt, "shape", ())),
                            "dtype": str(getattr(current_prompt, "dtype", "")),
                            "saved_at": datetime.now().isoformat(),
                            "training_stats": {
                                "iteration": training_stats.get("iteration", 0),
                                "current_loss": training_stats.get("current_loss", 0.0),
                                "best_loss": training_stats.get("best_loss", 0.0),
                                "loss_history": training_stats.get("loss_history", []),
                                "training_completed": training_stats.get("training_completed", False),
                                "total_updates": training_stats.get("total_updates", 0),
                            },
                            "version": "1.1",
                            "layer_id": self.name,
                            "layer_domain": self.domain,
                        }

                        with open(soft_prompt_meta_path, "w") as f:
                            json.dump(soft_prompt_metadata, f, indent=2)

                        components_info["has_soft_prompt"] = True
                        components_info["soft_prompt_shape"] = list(getattr(current_prompt, "shape", ()))
                        components_info["soft_prompt_iterations"] = training_stats.get("iteration", 0)
                    else:
                        components_info["has_soft_prompt"] = False
                except Exception as e:
                    self.logger.error(f"Failed to save soft prompt: {e}")
                    components_info["has_soft_prompt"] = False

            # Save memory state if available
            if getattr(self, "memory", None):
                memory_dir = temp_dir / "memory"
                memory_dir.mkdir(exist_ok=True)

                memory_state = self.memory.get_state()
                with open(memory_dir / "memory_state.json", "w") as f:
                    json.dump(memory_state, f, indent=2, default=str)

                embeddings = self.memory.get_all_embeddings()
                if embeddings is not None:
                    try:
                        import numpy as np  # type: ignore
                    except Exception as e:
                        raise RuntimeError(f"numpy is required to save embeddings: {e}")
                    np.save(memory_dir / "embeddings.npy", embeddings)

            # Save reranker if available
            if getattr(self, "reranker", None):
                reranker_state = self.reranker.get_state()
                with open(components_dir / "reranker_state.json", "w") as f:
                    json.dump(reranker_state, f, indent=2, default=str)

            # Save configuration
            config_dir = temp_dir / "config"
            config_dir.mkdir(exist_ok=True)
            self.config.save_to_file(config_dir / "toxo_config.yaml")

            layer_config = {
                "domain": self.domain,
                "task_type": self.task_type,
                "name": self.name,
                "description": self.description,
                "is_trained": self.is_trained,
                "training_history": getattr(self, "training_history", []),
                "performance_metrics": getattr(self.metadata, "performance_metrics", {}),
            }
            with open(config_dir / "layer_config.json", "w") as f:
                json.dump(layer_config, f, indent=2, default=str)

            if getattr(self, "master_prompt", None):
                try:
                    if hasattr(self.master_prompt, "to_dict"):
                        master_prompt_dict = self.master_prompt.to_dict()
                    elif hasattr(self.master_prompt, "__dict__"):
                        master_prompt_dict = self.master_prompt.__dict__
                    else:
                        master_prompt_dict = {"content": str(self.master_prompt)}

                    prompt_config = {
                        "master_prompt": master_prompt_dict,
                        "template_format": getattr(self.rich_prompt_engine, "template_format", "professional")
                        if getattr(self, "rich_prompt_engine", None)
                        else "professional",
                    }
                    with open(config_dir / "prompt_config.json", "w") as f:
                        json.dump(prompt_config, f, indent=2, default=str)
                except Exception as e:
                    self.logger.warning(f"Could not serialize master_prompt: {e}")
                    with open(config_dir / "prompt_config.json", "w") as f:
                        json.dump({"master_prompt": {"content": "Advanced prompt configuration"}, "template_format": "professional"}, f, indent=2)

            if os.environ.get("TOXO_ENABLE_LEGACY") == "1":
                self._create_standalone_toxo_module(temp_dir)

            metadata_dict = asdict(self.metadata)
            for key, value in list(metadata_dict.items()):
                if hasattr(value, "isoformat"):
                    metadata_dict[key] = value.isoformat()

            with open(temp_dir / "manifest.json", "w") as f:
                json.dump(
                    {
                        "format_version": "1.0",
                        "package_info": {
                            "name": self.name,
                            "domain": self.domain,
                            "task_type": self.task_type,
                            "description": self.description,
                            "created_at": self.metadata.created_at,
                            "updated_at": self.metadata.updated_at,
                            "version": self.metadata.version,
                        },
                        "components": {
                            "has_soft_prompt": bool(getattr(self, "soft_prompt_engine", None))
                            and self.soft_prompt_engine.get_current_prompt() is not None,
                            "has_memory": getattr(self, "memory", None) is not None,
                            "has_reranker": getattr(self, "reranker", None) is not None,
                            "has_rich_prompt": getattr(self, "rich_prompt_engine", None) is not None,
                            "has_training_data": components_info.get("has_training_data", False),
                            "has_memory_systems": components_info.get("has_memory_systems", False),
                            "has_cognitive_systems": components_info.get("has_cognitive_systems", False),
                            "has_agent_framework": components_info.get("has_agent_framework", False),
                            "has_processing_pipeline": components_info.get("has_processing_pipeline", False),
                            "has_enterprise_features": components_info.get("has_enterprise_features", False),
                            "training_examples_count": components_info.get("training_examples_count", 0),
                            "memory_types": components_info.get("memory_types", []),
                            "cognitive_types": components_info.get("cognitive_types", []),
                            "agent_types": components_info.get("agent_types", []),
                            "processing_types": components_info.get("processing_types", []),
                            "enterprise_types": components_info.get("enterprise_types", []),
                        },
                        "metadata": metadata_dict,
                    },
                    f,
                    indent=2,
                    default=str,
                )

            with zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as zf:
                for file_path in temp_dir.rglob("*"):
                    if file_path.is_file():
                        zf.write(file_path, file_path.relative_to(temp_dir))

            import shutil

            shutil.rmtree(temp_dir)
        except Exception:
            if temp_dir.exists():
                import shutil

                shutil.rmtree(temp_dir)
            raise

    def _save_training_data(self, temp_dir: Path, components_info: Dict[str, Any]):
        """Save comprehensive training data to .toxo file."""
        try:
            training_dir = temp_dir / "training"
            training_dir.mkdir(exist_ok=True)

            training_data = {
                "examples": getattr(self, "training_examples", []),
                "feedback": getattr(self, "feedback_data", []),
                "performance_metrics": getattr(getattr(self, "metadata", None), "performance_metrics", {}),
                "configurations": {
                    "domain": getattr(self, "domain", "general"),
                    "task_type": getattr(self, "task_type", "general"),
                    "training_config": getattr(self.config, "training", {}) if hasattr(self, "config") else {},
                },
                "soft_prompt_training": {},
                "evolutionary_training": {},
            }

            if getattr(self, "soft_prompt_engine", None):
                training_stats = self.soft_prompt_engine.get_training_stats()
                training_data["soft_prompt_training"] = {
                    "iteration": training_stats.get("iteration", 0),
                    "loss_history": training_stats.get("loss_history", []),
                    "best_loss": training_stats.get("best_loss", float("inf")),
                    "current_loss": training_stats.get("current_loss", 0.0),
                    "learning_rate_history": training_stats.get("learning_rate_history", []),
                    "momentum_history": training_stats.get("momentum_history", []),
                }

            with open(training_dir / "training_data.json", "w") as f:
                json.dump(training_data, f, indent=2, default=str)

            if training_data["examples"]:
                with open(training_dir / "training_examples.json", "w") as f:
                    json.dump(training_data["examples"], f, indent=2, default=str)

            if training_data["feedback"]:
                with open(training_dir / "feedback_data.json", "w") as f:
                    json.dump(training_data["feedback"], f, indent=2, default=str)

            components_info["has_training_data"] = True
            components_info["training_examples_count"] = len(training_data["examples"])
            components_info["feedback_count"] = len(training_data["feedback"])
        except Exception as e:
            self.logger.error(f"Failed to save training data: {e}")
            components_info["has_training_data"] = False

    def _save_memory_systems(self, temp_dir: Path, components_info: Dict[str, Any]):
        """Save all memory systems to .toxo file."""
        try:
            memory_dir = temp_dir / "memory"
            memory_dir.mkdir(exist_ok=True)

            memory_systems: Dict[str, Any] = {}

            if getattr(self, "memory", None):
                memory_systems["domain"] = self.memory.get_state()
                with open(memory_dir / "domain_memory.json", "w") as f:
                    json.dump(memory_systems["domain"], f, indent=2, default=str)

                embeddings = self.memory.get_all_embeddings()
                if embeddings is not None:
                    try:
                        import numpy as np  # type: ignore
                    except Exception as e:
                        raise RuntimeError(f"numpy is required to save embeddings: {e}")
                    np.save(memory_dir / "domain_embeddings.npy", embeddings)

            memory_components = [
                ("contextual_memory", "contextual"),
                ("episodic_memory", "episodic"),
                ("semantic_memory", "semantic"),
                ("working_memory", "working"),
                ("long_term_memory", "long_term"),
                ("procedural_memory", "procedural"),
            ]
            for attr_name, memory_type in memory_components:
                memory_component = getattr(self, attr_name, None)
                if memory_component and hasattr(memory_component, "get_state"):
                    try:
                        state = memory_component.get_state()
                        memory_systems[memory_type] = state
                        with open(memory_dir / f"{memory_type}_memory.json", "w") as f:
                            json.dump(state, f, indent=2, default=str)
                    except Exception as e:
                        self.logger.warning(f"Could not save {memory_type} memory: {e}")

            hierarchical = getattr(self, "hierarchical_memory", None)
            if hierarchical and hasattr(hierarchical, "get_state"):
                try:
                    hierarchical_state = hierarchical.get_state()
                    memory_systems["hierarchical"] = hierarchical_state
                    with open(memory_dir / "hierarchical_memory.json", "w") as f:
                        json.dump(hierarchical_state, f, indent=2, default=str)
                except Exception as e:
                    self.logger.warning(f"Could not save hierarchical memory: {e}")

            with open(memory_dir / "memory_systems.json", "w") as f:
                json.dump(memory_systems, f, indent=2, default=str)

            components_info["has_memory_systems"] = True
            components_info["memory_types"] = list(memory_systems.keys())
            components_info["memory_systems_count"] = len(memory_systems)
        except Exception as e:
            self.logger.error(f"Failed to save memory systems: {e}")
            components_info["has_memory_systems"] = False

    def _save_cognitive_systems(self, temp_dir: Path, components_info: Dict[str, Any]):
        """Save cognitive systems to .toxo file."""
        try:
            cognitive_dir = temp_dir / "cognitive"
            cognitive_dir.mkdir(exist_ok=True)

            cognitive_systems: Dict[str, Any] = {}

            for name in ["advanced_rag", "self_improvement", "feedback_learning", "performance_monitor"]:
                obj = getattr(self, name, None)
                if obj and hasattr(obj, "get_state"):
                    try:
                        state = obj.get_state()
                        cognitive_systems[name] = state
                        with open(cognitive_dir / f"{name}.json", "w") as f:
                            json.dump(state, f, indent=2, default=str)
                    except Exception as e:
                        self.logger.warning(f"Could not save {name}: {e}")

            with open(cognitive_dir / "cognitive_systems.json", "w") as f:
                json.dump(cognitive_systems, f, indent=2, default=str)

            components_info["has_cognitive_systems"] = True
            components_info["cognitive_types"] = list(cognitive_systems.keys())
            components_info["cognitive_systems_count"] = len(cognitive_systems)
        except Exception as e:
            self.logger.error(f"Failed to save cognitive systems: {e}")
            components_info["has_cognitive_systems"] = False

    def _save_agent_framework(self, temp_dir: Path, components_info: Dict[str, Any]):
        """Save agent framework to .toxo file."""
        try:
            agents_dir = temp_dir / "agents"
            agents_dir.mkdir(exist_ok=True)

            agent_systems: Dict[str, Any] = {}

            orchestrator = getattr(self, "orchestrator", None)
            if orchestrator and hasattr(orchestrator, "get_state"):
                try:
                    orchestrator_state = orchestrator.get_state()
                    agent_systems["orchestrator"] = orchestrator_state
                    with open(agents_dir / "orchestrator.json", "w") as f:
                        json.dump(orchestrator_state, f, indent=2, default=str)
                except Exception as e:
                    self.logger.warning(f"Could not save Agent Orchestrator: {e}")

            for agent_type in ["research", "creative", "analytical", "collaborative"]:
                agent = getattr(self, f"{agent_type}_agent", None)
                if agent and hasattr(agent, "get_state"):
                    try:
                        agent_state = agent.get_state()
                        agent_systems[agent_type] = agent_state
                        with open(agents_dir / f"{agent_type}_agent.json", "w") as f:
                            json.dump(agent_state, f, indent=2, default=str)
                    except Exception as e:
                        self.logger.warning(f"Could not save {agent_type} agent: {e}")

            comm = getattr(self, "agent_communication", None)
            if comm and hasattr(comm, "get_state"):
                try:
                    comm_state = comm.get_state()
                    agent_systems["communication"] = comm_state
                    with open(agents_dir / "communication.json", "w") as f:
                        json.dump(comm_state, f, indent=2, default=str)
                except Exception as e:
                    self.logger.warning(f"Could not save Agent Communication: {e}")

            with open(agents_dir / "agent_systems.json", "w") as f:
                json.dump(agent_systems, f, indent=2, default=str)

            components_info["has_agent_framework"] = True
            components_info["agent_types"] = list(agent_systems.keys())
            components_info["agent_systems_count"] = len(agent_systems)
        except Exception as e:
            self.logger.error(f"Failed to save agent framework: {e}")
            components_info["has_agent_framework"] = False

    def _save_processing_pipeline(self, temp_dir: Path, components_info: Dict[str, Any]):
        """Save processing pipeline to .toxo file."""
        try:
            processing_dir = temp_dir / "processing"
            processing_dir.mkdir(exist_ok=True)

            processing_systems: Dict[str, Any] = {}

            for attr, key, fname in [
                ("multimodal_processor", "multimodal", "multimodal_processor.json"),
                ("document_processor", "document", "document_processor.json"),
                ("advanced_document_processor", "advanced_document", "advanced_document_processor.json"),
            ]:
                obj = getattr(self, attr, None)
                if obj and hasattr(obj, "get_state"):
                    try:
                        state = obj.get_state()
                        processing_systems[key] = state
                        with open(processing_dir / fname, "w") as f:
                            json.dump(state, f, indent=2, default=str)
                    except Exception as e:
                        self.logger.warning(f"Could not save {attr}: {e}")

            with open(processing_dir / "processing_systems.json", "w") as f:
                json.dump(processing_systems, f, indent=2, default=str)

            components_info["has_processing_pipeline"] = True
            components_info["processing_types"] = list(processing_systems.keys())
            components_info["processing_systems_count"] = len(processing_systems)
        except Exception as e:
            self.logger.error(f"Failed to save processing pipeline: {e}")
            components_info["has_processing_pipeline"] = False

    def _save_enterprise_features(self, temp_dir: Path, components_info: Dict[str, Any]):
        """Save enterprise features to .toxo file."""
        try:
            enterprise_dir = temp_dir / "enterprise"
            enterprise_dir.mkdir(exist_ok=True)

            enterprise_systems: Dict[str, Any] = {}

            for attr, key, fname in [
                ("load_balancer", "load_balancer", "load_balancer.json"),
                ("cache_manager", "cache_manager", "cache_manager.json"),
                ("vector_store", "vector_store", "vector_store.json"),
                ("monitor", "monitor", "monitor.json"),
                ("analytics_manager", "analytics", "analytics.json"),
            ]:
                obj = getattr(self, attr, None)
                if obj and hasattr(obj, "get_state"):
                    try:
                        state = obj.get_state()
                        enterprise_systems[key] = state
                        with open(enterprise_dir / fname, "w") as f:
                            json.dump(state, f, indent=2, default=str)
                    except Exception as e:
                        self.logger.warning(f"Could not save {attr}: {e}")

            with open(enterprise_dir / "enterprise_systems.json", "w") as f:
                json.dump(enterprise_systems, f, indent=2, default=str)

            components_info["has_enterprise_features"] = True
            components_info["enterprise_types"] = list(enterprise_systems.keys())
            components_info["enterprise_systems_count"] = len(enterprise_systems)
        except Exception as e:
            self.logger.error(f"Failed to save enterprise features: {e}")
            components_info["has_enterprise_features"] = False

    def _create_standalone_toxo_module(self, temp_dir: Path):
        """Create an embedded shim for `.toxo` archives (legacy compatibility)."""
        shim = '''"""
TOXO embedded shim (legacy).

Generated into `.toxo` archives only when `TOXO_ENABLE_LEGACY=1`.
"""

from toxo.core.toxo_layer import ToxoLayer

__all__ = ["ToxoLayer"]
'''

        init_shim = '''"""
TOXO embedded package shim (legacy).
"""

from .toxo import ToxoLayer

__all__ = ["ToxoLayer"]
'''

        with open(temp_dir / "toxo.py", "w") as f:
            f.write(shim)

        with open(temp_dir / "__init__.py", "w") as f:
            f.write(init_shim)

    # The following methods are expected to exist on the layer and are left in the main class:
    # - _save_training_data
    # - _save_memory_systems
    # - _save_cognitive_systems
    # - _save_agent_framework
    # - _save_processing_pipeline
    # - _save_enterprise_features
    # - _create_standalone_toxo_module
    # - _save_json_format (implemented here)
    # - _save_pickle_format (implemented here)

