"""
High-level interface for model inference using Helium with streaming weights
"""
import os
from pathlib import Path
import json
import numpy as np
from typing import Dict, List, Optional, Union, Any, Tuple

from helium import HeliumMultiModal, HeliumModel, MultiModalModel
from helium.modality import ModalityType
from helium.tensor_ops import TensorOps
from helium.embedding import Embedding
from helium.positional_encoding import sinusoidal_positional_encoding
from helium.decoder_model import DecoderConfig
from helium.attention_utils import AttentionConfig, AttentionType
from helium.normalization import NormConfig, NormType
from helium.activations import softmax, gelu
from helium.safetensors import save_file

from .model_streaming import WeightManager, StreamBuffer, PrefetchManager
from .inference_runner import InferenceRunner

class HeliumStreamingModel:
    """High-level model interface with streaming support"""
    
    def __init__(self, 
                 model_url: str,
                 device_id: str = "vgpu0",
                 cache_dir: Optional[str] = None):
        self.model_url = model_url
        self.device_id = device_id
        self.cache_dir = cache_dir
        
        # Initialize streaming components
        self.runner = InferenceRunner(
            model_url=model_url,
            device_id=device_id,
            cache_dir=cache_dir
        )
        
        # Load model configuration
        self.config = self._load_config()
        
    def _load_config(self) -> Dict[str, Any]:
        """Load model configuration"""
        return self.runner.model_config
        
    def run(self, 
            inputs: Union[str, Dict[str, Any]], 
            **kwargs) -> Dict[str, Any]:
        """Run inference on inputs"""
        if isinstance(inputs, str):
            inputs = {'text': inputs}
            
        return self.runner.run_inference(inputs, **kwargs)
    
    def cleanup(self):
        """Cleanup resources"""
        self.runner.cleanup()
        
    @classmethod
    def from_pretrained(cls,
                       model_url: str,
                       device_id: str = "vgpu0",
                       **kwargs) -> "HeliumStreamingModel":
        """Create model from pretrained weights"""
        return cls(
            model_url=model_url,
            device_id=device_id,
            **kwargs
        )
        
        # Now create the real connection with the configured settings
        conn = duckdb.connect(db_url, config={'http_keep_alive': 'true'})
        conn.execute("INSTALL httpfs;")
        conn.execute("LOAD httpfs;")
        conn.execute("SET s3_endpoint='hf.co';")
        conn.execute("SET s3_use_ssl=true;")
        conn.execute("SET s3_url_style='path';")
        
        # Close temporary connection
        temp_conn.close()
        
        # Load config
        config = conn.execute(
            "SELECT config FROM model_configs WHERE model_key = ?",
            [model_key]
        ).fetchone()[0]
        self.config = json.loads(config)
        
        # Load state dict
        state_blob = conn.execute(
            "SELECT weights FROM model_weights WHERE model_key = ?",
            [model_key]
        ).fetchone()[0]
        
        state_dict = json.loads(state_blob)
        self.load_state_dict(state_dict)
        
    def to_device(self, device_id: str) -> None:
        """Move model to specified virtual GPU device"""
        self.device_id = device_id
        for module in self._modules.values():
            if hasattr(module, 'to_device'):
                module.to_device(device_id)
        
    def register_module(self, name: str, module: Any) -> None:
        self._modules[name] = module
        
    def register_parameter(self, name: str, param: np.ndarray) -> None:
        self._parameters[name] = param
        
    def register_buffer(self, name: str, buffer: np.ndarray) -> None:
        self._buffers[name] = buffer
        
    def state_dict(self) -> Dict[str, Any]:
        """Returns model state as a dictionary"""
        state = {}
        state.update(self._parameters)
        state.update(self._buffers)
        for name, module in self._modules.items():
            if hasattr(module, "state_dict"):
                state.update({
                    f"{name}.{k}": v 
                    for k, v in module.state_dict().items()
                })
        return state
        
    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
        """Loads model state from dictionary"""
        for name, param in state_dict.items():
            if "." in name:
                module_name, param_name = name.split(".", 1)
                if module_name in self._modules:
                    if hasattr(self._modules[module_name], "load_state_dict"):
                        self._modules[module_name].load_state_dict({param_name: param})
            else:
                if name in self._parameters:
                    self._parameters[name] = param
                elif name in self._buffers:
                    self._buffers[name] = param
                    
    def train(self, mode: bool = True) -> "HeliumStreamingModel":
        """Sets training mode"""
        self.training = mode
        for module in self._modules.values():
            if hasattr(module, "train"):
                module.train(mode)
        return self
        
    def eval(self) -> "HeliumModel":
        """Sets evaluation mode"""
        return self.train(False)
        
    def to_device(self, device_id: str) -> "HeliumModel":
        """Moves model to specified device"""
        for module in self._modules.values():
            if hasattr(module, "to_device"):
                module.to_device(device_id)
        return self

