"""
Lambda Infrastructure Configuration

This module provides the configuration classes for AWS Lambda functions.
The config classes are simple data containers that optionally provide
a get() method to return runtime APIs.
"""

import boto3
import inspect
import os
from typing import Optional, Dict, Any, List, Union, TYPE_CHECKING, Callable, TypeVar
from pathlib import Path
from pydantic import BaseModel, Field
from ...models.base import InfraConfig

if TYPE_CHECKING:
    from mypy_boto3_lambda import LambdaClient


class IAMRoleConfig(BaseModel):
    """Configuration for IAM role used by Lambda function."""
    logical_id: str = Field(..., description="CloudFormation logical ID for the role")
    managed_policies: List[str] = Field(default_factory=list, description="AWS managed policy ARNs")
    inline_policies: Dict[str, Dict[str, Any]] = Field(default_factory=dict, description="Inline policy documents")


class LambdaInfraConfig(InfraConfig):
    """
    Infrastructure configuration for an AWS Lambda function.
    
    This class serves as the primary entry point for accessing Lambda functions.
    It handles infrastructure configuration and provides direct access to
    boto3 Lambda client.
    """
    logical_id: str = Field(..., description="CloudFormation logical ID")
    function_name: str = Field(..., description="Lambda function name")
    code_directory: Path = Field(..., description="Directory containing Lambda function code")
    handler: str = Field(..., description="Function handler (e.g., 'index.handler')")
    runtime: str = Field(default="python3.11", description="Python runtime version")
    role: Optional[IAMRoleConfig] = Field(default=None, description="IAM execution role configuration")
    timeout: int = Field(default=3, ge=1, le=900, description="Function timeout in seconds")
    memory_size: int = Field(default=128, ge=128, le=10240, description="Memory size in MB")
    environment: Dict[str, str] = Field(default_factory=dict, description="Environment variables")
    tags: Dict[str, str] = Field(default_factory=dict, description="Resource tags")
    sqs_triggers: List[str] = Field(default_factory=list, description="SQS queue ARNs to trigger this function")
    sns_triggers: List[str] = Field(default_factory=list, description="SNS topic ARNs to trigger this function (internal use)")
    dependencies: Optional[Union[List[str], Path]] = Field(default=None, description="Dependencies as list of strings or path to requirements.txt file")

    def __init__(self, **data):
        super().__init__(**data)
        
        # Create default IAM role if not provided
        if self.role is None:
            self.role = IAMRoleConfig(
                logical_id=f"{self.logical_id}Role",
                managed_policies=["arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"]
            )
            
            # Add SQS permissions if SQS triggers are present
            if self.sqs_triggers:
                # Resolve queue names to CloudFormation references
                resolved_resources = []
                for queue_identifier in self.sqs_triggers:
                    if queue_identifier.startswith('arn:aws:sqs:'):
                        # Already an ARN
                        resolved_resources.append(queue_identifier)
                    else:
                        # Queue name - resolve to CloudFormation GetAtt reference
                        sanitized_name = queue_identifier.replace('-', '').replace('_', '')
                        resource_name = f"{sanitized_name}Queue"
                        resolved_resources.append({"Fn::GetAtt": [resource_name, "Arn"]})
                
                self.role.inline_policies["SQSPollingPolicy"] = {
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Effect": "Allow",
                            "Action": [
                                "sqs:ReceiveMessage",
                                "sqs:DeleteMessage",
                                "sqs:GetQueueAttributes"
                            ],
                            "Resource": resolved_resources
                        }
                    ]
                }

    def to_cloudformation_properties(self) -> Dict[str, Any]:
        """Convert to CloudFormation Lambda function properties."""
        properties = {
            "FunctionName": self.function_name,
            "Runtime": self.runtime,
            "Handler": self.handler,
            "Role": {"Fn::GetAtt": [self.role.logical_id, "Arn"]},
            "Timeout": self.timeout,
            "MemorySize": self.memory_size,
        }
        
        # Code will be handled by the packaging pipeline
        # Use the bucket name that will be created by CLI pre-deployment step
        properties["Code"] = {
            "S3Bucket": f"{self.logical_id.lower()}-deployment-bucket",
            "S3Key": f"lambda/{self.function_name}.zip"
        }
        
        # Add environment variables
        if self.environment:
            properties["Environment"] = {"Variables": self.environment}
        
        # Note: Lambda layers are not supported in the simplified approach
        # Dependencies are packaged directly into the ZIP file using pip install -t
        
        # Add tags
        if self.tags:
            properties["Tags"] = [{"Key": k, "Value": v} for k, v in self.tags.items()]
        
        return properties

    def get(self, session: Optional[boto3.Session] = None) -> 'LambdaClient':
        """
        Returns the boto3 Lambda client for this function configuration.
        
        Args:
            session: Optional boto3 session to use. If None, uses default session.
            
        Returns:
            boto3 Lambda client
        """
        if session is None:
            session = boto3.Session()
        
        return session.client('lambda')

    def __call__(self, *args, **kwargs):
        """
        Make the LambdaInfraConfig callable so it can execute the original function.
        This allows the decorated function to work normally while still providing access to the config.
        """
        if hasattr(self, '_original_function'):
            return self._original_function(*args, **kwargs)
        else:
            raise RuntimeError("No original function found in LambdaInfraConfig")


