import os
import torch
import torch.nn as nn
from typing import Dict, Any, List, Optional
from safetensors.torch import save_file, load_file

from xorfice.models.components.lora import get_lora_parameters, LoRALinear

class PersonalizationManager:
    """
    SOTA Personalization Manager for Xoron-Dev.
    Handles per-user LoRA weight isolation and on-the-fly online training.
    """
    def __init__(self, adapter_dir: str = "adapters", lr: float = 5e-5):
        self.adapter_dir = adapter_dir
        self.lr = lr
        os.makedirs(self.adapter_dir, exist_ok=True)

    def get_user_adapter_path(self, user_id: str) -> str:
        return os.path.join(self.adapter_dir, f"user_{user_id}.safetensors")

    def save_user_weights(self, model: nn.Module, user_id: str):
        """Saves current LoRA parameters for a specific user."""
        path = self.get_user_adapter_path(user_id)
        lora_params = {}
        for name, param in model.named_parameters():
            if 'lora_A' in name or 'lora_B' in name or 'magnitude' in name:
                # Store relative to module name for portability
                lora_params[name] = param.data.detach().cpu()
        
        save_file(lora_params, path)
        print(f"[Personalization] Saved weights for user {user_id} to {path}")

    def load_user_weights(self, model: nn.Module, user_id: str):
        """Loads LoRA parameters for a specific user into the model."""
        path = self.get_user_adapter_path(user_id)
        if not os.path.exists(path):
            print(f"[Personalization] No saved weights for user {user_id}. Using base model.")
            return

        state_dict = load_file(path)
        model_state_dict = model.state_dict()
        
        for name, tensor in state_dict.items():
            if name in model_state_dict:
                model_state_dict[name].copy_(tensor.to(model_state_dict[name].device))
        
        print(f"[Personalization] Loaded weights for user {user_id} from {path}")

    def train_on_fly(self, model: nn.Module, user_id: str, prompt_ids: torch.Tensor, response_ids: torch.Tensor):
        """
        SOTA On-the-fly Training: Performs a single optimization step for the user.
        Updates ONLY LoRA parameters.
        """
        # 1. Prepare inputs and labels for causal language modeling
        # input_ids = [prompt, response]
        # labels = [-100... -100, response]
        input_ids = torch.cat([prompt_ids, response_ids], dim=1)
        labels = input_ids.clone()
        labels[:, :prompt_ids.shape[1]] = -100 # Mask prompt tokens
        
        # 2. Setup Optimizer for LoRA params only
        lora_params = get_lora_parameters(model)
        if not lora_params:
            print("[Personalization] Warning: No LoRA parameters found in model. Skipping training.")
            return
            
        optimizer = torch.optim.AdamW(lora_params, lr=self.lr)
        
        # 3. Training Step
        model.train()
        optimizer.zero_grad()
        
        outputs = model(input_ids=input_ids, labels=labels)
        loss = outputs.loss
        
        if loss is not None:
            loss.backward()
            optimizer.step()
            print(f"[Personalization] Online training step for user {user_id} completed. Loss: {loss.item():.4f}")
            
            # 4. Save updated weights
            self.save_user_weights(model, user_id)
        
        model.eval()
class VoiceManager:
    """
    SOTA Voice Personalization Manager.
    Handles speaker embedding extraction for Zero-Shot Voice Cloning.
    """
    def __init__(self, voice_dir: str = "voices"):
        self.voice_dir = voice_dir
        self.user_voices: Dict[str, torch.Tensor] = {}
        os.makedirs(self.voice_dir, exist_ok=True)

    def get_voice_path(self, user_id: str) -> str:
        return os.path.join(self.voice_dir, f"voice_{user_id}.pt")

    def extract_speaker_embedding(self, model: nn.Module, audio_waveform: torch.Tensor) -> torch.Tensor:
        """
        Extracts a fixed-length speaker embedding from reference audio.
        Uses the model's internal audio encoder for consistency.
        """
        with torch.no_grad():
            # Real SOTA implementation: 
            # Xoron's encode_audio typically returns either a direct tensor of shape (B, T, D) 
            # or a tuple/dictionary containing the hidden states.
            # We extract it and take the mean across the sequence length (T) to get a global
            # speaker representation of shape (B, D).
            features = model.encode_audio(audio_waveform)
            
            if isinstance(features, tuple):
                features = features[0]
            elif isinstance(features, dict):
                features = features.get("last_hidden_state", features.get("audio_embeds"))
                
            if isinstance(features, torch.Tensor):
                # features shape: (Batch, Sequence/Time, Hidden_Dim)
                # Mean pooling across the temporal dimension gives a solid global speaker embedding
                embedding = features.mean(dim=1) # (B, D)
                return embedding
                
            raise TypeError(f"Unexpected output type from encode_audio: {type(features)}")

    def set_user_voice(self, user_id: str, embedding: torch.Tensor):
        """Saves speaker embedding for a specific user."""
        self.user_voices[user_id] = embedding.cpu()
        path = self.get_voice_path(user_id)
        torch.save(embedding.cpu(), path)
        print(f"[VoiceManager] Voice saved for user {user_id}")

    def load_user_voice(self, user_id: str) -> Optional[torch.Tensor]:
        """Loads speaker embedding for a specific user."""
        if user_id in self.user_voices:
            return self.user_voices[user_id]
        
        path = self.get_voice_path(user_id)
        if os.path.exists(path):
            embedding = torch.load(path)
            self.user_voices[user_id] = embedding
            return embedding
        return None

    def list_available_voices(self) -> List[str]:
        return [f.replace("voice_", "").replace(".pt", "") for f in os.listdir(self.voice_dir) if f.endswith(".pt")]
        
    def delete_voice(self, user_id: str) -> bool:
        """Deletes a user's voice embedding securely from disk and cache."""
        if user_id in self.user_voices:
            del self.user_voices[user_id]
        path = self.get_voice_path(user_id)
        if os.path.exists(path):
            os.remove(path)
            print(f"[VoiceManager] Voice deleted securely for user {user_id}")
            return True
        return False
