"""
AWS Infrastructure for Amar Mascotas

Deploys:
- VPC with public/private subnets
- EC2 instance for Django app + Celery
- RDS PostgreSQL (PostGIS via extension)
- ElastiCache Redis
- Security Groups
- (Optional) ALB, Route53

Estimated cost: ~$93/month

NOTE: AWS is more complex but offers more services and better scaling options.
"""

import pulumi
import pulumi_aws as aws
import sys
sys.path.append("..")
from shared.config import get_config, APP_SERVER_INIT_SCRIPT

# Load configuration
cfg = get_config()

# Get current region and availability zones
region = aws.get_region()
azs = aws.get_availability_zones(state="available")
az1 = azs.names[0]
az2 = azs.names[1] if len(azs.names) > 1 else azs.names[0]

# =============================================================================
# NETWORKING - VPC
# =============================================================================

vpc = aws.ec2.Vpc(
    f"{cfg.resource_prefix}-vpc",
    cidr_block="10.0.0.0/16",
    enable_dns_hostnames=True,
    enable_dns_support=True,
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-vpc"},
)

# Internet Gateway (for public internet access)
igw = aws.ec2.InternetGateway(
    f"{cfg.resource_prefix}-igw",
    vpc_id=vpc.id,
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-igw"},
)

# Public subnets (for EC2, load balancer)
public_subnet_1 = aws.ec2.Subnet(
    f"{cfg.resource_prefix}-public-1",
    vpc_id=vpc.id,
    cidr_block="10.0.1.0/24",
    availability_zone=az1,
    map_public_ip_on_launch=True,
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-public-1"},
)

public_subnet_2 = aws.ec2.Subnet(
    f"{cfg.resource_prefix}-public-2",
    vpc_id=vpc.id,
    cidr_block="10.0.2.0/24",
    availability_zone=az2,
    map_public_ip_on_launch=True,
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-public-2"},
)

# Private subnets (for RDS, ElastiCache)
private_subnet_1 = aws.ec2.Subnet(
    f"{cfg.resource_prefix}-private-1",
    vpc_id=vpc.id,
    cidr_block="10.0.10.0/24",
    availability_zone=az1,
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-private-1"},
)

private_subnet_2 = aws.ec2.Subnet(
    f"{cfg.resource_prefix}-private-2",
    vpc_id=vpc.id,
    cidr_block="10.0.11.0/24",
    availability_zone=az2,
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-private-2"},
)

# Route table for public subnets
public_rt = aws.ec2.RouteTable(
    f"{cfg.resource_prefix}-public-rt",
    vpc_id=vpc.id,
    routes=[
        aws.ec2.RouteTableRouteArgs(
            cidr_block="0.0.0.0/0",
            gateway_id=igw.id,
        ),
    ],
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-public-rt"},
)

# Associate route table with public subnets
aws.ec2.RouteTableAssociation(
    f"{cfg.resource_prefix}-public-1-rta",
    subnet_id=public_subnet_1.id,
    route_table_id=public_rt.id,
)

aws.ec2.RouteTableAssociation(
    f"{cfg.resource_prefix}-public-2-rta",
    subnet_id=public_subnet_2.id,
    route_table_id=public_rt.id,
)

# =============================================================================
# SECURITY GROUPS
# =============================================================================

# App server security group
app_sg = aws.ec2.SecurityGroup(
    f"{cfg.resource_prefix}-app-sg",
    vpc_id=vpc.id,
    description="Security group for app server",
    ingress=[
        # SSH
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=22,
            to_port=22,
            cidr_blocks=cfg.allowed_ssh_ips or ["0.0.0.0/0"],
            description="SSH access",
        ),
        # HTTP
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=80,
            to_port=80,
            cidr_blocks=["0.0.0.0/0"],
            description="HTTP",
        ),
        # HTTPS
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=443,
            to_port=443,
            cidr_blocks=["0.0.0.0/0"],
            description="HTTPS",
        ),
    ],
    egress=[
        aws.ec2.SecurityGroupEgressArgs(
            protocol="-1",
            from_port=0,
            to_port=0,
            cidr_blocks=["0.0.0.0/0"],
            description="Allow all outbound",
        ),
    ],
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-app-sg"},
)

# Database security group (only accessible from app server)
db_sg = aws.ec2.SecurityGroup(
    f"{cfg.resource_prefix}-db-sg",
    vpc_id=vpc.id,
    description="Security group for RDS",
    ingress=[
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=5432,
            to_port=5432,
            security_groups=[app_sg.id],
            description="PostgreSQL from app",
        ),
    ],
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-db-sg"},
)

# Redis security group (only accessible from app server)
redis_sg = aws.ec2.SecurityGroup(
    f"{cfg.resource_prefix}-redis-sg",
    vpc_id=vpc.id,
    description="Security group for ElastiCache",
    ingress=[
        aws.ec2.SecurityGroupIngressArgs(
            protocol="tcp",
            from_port=6379,
            to_port=6379,
            security_groups=[app_sg.id],
            description="Redis from app",
        ),
    ],
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-redis-sg"},
)

# =============================================================================
# DATABASE - RDS PostgreSQL
# =============================================================================

