from pydantic import BaseModel, Field, model_validator, field_validator, PrivateAttr
from typing import Union, Optional, Dict, List, Any, ClassVar, Tuple
from enum import Enum
from datetime import datetime
from pathlib import Path
import logging

from ...core.base.config_base import BasePipelineConfig

logger = logging.getLogger(__name__)


class VariableType(str, Enum):
    NUMERIC = "NUMERIC"
    TEXT = "TEXT"
    
    @classmethod
    def _missing_(cls, value: str) -> Optional['VariableType']:
        """Handle string values"""
        try:
            return cls(value.upper())
        except ValueError:
            return None

    def __str__(self) -> str:
        """String representation"""
        return self.value


def create_inference_variable_list(
    numeric_fields: List[str] = None,
    text_fields: List[str] = None,
    output_format: str = 'dict'
) -> Union[Dict[str, Union[VariableType, str]], List[List[str]]]:
    """
    Create an inference variable list for model input variables using separate lists for numeric and text fields.
    This is a helper function that can be used standalone or within RegistrationConfig.
    
    Args:
        numeric_fields: List of field names that should be treated as NUMERIC
        text_fields: List of field names that should be treated as TEXT
        output_format: Format for storing variable list - either 'dict' or 'list'
        
    Returns:
        A dictionary mapping variable names to their types, or a list of [name, type] pairs,
        depending on the output_format parameter
    """
    # Initialize with empty lists if not provided
    numeric_fields = numeric_fields or []
    text_fields = text_fields or []
    
    # Validate inputs are lists of strings
    for field_name in numeric_fields:
        if not isinstance(field_name, str):
            raise ValueError(f"Field name must be string, got {type(field_name)} for: {field_name}")
            
    for field_name in text_fields:
        if not isinstance(field_name, str):
            raise ValueError(f"Field name must be string, got {type(field_name)} for: {field_name}")
    
    # Check for duplicates between numeric and text fields
    common_fields = set(numeric_fields) & set(text_fields)
    if common_fields:
        raise ValueError(f"Fields cannot be both numeric and text: {common_fields}")
    
    # Validate output format
    if output_format not in ['dict', 'list']:
        raise ValueError(f"Output format must be 'dict' or 'list', got: {output_format}")
    
    # Create the variable list in the requested format
    if output_format == 'dict':
        # Dictionary format - map field names to their types
        result = {}
        
        # Add numeric fields
        for field_name in numeric_fields:
            result[field_name] = VariableType.NUMERIC
            
        # Add text fields
        for field_name in text_fields:
            result[field_name] = VariableType.TEXT
            
    else:  # 'list' format
        # List format - create list of [name, type] pairs
        result = []
        
        # Add numeric fields
        for field_name in numeric_fields:
            result.append([field_name, VariableType.NUMERIC.value])
            
        # Add text fields
        for field_name in text_fields:
            result.append([field_name, VariableType.TEXT.value])
        
    return result


