from typing import List, Union, Any, Dict, Optional
from pydantic import BaseModel, Field


class ResourceProperties(BaseModel):
    """Base class for CloudFormation resource properties"""
    pass


# AWS Lambda Properties
class LambdaFunctionProperties(ResourceProperties):
    """Properties for AWS::Lambda::Function"""
    FunctionName: str
    Runtime: str
    Handler: str
    Code: Dict[str, Any]
    Role: Dict[str, Any]  # Ref to IAM role or GetAtt
    Timeout: Optional[int] = 3
    MemorySize: Optional[int] = 128
    Environment: Optional[Dict[str, Any]] = None
    Tags: Optional[List[Dict[str, str]]] = Field(default_factory=list)


class LambdaPermissionProperties(ResourceProperties):
    """Properties for AWS::Lambda::Permission"""
    FunctionName: Dict[str, str]  # Ref to function
    Action: str
    Principal: str
    SourceArn: Optional[str] = None


class LambdaEventSourceMappingProperties(ResourceProperties):
    """Properties for AWS::Lambda::EventSourceMapping"""
    EventSourceArn: Union[str, Dict[str, Any]]  # String ARN or CloudFormation function
    FunctionName: Dict[str, str]  # Ref to function
    BatchSize: int = 10


# AWS IAM Properties
class IAMRoleProperties(ResourceProperties):
    """Properties for AWS::IAM::Role"""
    AssumeRolePolicyDocument: Dict[str, Any]
    ManagedPolicyArns: Optional[List[str]] = Field(default_factory=list)
    Policies: Optional[List[Dict[str, Any]]] = Field(default_factory=list)


# AWS SNS Properties
class SNSTopicProperties(ResourceProperties):
    """Properties for AWS::SNS::Topic"""
    TopicName: Optional[str] = None
    DisplayName: Optional[str] = None
    FifoTopic: Optional[bool] = None
    ContentBasedDeduplication: Optional[bool] = None
    KmsMasterKeyId: Optional[str] = None
    Tags: Optional[List[Dict[str, Any]]] = None


class S3BucketProperties(ResourceProperties):
    """Properties for AWS::S3::Bucket"""
    BucketName: str
    PublicAccessBlockConfiguration: Optional[Dict[str, Any]] = None
    BucketEncryption: Optional[Dict[str, Any]] = None


class SNSSubscriptionProperties(ResourceProperties):
    """Properties for AWS::SNS::Subscription"""
    Protocol: str
    TopicArn: str
    Endpoint: Dict[str, Any]  # GetAtt to Lambda function


# AWS SQS Properties
class SQSQueueProperties(ResourceProperties):
    """Properties for AWS::SQS::Queue"""
    QueueName: str
    FifoQueue: Optional[bool] = None
    ContentBasedDeduplication: Optional[bool] = None


class CloudFormationIntrinsicFunction(BaseModel):
    """Base class for CloudFormation intrinsic functions"""
    pass


class Ref(CloudFormationIntrinsicFunction):
    """CloudFormation Ref intrinsic function"""
    Ref: str


class GetAtt(CloudFormationIntrinsicFunction):
    """CloudFormation Fn::GetAtt intrinsic function"""
    class Config:
        fields = {"Fn::GetAtt": "get_att"}
    
    get_att: List[str]


class OutputValue(BaseModel):
    """CloudFormation output value - can be a Ref or GetAtt"""
    ref: Optional[str] = None
    get_att: Optional[List[str]] = None
    
    def model_dump(self, **kwargs):
        if self.ref:
            return {"Ref": self.ref}
        elif self.get_att:
            return {"Fn::GetAtt": self.get_att}
        else:
            return {}


class OutputExport(BaseModel):
    """CloudFormation output export"""
    name: str


class Output(BaseModel):
    """CloudFormation output"""
    description: str
    value: OutputValue
    export: OutputExport


class Resource(BaseModel):
    """CloudFormation resource"""
    logical_id: str
    type: str
    properties: ResourceProperties
    
    def model_dump(self, **kwargs):
        """Custom serialization to properly handle properties field"""
        data = super().model_dump(**kwargs)
        # Ensure properties are serialized correctly and exclude None values
        if hasattr(self.properties, 'model_dump'):
            props_data = self.properties.model_dump(**kwargs)
            # Filter out None values and empty collections to match CloudFormation expectations
            filtered_props = {}
            for k, v in props_data.items():
                if v is not None:
                    # Skip empty dictionaries and lists
                    if isinstance(v, (dict, list)) and len(v) == 0:
                        continue
                    filtered_props[k] = v
            data['properties'] = filtered_props
        return data


class CloudFormationTemplate(BaseModel):
    """Complete CloudFormation template"""
    aws_template_format_version: str = "2010-09-09"
    description: str
    resources: Dict[str, Resource] = Field(default_factory=dict)
    outputs: Dict[str, Output] = Field(default_factory=dict)
    
    class Config:
        fields = {
            "aws_template_format_version": "AWSTemplateFormatVersion"
        }


class Stack(BaseModel):
    """Stack representation with resources and outputs"""
    name: str
    resources: List[Resource] = Field(default_factory=list)
    outputs: List[Output] = Field(default_factory=list)

    def add_resource(self, resource: Resource) -> 'Stack':
        """Add a resource to the stack"""
        self.resources.append(resource)
        return self

    def add_output(self, output: Output) -> 'Stack':
        """Add an output to the stack"""
        self.outputs.append(output)
        return self

    def add(self, item: Union[Resource, Output]) -> 'Stack':
        """Add either a resource or output to the stack"""
        if isinstance(item, Resource):
            self.resources.append(item)
        elif isinstance(item, Output):
            self.outputs.append(item)
        return self
    
    def to_cloudformation_template(self) -> CloudFormationTemplate:
        """Convert stack to CloudFormation template format"""
        # Convert list-based resources to dict format
        resources_dict = {}
        for resource in self.resources:
            resources_dict[resource.logical_id] = resource
        
        # Convert list-based outputs to dict format  
        outputs_dict = {}
        for i, output in enumerate(self.outputs):
            output_name = f"Output{i}"  # Default naming, should be overridden
            outputs_dict[output_name] = output
            
        return CloudFormationTemplate(
            description=f"CloudFormation template for stack {self.name}",
            resources=resources_dict,
            outputs=outputs_dict
        )
    