"""
Configuration management for Toxo.

This module handles all configuration aspects including:
- API credentials and endpoints
- Training parameters and optimization settings
- Memory and storage configurations
- Performance and scaling settings
"""

import os
import json
try:
    import yaml  # type: ignore
except Exception:
    yaml = None
from pathlib import Path
from typing import Dict, Any, Optional, Union, List
from dataclasses import dataclass, field, asdict
from pydantic import BaseModel, Field, validator
import logging

from ..utils.exceptions import ConfigError, ValidationError


@dataclass
class GeminiConfig:
    """Configuration for Gemini API integration."""
    api_key: Optional[str] = None  # Will be dynamically set
    model: str = "gemini-2.5-flash-lite"
    temperature: float = 0.7
    max_tokens: int = 8192
    timeout: int = 30
    max_retries: int = 3  # Fixed: was retry_attempts
    retry_delay: float = 1.0
    test_mode: bool = False  # Set to False since we have a real API key
    
    def __post_init__(self):
        """Validate Gemini configuration."""
        # In test mode, we don't require API keys
        if hasattr(self, '_test_mode') and self._test_mode:
            if not self.api_key:
                self.api_key = "test_key_for_demo"
            return
        
        # Use SecurityManager if available, else env vars
        try:
            from .security import get_security_manager
            security_manager = get_security_manager()
            self.api_key = security_manager.get_api_key("gemini", self.api_key)
        except Exception:
            pass  # Fall through to env/param handling
        
        if not self.api_key:
            self.api_key = (
                os.getenv("GEMINI_API_KEY") or 
                os.getenv("GOOGLE_API_KEY")
            )
        
        # API key is intentionally optional at config creation time; it must be provided
        # before making network calls (e.g. via `layer.setup_api_key(...)`).
        if self.api_key and not self.api_key.startswith(("AIza", "test")):
            raise ConfigError("Invalid Gemini API key format. Key should start with 'AIza'.")
        
        # Add validation
        if not isinstance(self.temperature, (int, float)) or not 0 <= self.temperature <= 1:
            raise ValidationError("Temperature must be between 0 and 1")
        if not isinstance(self.max_tokens, int) or self.max_tokens <= 0:
            raise ValidationError("Max tokens must be a positive integer")
        if not isinstance(self.max_retries, int) or self.max_retries < 0:
            raise ValidationError("Max retries must be a non-negative integer")


@dataclass 
class TrainingConfig:
    """Configuration for training parameters."""
    # SPSA (Simultaneous Perturbation Stochastic Approximation) settings
    learning_rate: float = 0.05
    perturbation_scale: float = 0.1
    batch_size: int = 16
    max_epochs: int = 10
    early_stopping: bool = True  # Fixed: was early_stopping_patience
    patience: int = 3  # Fixed: renamed from early_stopping_patience
    
    # Evolutionary optimization settings
    population_size: int = 20
    mutation_rate: float = 0.1
    crossover_rate: float = 0.8
    elite_ratio: float = 0.2
    
    # Advanced training settings
    curriculum_learning: bool = True
    adaptive_learning_rate: bool = True
    momentum: float = 0.9
    gradient_clipping: float = 2.0
    
    # Validation and metrics
    validation_split: float = 0.2
    metrics: List[str] = field(default_factory=lambda: ["accuracy", "bleu", "rouge"])
    save_best_only: bool = True
    
    def __post_init__(self):
        """Validate training configuration."""
        if not isinstance(self.learning_rate, (int, float)) or not 0 < self.learning_rate <= 1:
            raise ValidationError("Learning rate must be between 0 and 1 (exclusive of 0, inclusive of 1)")
        if not isinstance(self.perturbation_scale, (int, float)) or not 0 < self.perturbation_scale < 1:
            raise ValidationError("Perturbation scale must be between 0 and 1")
        if not isinstance(self.batch_size, int) or self.batch_size < 1:
            raise ValidationError("Batch size must be a positive integer")
        if not isinstance(self.max_epochs, int) or self.max_epochs < 1:
            raise ValidationError("Max epochs must be a positive integer")
        if not isinstance(self.patience, int) or self.patience < 1:
            raise ValidationError("Patience must be a positive integer")