# Type variable for the decorated function
F = TypeVar('F', bound=Callable[..., Any])

def handler(
    function_name: Optional[str] = None,
    logical_id: Optional[str] = None,
    runtime: str = "python3.11",
    timeout: int = 3,
    memory_size: int = 128,
    environment: Optional[Dict[str, str]] = None,
    tags: Optional[Dict[str, str]] = None,
    dependencies: Optional[Union[List[str], Path]] = None,
    role: Optional[IAMRoleConfig] = None,
    stack: Optional[str] = None,
    # Event source parameters
    input: Optional[Union['SQSQueueConfig', str, List[Union['SQSQueueConfig', str]]]] = None,
) -> Callable[[F], F]:
    """
    Decorator to convert a Python function into a Lambda function configuration.
    
    Usage with named event sources:
        # SQS input queue automatically creates SQS trigger
        input_queue = SQSQueueConfig(name="input-messages")
        
        @handler(input=input_queue)
        def process_message(event, context):
            # Automatically triggered by input_queue
            return {"statusCode": 200, "body": "processed"}
        
        # Multiple inputs (SQS + SNS)
        alerts_queue = SQSQueueConfig(name="alerts")
        notifications_topic = SNSTopicInfraConfig(name="notifications")
        
        @handler(input=[alerts_queue, notifications_topic])
        def handle_multiple_sources(event, context):
            # Triggered by both SQS queue and SNS topic
            return {"statusCode": 200}
    
    Args:
        function_name: Name of the Lambda function (defaults to Python function name)
        logical_id: CloudFormation logical ID (defaults to function_name)
        runtime: Python runtime version
        timeout: Function timeout in seconds
        memory_size: Memory size in MB
        environment: Environment variables
        tags: Resource tags
        dependencies: Dependencies as list or path to requirements.txt
        role: Custom IAM role configuration
        stack: CloudFormation stack name (defaults to "IICYStack")
        input: SQS/SNS resources for input (creates event source mappings/triggers)
    """
    def decorator(func: Callable) -> 'LambdaInfraConfig':
        # Get the source file path
        source_file = Path(inspect.getfile(func))
        source_dir = source_file.parent
        
        # Import config classes locally to avoid circular imports
        from ...queue import SQSQueueConfig
        from ...sns import SNSTopicInfraConfig
        
        # Process named event source parameters
        detected_sqs_triggers = []
        detected_sns_triggers = []
        
        def process_input_source(source):
            """Process a single input source (SQS, SNS, or string)."""
            if isinstance(source, SQSQueueConfig):
                detected_sqs_triggers.append(source.name)
            elif isinstance(source, SNSTopicInfraConfig):
                # For SNS topics, we need to use the ARN or topic name
                # If topic_name exists, construct ARN, otherwise use name
                if hasattr(source, 'topic_name') and source.topic_name:
                    # This will be resolved to ARN during CloudFormation generation
                    detected_sns_triggers.append(source.topic_name)
                else:
                    detected_sns_triggers.append(source.name)
            elif isinstance(source, str):
                # String could be queue name or SNS ARN
                if source.startswith('arn:aws:sns:'):
                    detected_sns_triggers.append(source)
                else:
                    # Assume it's a queue name
                    detected_sqs_triggers.append(source)
        
        # Process input parameter (creates triggers for SQS/SNS)
        if input is not None:
            if isinstance(input, list):
                for source in input:
                    process_input_source(source)
            else:
                process_input_source(input)
        
        # Use function name as default if none provided
        detected_function_name = function_name if function_name is not None else func.__name__
        
        # Create logical_id if not provided
        config_logical_id = logical_id
        if config_logical_id is None:
            config_logical_id = detected_function_name.replace('-', '').replace('_', '')
        
        # Get configuration defaults
        from ...config import get_config
        
        config_obj = get_config()
        
        # Create LambdaInfraConfig
        config = LambdaInfraConfig(
            logical_id=config_logical_id,
            function_name=detected_function_name,
            code_directory=source_dir,
            handler=f"{source_file.stem}.{func.__name__}",
            runtime=runtime or config_obj.lambda_.runtime,
            timeout=timeout or config_obj.lambda_.timeout,
            memory_size=memory_size or config_obj.lambda_.memory,
            environment=environment or {},
            tags=tags or {},
            dependencies=dependencies,
            sqs_triggers=detected_sqs_triggers,
            sns_triggers=detected_sns_triggers,
            role=role,
            stack=stack or "IICYStack"
        )
        
        # Store the original function in the config
        config._original_function = func
        
        # Return the config (which is now callable)
        return config
    
    return decorator