# Subnet group for RDS (requires at least 2 AZs)
db_subnet_group = aws.rds.SubnetGroup(
    f"{cfg.resource_prefix}-db-subnet-group",
    subnet_ids=[private_subnet_1.id, private_subnet_2.id],
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-db-subnet-group"},
)

# RDS PostgreSQL instance
# Note: PostGIS is available as an extension, enable after creation
db_instance = aws.rds.Instance(
    f"{cfg.resource_prefix}-db",
    identifier=f"{cfg.resource_prefix}-db",
    engine="postgres",
    engine_version=cfg.db_version,
    instance_class="db.t3.micro",      # $25/mo - smallest
    allocated_storage=20,
    storage_type="gp3",
    db_name=cfg.db_name,
    username=cfg.db_user,
    password=pulumi.Config().require_secret("db_password"),  # Set via: pulumi config set --secret db_password xxx
    vpc_security_group_ids=[db_sg.id],
    db_subnet_group_name=db_subnet_group.name,
    publicly_accessible=False,
    skip_final_snapshot=True,           # Set False for production!
    backup_retention_period=7,
    multi_az=False,                     # Set True for HA ($$$)
    tags=cfg.tags,
)

# =============================================================================
# CACHE - ElastiCache Redis
# =============================================================================

# Subnet group for ElastiCache
redis_subnet_group = aws.elasticache.SubnetGroup(
    f"{cfg.resource_prefix}-redis-subnet-group",
    subnet_ids=[private_subnet_1.id, private_subnet_2.id],
    tags=cfg.tags,
)

# ElastiCache Redis cluster
redis_cluster = aws.elasticache.Cluster(
    f"{cfg.resource_prefix}-redis",
    cluster_id=f"{cfg.resource_prefix}-redis",
    engine="redis",
    engine_version="7.0",
    node_type="cache.t3.micro",         # $15/mo - smallest
    num_cache_nodes=1,
    port=6379,
    subnet_group_name=redis_subnet_group.name,
    security_group_ids=[redis_sg.id],
    tags=cfg.tags,
)

# =============================================================================
# COMPUTE - EC2 Instance
# =============================================================================

# Get latest Ubuntu 22.04 AMI
ubuntu_ami = aws.ec2.get_ami(
    most_recent=True,
    owners=["099720109477"],  # Canonical
    filters=[
        aws.ec2.GetAmiFilterArgs(
            name="name",
            values=["ubuntu/images/hvm-ssd/ubuntu-jammy-22.04-amd64-server-*"],
        ),
        aws.ec2.GetAmiFilterArgs(
            name="virtualization-type",
            values=["hvm"],
        ),
    ],
)

# Key pair (import your existing key or create new)
# key_pair = aws.ec2.KeyPair(
#     f"{cfg.resource_prefix}-key",
#     public_key=open("~/.ssh/id_rsa.pub").read(),
#     tags=cfg.tags,
# )

# EC2 instance
ec2_instance = aws.ec2.Instance(
    f"{cfg.resource_prefix}-app",
    ami=ubuntu_ami.id,
    instance_type="t3.medium",          # $35/mo - 4GB RAM, 2 vCPU
    subnet_id=public_subnet_1.id,
    vpc_security_group_ids=[app_sg.id],
    # key_name=key_pair.key_name,       # Uncomment when key_pair is defined
    user_data=APP_SERVER_INIT_SCRIPT,
    root_block_device=aws.ec2.InstanceRootBlockDeviceArgs(
        volume_size=30,
        volume_type="gp3",
    ),
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-app"},
)

# Elastic IP (static public IP)
eip = aws.ec2.Eip(
    f"{cfg.resource_prefix}-eip",
    instance=ec2_instance.id,
    domain="vpc",
    tags={**cfg.tags, "Name": f"{cfg.resource_prefix}-eip"},
)

# =============================================================================
# OPTIONAL: Application Load Balancer (uncomment if needed)
# =============================================================================

# alb = aws.lb.LoadBalancer(
#     f"{cfg.resource_prefix}-alb",
#     load_balancer_type="application",
#     security_groups=[app_sg.id],
#     subnets=[public_subnet_1.id, public_subnet_2.id],
#     tags=cfg.tags,
# )

# =============================================================================
# OUTPUTS
# =============================================================================

pulumi.export("ec2_public_ip", eip.public_ip)
pulumi.export("ec2_private_ip", ec2_instance.private_ip)
pulumi.export("db_endpoint", db_instance.endpoint)
pulumi.export("db_name", cfg.db_name)
pulumi.export("db_user", cfg.db_user)
pulumi.export("redis_endpoint", redis_cluster.cache_nodes[0].address)
pulumi.export("redis_port", redis_cluster.port)

# Generate .env content
pulumi.export("env_file", pulumi.Output.all(
    db_instance.endpoint,
    redis_cluster.cache_nodes[0].address,
    redis_cluster.port,
).apply(lambda args: f"""
# Generated by Pulumi - AWS
DB_HOST={args[0].split(':')[0]}
DB_PORT=5432
DB_NAME={cfg.db_name}
DB_USER={cfg.db_user}
DB_PASSWORD=<set via pulumi config>
CELERY_BROKER_URL=redis://{args[1]}:{args[2]}/0
CELERY_RESULT_BACKEND=redis://{args[1]}:{args[2]}/0
"""))