@dataclass
class MemoryConfig:
    """Configuration for memory system."""
    # Memory hierarchy settings - Fixed: updated default values to match tests
    working_memory_size: int = 1000
    session_memory_size: int = 5000  # Fixed: was 10000
    domain_memory_size: int = 10000  # Fixed: was 100000
    global_memory_size: int = 100000  # Fixed: was 1000000
    
    # Memory consolidation settings
    consolidation_threshold: float = 0.8
    consolidation_frequency: int = 100  # interactions
    decay_rate: float = 0.01
    
    # Vector database settings
    embedding_dimension: int = 768
    similarity_threshold: float = 0.7
    max_retrievals: int = 10
    
    # Storage settings
    persistence_enabled: bool = True
    backup_frequency: int = 1000  # interactions
    compression_enabled: bool = True
    
    def __post_init__(self):
        """Validate memory configuration."""
        # Add proper validation for negative values
        if not isinstance(self.working_memory_size, int) or self.working_memory_size < 1:
            raise ValidationError("Working memory size must be a positive integer")
        if not isinstance(self.session_memory_size, int) or self.session_memory_size < 1:
            raise ValidationError("Session memory size must be a positive integer")
        if not isinstance(self.domain_memory_size, int) or self.domain_memory_size < 1:
            raise ValidationError("Domain memory size must be a positive integer")
        if not isinstance(self.embedding_dimension, int) or self.embedding_dimension not in [384, 512, 768, 1024, 1536]:
            raise ValidationError("Embedding dimension must be one of: 384, 512, 768, 1024, 1536")


@dataclass
class PromptConfig:
    """Configuration for prompt engineering."""
    # Soft prompt settings - Fixed: updated default value to match tests
    soft_prompt_length: int = 10  # Fixed: was 50
    soft_prompt_init: str = "random"  # random, zeros, ones, learned
    
    # Template settings
    use_system_prompt: bool = True
    use_few_shot_examples: bool = True
    max_examples: int = 5
    template_format: str = "structured"  # Added missing attribute
    
    # Dynamic prompt settings
    context_window_size: int = 4096
    context_compression: bool = True
    adaptive_context: bool = True
    
    # Advanced prompt features
    multi_step_reasoning: bool = True
    chain_of_thought: bool = True
    self_reflection: bool = True
    
    def __post_init__(self):
        """Validate prompt configuration."""
        if not isinstance(self.soft_prompt_length, int) or self.soft_prompt_length < 1 or self.soft_prompt_length > 200:
            raise ValidationError("Soft prompt length must be between 1 and 200")
        if self.soft_prompt_init not in ["random", "zeros", "ones", "learned"]:
            raise ValidationError("Soft prompt init must be one of: random, zeros, ones, learned")
        if self.template_format not in ["structured", "freeform", "json", "yaml"]:
            raise ValidationError("Template format must be one of: structured, freeform, json, yaml")


@dataclass
class PerformanceConfig:
    """Configuration for performance optimization."""
    # Caching settings
    response_caching: bool = True
    cache_ttl: int = 3600  # seconds
    cache_size: int = 10000
    
    # Concurrency settings
    max_concurrent_requests: int = 10
    request_timeout: int = 30
    connection_pool_size: int = 20
    
    # Optimization settings
    batch_processing: bool = True
    streaming_responses: bool = False
    compression: bool = True
    
    # Resource management
    max_memory_usage: int = 4096  # MB
    cleanup_frequency: int = 3600  # seconds
    
    def __post_init__(self):
        """Validate performance configuration."""
        if not isinstance(self.max_concurrent_requests, int) or self.max_concurrent_requests < 1:
            raise ValidationError("Max concurrent requests must be a positive integer")
        if not isinstance(self.cache_size, int) or self.cache_size < 0:
            raise ValidationError("Cache size must be non-negative")


@dataclass
class SecurityConfig:
    """Configuration for security settings."""
    # API security
    api_key_validation: bool = True
    request_signing: bool = False
    encrypt_storage: bool = True
    
    # Data protection
    pii_detection: bool = True
    data_anonymization: bool = False
    audit_logging: bool = True
    
    # Access control
    role_based_access: bool = False
    session_management: bool = True
    token_expiry: int = 86400  # seconds