class RegistrationConfig(BasePipelineConfig):
    """
    Configuration for model registration step, following the three-tier categorization:
    
    Tier 1: Essential User Inputs - fields that users must explicitly provide
    Tier 2: System Inputs - fields with reasonable defaults that users can override
    Tier 3: Derived Fields - private fields with read-only property access
    """
   
    # ===== Essential User Inputs (Tier 1) =====
    # These are fields that users must explicitly provide
    
    model_owner: str = Field(
        description="Team ID of model owner"
    )
    
    model_domain: str = Field(
        description="Domain for model registration"
    )
    
    model_objective: str = Field(
        description="Objective of model registration"
    )
    
    framework: str = Field(
        description="ML framework used for the model"
    )
    
    inference_entry_point: str = Field(
        description="Entry point script for inference"
    )
    
    source_model_inference_input_variable_list: Union[Dict[str, Union[VariableType, str]], List[List[str]]] = Field(
        default_factory=dict,
        description="Input variables and their types. Can be either:\n"
                   "1. Dictionary: {'var1': 'NUMERIC', 'var2': 'TEXT'}\n"
                   "2. List of pairs: [['var1', 'NUMERIC'], ['var2', 'TEXT']]"
    )
    
    # ===== System Inputs with Defaults (Tier 2) =====
    # These are fields with reasonable defaults that users can override
    
    inference_instance_type: str = Field(
        default='ml.m5.large', 
        description="Instance type for inference endpoint/transform job"
    )
    
    source_model_inference_content_types: List[str] = Field(
        default=["text/csv"],
        description="Content type for model inference input. Must be exactly ['text/csv'] or ['application/json']"
    )
    
    source_model_inference_response_types: List[str] = Field(
        default=["application/json"],
        description="Response type for model inference output. Must be exactly ['text/csv'] or ['application/json']"
    )

    source_model_inference_output_variable_list: Dict[str, VariableType] = Field(
        default={
            'legacy-score': VariableType.NUMERIC
        },
        description="Dictionary of output variables and their types (NUMERIC or TEXT)"
    )
    
    # ===== Derived Fields (Tier 3) =====
    # These are fields calculated from other fields, stored in private attributes
    # with public read-only properties for access
    
    # Removed _source_model_inference_input_variable_list as it's now a field
    _variable_schema: Optional[Dict[str, Dict[str, List[Dict[str, str]]]]] = PrivateAttr(default=None)
    
    # Update to Pydantic V2 style model_config
    model_config = {
        'arbitrary_types_allowed': True,
        'validate_assignment': True,
        'extra': 'allow',  # Accept metadata fields during deserialization
        'json_encoders': {
            VariableType: lambda v: v.value
        }
    }
    
    # ===== Property Accessors for Derived Fields =====
    # (No property accessor needed for source_model_inference_input_variable_list since it's now a field)
    
    # ===== Validators =====
    
    @field_validator('inference_instance_type')
    @classmethod
    def validate_inference_instance_type(cls, v: str) -> str:
        """Validate the inference instance type"""
        if not v.startswith('ml.'):
            raise ValueError(f"Invalid inference instance type: {v}. Must start with 'ml.'")
        return v

    @field_validator('framework')
    @classmethod
    def validate_framework(cls, v: str) -> str:
        """Validate the ML framework"""
        valid_frameworks = ['xgboost', 'sklearn', 'pytorch', 'tensorflow']
        if v.lower() not in valid_frameworks:
            raise ValueError(f"Framework must be one of {valid_frameworks}")
        return v.lower()
    
    # ===== Model Validation =====
    
    @model_validator(mode='after')
    def initialize_derived_fields(self) -> 'RegistrationConfig':
        """Initialize all derived fields once after validation."""
        # Call parent validator first
        super().initialize_derived_fields()
        return self
        
    @model_validator(mode='after')
    def validate_registration_configs(self) -> 'RegistrationConfig':
        """Validate registration-specific configurations"""
        # Validate inference entry point
        if self.source_dir and not self.source_dir.startswith('s3://'):
            entry_point_path = Path(self.source_dir) / self.inference_entry_point
            if not entry_point_path.exists():
                raise ValueError(f"Inference entry point script not found: {entry_point_path}")
        
        return self

    @field_validator('source_model_inference_content_types', 'source_model_inference_response_types')
    @classmethod
    def validate_content_types(cls, v: List[str]) -> List[str]:
        """Validate content and response types"""
        valid_types = [["text/csv"], ["application/json"]]
        if v not in valid_types:
            raise ValueError(f"Content/Response types must be one of {valid_types}")
        return v
    
    @field_validator('source_model_inference_input_variable_list')
    @classmethod
    def validate_input_variable_list(
        cls, 
        v: Union[Dict[str, Union[VariableType, str]], List[List[str]]]
    ) -> Union[Dict[str, Union[VariableType, str]], List[List[str]]]:
        """
        Validate input variable list format.
        
        Args:
            v: Either a dictionary of variable names to types,
               or a list of [variable_name, variable_type] pairs
               
        Returns:
            Validated input variable list
        """
        if v is None:
            return {}  # Return empty dict as default
        
        # Handle dictionary format
        if isinstance(v, dict):
            result = {}
            for key, value in v.items():
                if not isinstance(key, str):
                    raise ValueError(f"Key must be string, got {type(key)} for key: {key}")
                
                # Convert string values to VariableType enum
                if isinstance(value, str):
                    try:
                        value = VariableType(value.upper())
                    except ValueError:
                        raise ValueError(f"Value must be 'NUMERIC' or 'TEXT', got: {value}")
                elif isinstance(value, VariableType):
                    # Keep VariableType as is
                    pass
                else:
                    raise ValueError(f"Value must be string or VariableType, got: {type(value)}")
                
                result[key] = value
            return result

        # Handle list format
        elif isinstance(v, list):
            for item in v:
                if not isinstance(item, list) or len(item) != 2:
                    raise ValueError("Each item must be a list of [variable_name, variable_type]")
                
                var_name, var_type = item
                if not isinstance(var_name, str):
                    raise ValueError(f"Variable name must be string, got {type(var_name)}")
                
                if not isinstance(var_type, str):
                    raise ValueError(f"Type must be string, got {type(var_type)}")
                
                if var_type.upper() not in ["NUMERIC", "TEXT"]:
                    raise ValueError(f"Type must be 'NUMERIC' or 'TEXT', got: {var_type}")
            
            return v

        else:
            raise ValueError("Must be either a dictionary or a list of pairs")
            
    @field_validator('source_model_inference_output_variable_list')
    @classmethod
    def validate_output_variable_list(cls, v: Dict[str, Union[VariableType, str]]) -> Dict[str, str]:
        """Validate variable lists and convert to string values"""
        if not v:  # If empty dictionary
            return v
        
        result = {}
        for key, value in v.items():
            # Validate key is a string
            if not isinstance(key, str):
                raise ValueError(f"Key must be string, got {type(key)} for key: {key}")
        
            # Convert VariableType to string or validate string value
            if isinstance(value, VariableType):
                result[key] = value.value
            elif isinstance(value, str) and value in [vt.value for vt in VariableType]:
                result[key] = value
            else:
                raise ValueError(f"Value must be either 'NUMERIC' or 'TEXT', got: {value}")
    
        return result
    
    # ===== Property Accessors for Derived Method Results =====
        
    @property
    def variable_schema(self) -> Dict[str, Dict[str, List[Dict[str, str]]]]:
        """Generate variable schema for model registration"""
        if self._variable_schema is None:
            schema = {
                "input": {"variables": []},
                "output": {"variables": []}
            }
            
            # Handle input variables in either format
            input_vars = self.source_model_inference_input_variable_list
            if isinstance(input_vars, dict):
                # Dictionary format
                for var_name, var_type in input_vars.items():
                    schema["input"]["variables"].append({
                        "name": var_name,
                        "type": var_type if isinstance(var_type, str) else var_type.value
                    })
            elif isinstance(input_vars, list):
                # List format
                for var_name, var_type in input_vars:
                    schema["input"]["variables"].append({
                        "name": var_name,
                        "type": var_type
                    })
                
            # Add output variables
            for name, var_type in self.source_model_inference_output_variable_list.items():
                schema["output"]["variables"].append({
                    "name": name,
                    "type": var_type if isinstance(var_type, str) else var_type.value
                })
                
            self._variable_schema = schema
        
        return self._variable_schema
    
    # ===== Legacy Methods =====
        
    def get_variable_schema(self) -> Dict[str, Dict[str, List[Dict[str, str]]]]:
        """Legacy method that forwards to the property"""
        return self.variable_schema
    
    # ===== Serialization =====
    
    def model_dump(self, **kwargs) -> Dict[str, Any]:
        """Override model_dump to include derived properties."""
        data = super().model_dump(**kwargs)
        
        # Add derived properties to output
        data["variable_schema"] = self.variable_schema
        
        # Process variable lists for proper serialization
        if 'source_model_inference_output_variable_list' in data:
            data['source_model_inference_output_variable_list'] = {
                k: v.value if isinstance(v, VariableType) else v
                for k, v in data['source_model_inference_output_variable_list'].items()
            }
        
        if 'source_model_inference_input_variable_list' in data:
            input_vars = data['source_model_inference_input_variable_list']
            if isinstance(input_vars, dict):
                data['source_model_inference_input_variable_list'] = {
                    k: v.value if isinstance(v, VariableType) else v
                    for k, v in input_vars.items()
                }
            
        return data
    
    # ===== Methods for working with input variable lists =====
    
    def set_source_model_inference_input_variable_list(
        self,
        numeric_fields: List[str] = None,
        text_fields: List[str] = None,
        output_format: str = 'dict'
    ) -> None:
        """
        Set the input variable list for model inference using separate lists for numeric and text fields.
        
        Args:
            numeric_fields: List of field names that should be treated as NUMERIC
            text_fields: List of field names that should be treated as TEXT
            output_format: Format for storing variable list - either 'dict' or 'list'
        """
        # Use the standalone function to create the variable list
        result = create_inference_variable_list(
            numeric_fields=numeric_fields,
            text_fields=text_fields,
            output_format=output_format
        )
        
        # Set the input variable list (now a regular field, not a private attribute)
        self.source_model_inference_input_variable_list = result
            
        # Invalidate cached schema so it will be regenerated
        self._variable_schema = None
