"""Domain Adaptation Manager for integrating with ModelManager.

This module provides the main integration point between the domain adaptation system
and the existing ModelManager, enabling seamless domain-specific transcription.
"""

import logging
import os
from pathlib import Path
from typing import Dict, Optional, Any, List, Union

import torch
from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments

from .domain_adaptation import DomainAdapter, DomainDetector
from .model_manager import ModelManager

logger = logging.getLogger(__name__)


class DomainAdaptationManager:
    """Main manager for domain adaptation system integration.
    
    Integrates LoRA adapters and domain detection with the existing ModelManager
    to provide domain-specific transcription capabilities.
    """
    
    def __init__(self, model_manager: Optional[ModelManager] = None):
        """Initialize the domain adaptation manager.
        
        Args:
            model_manager: Optional ModelManager instance, creates new one if None
        """
        self.model_manager = model_manager or ModelManager()
        self.domain_adapter = DomainAdapter()
        self.domain_detector = DomainDetector()
        
        # Initialize with pre-trained domain adapters
        self._load_default_adapters()
        
        logger.info("Domain adaptation manager initialized")
    
    def _load_default_adapters(self) -> None:
        """Load default domain adapters if they exist."""
        default_domains = {
            "technical": "models/adapters/technical_adapter",
            "medical": "models/adapters/medical_adapter", 
            "academic": "models/adapters/academic_adapter"
        }
        
        for domain_name, path in default_domains.items():
            if os.path.exists(path):
                try:
                    self.domain_adapter.load_adapter(domain_name, path)
                    logger.info(f"Loaded default adapter for domain: {domain_name}")
                except Exception as e:
                    logger.warning(f"Failed to load default adapter for {domain_name}: {e}")
            else:
                logger.info(f"No default adapter found for domain: {domain_name}")
    
    def transcribe_with_domain_adaptation(
        self, 
        audio: Union[str, bytes, Path], 
        auto_detect: bool = True, 
        domain: Optional[str] = None,
        **kwargs
    ) -> str:
        """Transcribe audio with appropriate domain adaptation.
        
        Args:
            audio: Audio input (file path, bytes, or Path object)
            auto_detect: Whether to automatically detect domain from initial transcription
            domain: Specific domain to use (overrides auto_detect if provided)
            **kwargs: Additional arguments for transcription
            
        Returns:
            Domain-adapted transcription
        """
        logger.info("Starting domain-adapted transcription")
        
        # Get initial transcription from base model
        initial_transcription = self.model_manager.transcribe(audio, **kwargs)
        
        # Determine domain to use
        if domain is not None:
            detected_domain = domain
            logger.info(f"Using specified domain: {detected_domain}")
        elif auto_detect:
            detected_domain = self.domain_detector.detect_domain(initial_transcription)
            logger.info(f"Auto-detected domain: {detected_domain}")
        else:
            detected_domain = "general"
            logger.info("Using general domain (no adaptation)")
        
        # Use domain-specific adapter if available and not general
        if detected_domain != "general" and detected_domain in self.domain_adapter.domain_adapters:
            try:
                adapter_model = self.domain_adapter.switch_adapter(detected_domain)
                # For now, return the initial transcription with domain info
                # In a full implementation, this would use the adapter for re-transcription
                enhanced_transcription = f"[{detected_domain.upper()}] {initial_transcription}"
                logger.info(f"Applied {detected_domain} domain adaptation")
                return enhanced_transcription
            except Exception as e:
                logger.warning(f"Failed to apply domain adaptation: {e}")
                return initial_transcription
        else:
            return initial_transcription
    
    def train_custom_domain(
        self, 
        domain_name: str, 
        training_data: Any,
        output_dir: str = "models/adapters"
    ) -> None:
        """Train a new domain adapter on custom data.
        
        Args:
            domain_name: Name of the new domain
            training_data: Training dataset
            output_dir: Directory to save the trained adapter
        """
        logger.info(f"Training custom domain adapter: {domain_name}")
        
        # Create new adapter if it doesn't exist
        if domain_name not in self.domain_adapter.domain_adapters:
            self.domain_adapter.create_adapter(domain_name)
        
        adapter_model = self.domain_adapter.domain_adapters[domain_name]
        
        # Set up trainer
        trainer = self._setup_trainer(adapter_model, output_dir)
        
        # Train the adapter
        trainer.train(training_data)
        
        # Save the trained adapter
        save_path = os.path.join(output_dir, f"{domain_name}_adapter")
        self.domain_adapter.save_adapter(domain_name, save_path)
        
        logger.info(f"Custom domain adapter training completed: {domain_name}")
    
    def _setup_trainer(self, model: Any, output_dir: str) -> Seq2SeqTrainer:
        """Set up a trainer for adapter fine-tuning.
        
        Args:
            model: Model to train
            output_dir: Output directory for training artifacts
            
        Returns:
            Configured Seq2SeqTrainer
        """
        training_args = Seq2SeqTrainingArguments(
            output_dir=output_dir,
            per_device_train_batch_size=8,
            gradient_accumulation_steps=4,
            learning_rate=5e-5,
            num_train_epochs=3,
            save_strategy="epoch",
            logging_steps=100,
            evaluation_strategy="epoch",
            save_total_limit=2,
            load_best_model_at_end=True,
            metric_for_best_model="eval_loss",
            greater_is_better=False,
            warmup_steps=500,
            weight_decay=0.01,
            fp16=torch.cuda.is_available(),
            dataloader_pin_memory=False,
            remove_unused_columns=False,
            push_to_hub=False
        )
        
        return Seq2SeqTrainer(
            model=model,
            args=training_args,
            # Additional trainer parameters would be configured here
        )
    
    def get_available_domains(self) -> List[str]:
        """Get list of available domain adapters.
        
        Returns:
            List of domain names with available adapters
        """
        return self.domain_adapter.list_adapters()
    
    def get_domain_info(self, domain_name: str) -> Dict[str, Any]:
        """Get information about a specific domain adapter.
        
        Args:
            domain_name: Name of the domain
            
        Returns:
            Dictionary with domain adapter information
        """
        return self.domain_adapter.get_adapter_info(domain_name)
    
    def train_domain_detector(
        self, 
        texts: List[str], 
        domain_labels: List[str],
        save_path: Optional[str] = None
    ) -> Dict[str, Any]:
        """Train the domain detector on labeled data.
        
        Args:
            texts: Training text samples
            domain_labels: Corresponding domain labels
            save_path: Optional path to save the trained detector
            
        Returns:
            Training metrics
        """
        logger.info(f"Training domain detector on {len(texts)} samples")
        
        # Train the detector
        self.domain_detector.train(texts, domain_labels)
        
        # Save if path provided
        if save_path:
            self.domain_detector.save_model(save_path)
        
        logger.info("Domain detector training completed")
        
        return {"status": "completed", "samples": len(texts)}
    
    def evaluate_domain_detector(
        self, 
        test_texts: List[str], 
        test_labels: List[str]
    ) -> Dict[str, Any]:
        """Evaluate the domain detector performance.
        
        Args:
            test_texts: Test text samples
            test_labels: True domain labels
            
        Returns:
            Evaluation metrics
        """
        return self.domain_detector.evaluate(test_texts, test_labels)
    
    def detect_domain(self, text: str, threshold: float = 0.6) -> str:
        """Detect the domain of a given text.
        
        Args:
            text: Text to classify
            threshold: Confidence threshold
            
        Returns:
            Detected domain name
        """
        return self.domain_detector.detect_domain(text, threshold)
    
    def get_domain_probabilities(self, text: str) -> Dict[str, float]:
        """Get probability scores for all domains.
        
        Args:
            text: Text to classify
            
        Returns:
            Dictionary mapping domain names to probability scores
        """
        return self.domain_detector.get_domain_probabilities(text)
    
    def remove_domain_adapter(self, domain_name: str) -> None:
        """Remove a domain adapter from memory.
        
        Args:
            domain_name: Name of the domain to remove
        """
        self.domain_adapter.remove_adapter(domain_name)
        logger.info(f"Removed domain adapter: {domain_name}")
    
    def cleanup(self) -> None:
        """Clean up resources and free memory."""
        logger.info("Cleaning up domain adaptation manager")
        
        # Remove all adapters
        for domain_name in list(self.domain_adapter.domain_adapters.keys()):
            self.domain_adapter.remove_adapter(domain_name)
        
        # Clear CUDA cache
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        logger.info("Domain adaptation manager cleanup completed")