class ToxoConfig:
    """
    Main configuration class for Toxo platform.
    
    This class manages all configuration aspects and provides methods for
    loading, saving, and validating configurations.
    """
    
    def __init__(
        self,
        gemini: Optional[GeminiConfig] = None,
        training: Optional[TrainingConfig] = None,
        memory: Optional[MemoryConfig] = None,
        prompt: Optional[PromptConfig] = None,
        performance: Optional[PerformanceConfig] = None,
        security: Optional[SecurityConfig] = None,
        config_path: Optional[Union[str, Path]] = None,
        data_dir: Optional[Union[str, Path]] = None,
        **kwargs
    ):
        """
        Initialize Toxo configuration.
        
        Args:
            gemini: GeminiConfig instance
            training: TrainingConfig instance
            memory: MemoryConfig instance
            prompt: PromptConfig instance
            performance: PerformanceConfig instance
            security: SecurityConfig instance
            config_path: Path to configuration file
            data_dir: Data directory path
            **kwargs: Additional configuration parameters
        """
        # Store test_mode first
        test_mode = kwargs.get('test_mode', False)
        
        # Set data directory
        self.data_dir = Path(data_dir) if data_dir else Path.home() / ".toxo" / "data"
        self.data_dir.mkdir(parents=True, exist_ok=True)
        
        if config_path and Path(config_path).exists():
            self.load_from_file(config_path)
        else:
            # Initialize sub-configurations
            if gemini is None:
                gemini = GeminiConfig()
                if test_mode:
                    gemini._test_mode = test_mode
            
            self.gemini = gemini or GeminiConfig()
            self.training = training or TrainingConfig()
            self.memory = memory or MemoryConfig()
            self.prompt = prompt or PromptConfig()
            self.performance = performance or PerformanceConfig()
            self.security = security or SecurityConfig()
        
        # Override with kwargs
        if kwargs:
            self.update(**kwargs)
        
        # Validate final configuration
        self.validate()
    
    def load_from_file(self, config_path: Union[str, Path]) -> None:
        """
        Load configuration from file.
        
        Args:
            config_path: Path to configuration file (JSON or YAML)
        """
        config_path = Path(config_path)
        
        if not config_path.exists():
            raise ConfigError(f"Configuration file not found: {config_path}")
        
        try:
            with open(config_path, 'r') as f:
                if config_path.suffix.lower() in ['.yaml', '.yml']:
                    config_data = yaml.safe_load(f)
                elif config_path.suffix.lower() == '.json':
                    config_data = json.load(f)
                else:
                    raise ConfigError(f"Unsupported config file format: {config_path.suffix}")
            
            # Create component configs from loaded data
            if "gemini" in config_data:
                self.gemini = GeminiConfig(**config_data["gemini"])
            if "training" in config_data:
                self.training = TrainingConfig(**config_data["training"])
            if "memory" in config_data:
                self.memory = MemoryConfig(**config_data["memory"])
            if "prompt" in config_data:
                self.prompt = PromptConfig(**config_data["prompt"])
            if "performance" in config_data:
                self.performance = PerformanceConfig(**config_data["performance"])
            if "security" in config_data:
                self.security = SecurityConfig(**config_data["security"])
            
        except Exception as e:
            raise ConfigError(f"Failed to load configuration: {str(e)}")
    
    def save_to_file(self, config_path: Union[str, Path], format: str = "yaml") -> None:
        """
        Save configuration to file.
        
        Args:
            config_path: Path to save configuration
            format: File format ('yaml' or 'json')
        """
        config_path = Path(config_path)
        config_data = self.to_dict()
        
        try:
            with open(config_path, 'w') as f:
                if format.lower() == 'yaml':
                    yaml.dump(config_data, f, default_flow_style=False, indent=2)
                elif format.lower() == 'json':
                    json.dump(config_data, f, indent=2)
                else:
                    raise ConfigError(f"Unsupported format: {format}")
                    
        except Exception as e:
            raise ConfigError(f"Failed to save configuration: {str(e)}")
    
    def update(self, **kwargs) -> None:
        """
        Update configuration with new values.
        
        Args:
            **kwargs: Configuration parameters to update
        """
        for section, values in kwargs.items():
            if hasattr(self, section) and isinstance(values, dict):
                section_config = getattr(self, section)
                for key, value in values.items():
                    if hasattr(section_config, key):
                        setattr(section_config, key, value)
                        # Re-validate the section after updates
                        if hasattr(section_config, '__post_init__'):
                            section_config.__post_init__()
                    else:
                        logging.warning(f"Unknown config parameter: {section}.{key}")
            else:
                logging.warning(f"Unknown config section: {section}")
    
    def validate(self) -> None:
        """Validate the entire configuration."""
        try:
            # Each dataclass validates itself in __post_init__
            # Additional cross-section validation can be added here
            
            # Example: Check memory settings compatibility
            if self.memory.working_memory_size > self.memory.session_memory_size:
                raise ValidationError("Working memory size cannot exceed session memory size")
            
            # Check performance vs memory settings
            estimated_memory = (
                self.memory.embedding_dimension * 
                self.memory.domain_memory_size * 4  # 4 bytes per float
            ) // (1024 * 1024)  # Convert to MB
            
            if estimated_memory > self.performance.max_memory_usage:
                logging.warning(
                    f"Estimated memory usage ({estimated_memory}MB) may exceed limit "
                    f"({self.performance.max_memory_usage}MB)"
                )
                
        except Exception as e:
            raise ValidationError(f"Configuration validation failed: {str(e)}")
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert configuration to dictionary."""
        return {
            "gemini": asdict(self.gemini),
            "training": asdict(self.training),
            "memory": asdict(self.memory),
            "prompt": asdict(self.prompt),
            "performance": asdict(self.performance),
            "security": asdict(self.security),
        }
    
    @classmethod
    def from_dict(cls, config_dict: Dict[str, Any]) -> "ToxoConfig":
        """Create configuration from dictionary."""
        # Extract component configs if they exist
        component_configs = {}
        
        if "gemini" in config_dict:
            component_configs["gemini"] = GeminiConfig(**config_dict["gemini"])
        if "training" in config_dict:
            component_configs["training"] = TrainingConfig(**config_dict["training"])
        if "memory" in config_dict:
            component_configs["memory"] = MemoryConfig(**config_dict["memory"])
        if "prompt" in config_dict:
            component_configs["prompt"] = PromptConfig(**config_dict["prompt"])
        if "performance" in config_dict:
            component_configs["performance"] = PerformanceConfig(**config_dict["performance"])
        if "security" in config_dict:
            component_configs["security"] = SecurityConfig(**config_dict["security"])
        
        return cls(**component_configs)
    
    @classmethod
    def from_file(cls, config_path: Union[str, Path]) -> "ToxoConfig":
        """Load configuration from file."""
        return cls(config_path=config_path)
    
    @classmethod
    def load_default(cls) -> "ToxoConfig":
        """Load default configuration."""
        return cls()
    
    def __repr__(self) -> str:
        """String representation of configuration."""
        return f"ToxoConfig(gemini={self.gemini.model}, training_lr={self.training.learning_rate})"


def get_default_config(test_mode: bool = False) -> ToxoConfig:
    """
    Get default Toxo configuration with dynamic API key handling.
    
    Args:
        test_mode: Enable test mode for demos and testing
        
    Returns:
        ToxoConfig: Default configuration
    """
    # Import here to avoid circular imports
    try:
        from .api_manager import api_manager
        
        if test_mode:
            # For test mode, use a test configuration
            gemini_config = GeminiConfig()
            gemini_config.test_mode = True
            gemini_config.api_key = "test_key_for_demo"
        else:
            # Use API manager for dynamic key handling
            return api_manager.create_config()
    except ImportError:
        # Fallback if api_manager is not available
        gemini_config = GeminiConfig()
        gemini_config.test_mode = test_mode
        if test_mode:
            gemini_config.api_key = "test_key_for_demo"
    
    return ToxoConfig(
        gemini=gemini_config,
        training=TrainingConfig(),
        memory=MemoryConfig(),
        prompt=PromptConfig(),
        performance=PerformanceConfig(),
        security=SecurityConfig()
    )


def load_config(config_path: Union[str, Path]) -> ToxoConfig:
    """
    Load configuration from file.
    
    Args:
        config_path: Path to configuration file
        
    Returns:
        ToxoConfig instance
    """
    return ToxoConfig.from_file(config_path) 