def main():
    """Example usage of streaming model"""
    
    def __init__(
        self,
        hidden_size: int = 1024,
        num_heads: int = 16,
        num_layers: int = 12,
        vocab_size: int = 50257,
        max_seq_len: int = 2048,
        device_id: str = "vgpu0"
    ):
        super().__init__()
        
        # Save config
        self.config = {
            "hidden_size": hidden_size,
            "num_heads": num_heads,
            "num_layers": num_layers,
            "vocab_size": vocab_size,
            "max_seq_len": max_seq_len
        }
        
        # Initialize virtual GPU system
        self.system = HeliumMultiModal(
            num_tensor_cores=1,
            memory_size=None  # Unlimited VRAM
        )
        self.device_id = device_id
        
        # Text components
        # Get the virtual GPU device
        self.system = HeliumMultiModal(num_tensor_cores=1)
        
        # Initialize components with device
        driver = self.system.gpu.tensor_cores[0]
        
        self.register_module("text_embedding", Embedding(
            vocab_size=vocab_size,
            embedding_dim=hidden_size,
            driver=driver,
            prefix="text_embed"
        ))
        
        # Generate positional encodings
        pos_enc = sinusoidal_positional_encoding(
            seq_len=max_seq_len,
            hidden_dim=hidden_size,
            driver=driver,
            prefix="pos_enc"
        )
        self.register_buffer("positional_encoding", pos_enc)
        
        # Decoder components
        decoder_config = DecoderConfig(
            output_modalities=[ModalityType.TEXT],
            hidden_dim=hidden_size,
            num_layers=num_layers,
            num_heads=num_heads,
            intermediate_size=hidden_size * 4,
            max_seq_len={ModalityType.TEXT: max_seq_len},
            vocab_size=vocab_size,
            use_cache=True
        )
        
        # Attention configuration
        attn_config = AttentionConfig(
            attention_type=AttentionType.SELF,
            hidden_size=hidden_size,
            num_heads=num_heads,
            head_dim=hidden_size // num_heads,
            dropout=0.1
        )
        self.register_buffer("attention_config", attn_config)
        
        # Normalization configuration
        norm_config = NormConfig(
            norm_type=NormType.LAYER,
            hidden_size=hidden_size,
            eps=1e-5
        )
        self.register_buffer("norm_config", norm_config)
        
        # Initialize weights
        self.register_parameter(
            "qkv_weights",
            np.random.randn(3, hidden_size, hidden_size).astype(np.float32) * 0.02
        )
        
        self.register_parameter(
            "norm_weight",
            np.ones(hidden_size).astype(np.float32)
        )
        
        self.register_parameter(
            "norm_bias",
            np.zeros(hidden_size).astype(np.float32)
        )
        
        # Cross-modal fusion weights
        self.register_parameter(
            "fusion_weight", 
            np.random.randn(hidden_size, hidden_size).astype(np.float32)
        )
        
    def forward(
        self,
        input_dict: Dict[str, np.ndarray],
        return_dict: bool = True
    ) -> Union[np.ndarray, Dict[str, np.ndarray]]:
        """Forward pass"""
        outputs = {}
        
        # Process each modality
        for modality, inputs in input_dict.items():
            if modality == "text":
                # Text processing
                embeds = self._modules["text_embedding"](inputs)
                pos_embeds = embeds + self._buffers["positional_encoding"][:inputs.shape[1]]
                
                # Layer normalization
                mean = pos_embeds.mean(axis=-1, keepdims=True)
                var = ((pos_embeds - mean) ** 2).mean(axis=-1, keepdims=True)
                hidden = (pos_embeds - mean) / np.sqrt(var + self._buffers["norm_config"].eps)
                hidden = hidden * self._parameters["norm_weight"] + self._parameters["norm_bias"]
                
                # Self attention
                qkv = np.einsum('...d,hdi->...hi', hidden, self._parameters["qkv_weights"])
                q, k, v = np.split(qkv, 3, axis=-2)
                
                # Scaled dot-product attention
                attn_weights = np.matmul(q, k.transpose(-2, -1)) / np.sqrt(hidden.shape[-1])
                attn_weights = softmax(attn_weights, axis=-1)
                attn_output = np.matmul(attn_weights, v)
                
                # Apply GELU activation
                hidden = gelu(attn_output)
                outputs["text_features"] = hidden
                
            elif modality == "image":
                # Image processing
                outputs["image_features"] = self.system.process_batch({
                    ModalityType.IMAGE: inputs
                })
                
            elif modality == "audio":
                # Audio processing  
                outputs["audio_features"] = self.system.process_batch({
                    ModalityType.AUDIO: inputs
                })
                
        # Fuse modalities if multiple present
        if len(outputs) > 1:
            fusion = sum(outputs.values())
            fusion = fusion @ self._parameters["fusion_weight"]
            outputs["fused_features"] = fusion
            
        return outputs if return_dict else fusion
        
    def generate(
        self, 
        inputs: Union[np.ndarray, Dict[str, np.ndarray]],
        max_length: int = 100,
        **kwargs
    ) -> np.ndarray:
        """Generate sequence"""
        if isinstance(inputs, dict):
            # Get fused representation for multi-modal input
            hidden = self.forward(inputs, return_dict=False)
        else:
            # Single modality (text) input
            embeds = self._modules["text_embedding"](inputs)
            pos_embeds = self._modules["pos_encoding"](embeds)
            hidden = self._modules["decoder"](pos_embeds)
            
        # Auto-regressive generation
        generated = []
        for _ in range(max_length):
            next_token = self._modules["decoder"].predict_next(hidden)
            generated.append(next_token)
            
            # Update hidden state
            next_embeds = self._modules["text_embedding"](next_token)
            next_pos = self._modules["pos_encoding"](next_embeds)
            hidden = self._modules["decoder"](next_pos, hidden)
            
        return np.array(generated)
        
    def save_pretrained(self, path: str) -> None:
        """Save model weights and config"""
        os.makedirs(path, exist_ok=True)
        
        # Save config
        with open(os.path.join(path, "config.json"), "w") as f:
            json.dump(self.config, f, indent=2)
            
        # Save weights
        save_file(self.state_dict(), os.path.join(path, "model.safetensors"))
        
    @classmethod
    def from_pretrained(
        cls,
        model_id: str = "openai-oss-20b",
        device_id: str = "vgpu0",
        cache_dir: Optional[str] = None,
        **kwargs
    ) -> "MultiModalModel":
        """Load pretrained model from HuggingFace Hub"""
        from .model_loader import download_model, store_in_device_db
        from config import get_db_url
        
        # Download model from HuggingFace
        local_path = download_model(model_id, cache_dir)
        
        # Store in device DB
        db_url = get_db_url()
        model_key = store_in_device_db(local_path, db_url)
        
        # Initialize model on virtual GPU
        model = cls()
        model.load_state_from_db(model_key, device_id)
        model.to_device(device_id)
        return model
        device_db_url = get_db_url()
        store_in_device_db(local_path, device_db_url, model_id)
        
        # Connect to device DB
        conn = duckdb.connect(device_db_url)
        
        # Load config from DB
        config = conn.execute(
            "SELECT config FROM model_configs WHERE model_id = ?",
            [model_id]
        ).fetchone()[0]
        config = json.loads(config)
        
        # Create model
        model = cls(**config, device_id=device_id, **kwargs)
        state_dict = load_file(os.path.join(path, "model.safetensors"))
        model.load_state_dict(state_dict)
        
        return model

def main():
    """Example usage"""
    # Create model
    model = MultiModalModel(
        hidden_size=1024,
        num_heads=16,
        num_layers=12,
        device_id="vgpu0"
    )
    
    # Example inputs
    inputs = {
        "text": np.random.randint(0, 50257, (1, 64)),  # token_ids
        "image": np.random.randn(1, 3, 224, 224),      # image tensor
        "audio": np.random.randn(1, 1, 16000)          # audio waveform
    }
    
    # Inference
    model.eval()  # Set to evaluation mode
    outputs = model(inputs)
        
    print("Output features shapes:")
    for k, v in outputs.items():
        print(f"  {k}: {v.shape}")
        
    # Generate from multi-modal context
    generated = model.generate(inputs, max_length=20)
    print("\nGenerated sequence shape:", generated.shape)
    
    # Save model
    model.save_pretrained("model_checkpoint")
    
    # Load model
    loaded_model = MultiModalModel.from_pretrained("model_checkpoint")
    print("\nSuccessfully loaded model from checkpoint")

if __name__ == "__main__":
    main()