# LocalStack Resource Provider Scaffolding v2
from __future__ import annotations

import base64
from pathlib import Path
from typing import TypedDict

import localstack.services.cloudformation.provider_utils as util
from localstack.services.cloudformation.resource_provider import (
    OperationStatus,
    ProgressEvent,
    ResourceProvider,
    ResourceRequest,
)
from localstack.utils.strings import to_str


class EC2InstanceProperties(TypedDict):
    AdditionalInfo: str | None
    Affinity: str | None
    AvailabilityZone: str | None
    BlockDeviceMappings: list[BlockDeviceMapping] | None
    CpuOptions: CpuOptions | None
    CreditSpecification: CreditSpecification | None
    DisableApiTermination: bool | None
    EbsOptimized: bool | None
    ElasticGpuSpecifications: list[ElasticGpuSpecification] | None
    ElasticInferenceAccelerators: list[ElasticInferenceAccelerator] | None
    EnclaveOptions: EnclaveOptions | None
    HibernationOptions: HibernationOptions | None
    HostId: str | None
    HostResourceGroupArn: str | None
    IamInstanceProfile: str | None
    Id: str | None
    ImageId: str | None
    InstanceInitiatedShutdownBehavior: str | None
    InstanceType: str | None
    Ipv6AddressCount: int | None
    Ipv6Addresses: list[InstanceIpv6Address] | None
    KernelId: str | None
    KeyName: str | None
    LaunchTemplate: LaunchTemplateSpecification | None
    LicenseSpecifications: list[LicenseSpecification] | None
    Monitoring: bool | None
    NetworkInterfaces: list[NetworkInterface] | None
    PlacementGroupName: str | None
    PrivateDnsName: str | None
    PrivateDnsNameOptions: PrivateDnsNameOptions | None
    PrivateIp: str | None
    PrivateIpAddress: str | None
    PropagateTagsToVolumeOnCreation: bool | None
    PublicDnsName: str | None
    PublicIp: str | None
    RamdiskId: str | None
    SecurityGroupIds: list[str] | None
    SecurityGroups: list[str] | None
    SourceDestCheck: bool | None
    SsmAssociations: list[SsmAssociation] | None
    SubnetId: str | None
    Tags: list[Tag] | None
    Tenancy: str | None
    UserData: str | None
    Volumes: list[Volume] | None


class Ebs(TypedDict):
    DeleteOnTermination: bool | None
    Encrypted: bool | None
    Iops: int | None
    KmsKeyId: str | None
    SnapshotId: str | None
    VolumeSize: int | None
    VolumeType: str | None


class BlockDeviceMapping(TypedDict):
    DeviceName: str | None
    Ebs: Ebs | None
    NoDevice: dict | None
    VirtualName: str | None


class InstanceIpv6Address(TypedDict):
    Ipv6Address: str | None


class ElasticGpuSpecification(TypedDict):
    Type: str | None


class ElasticInferenceAccelerator(TypedDict):
    Type: str | None
    Count: int | None


class Volume(TypedDict):
    Device: str | None
    VolumeId: str | None


class LaunchTemplateSpecification(TypedDict):
    Version: str | None
    LaunchTemplateId: str | None
    LaunchTemplateName: str | None


class EnclaveOptions(TypedDict):
    Enabled: bool | None


class PrivateIpAddressSpecification(TypedDict):
    Primary: bool | None
    PrivateIpAddress: str | None


class NetworkInterface(TypedDict):
    DeviceIndex: str | None
    AssociateCarrierIpAddress: bool | None
    AssociatePublicIpAddress: bool | None
    DeleteOnTermination: bool | None
    Description: str | None
    GroupSet: list[str] | None
    Ipv6AddressCount: int | None
    Ipv6Addresses: list[InstanceIpv6Address] | None
    NetworkInterfaceId: str | None
    PrivateIpAddress: str | None
    PrivateIpAddresses: list[PrivateIpAddressSpecification] | None
    SecondaryPrivateIpAddressCount: int | None
    SubnetId: str | None


class Tag(TypedDict):
    Key: str | None
    Value: str | None


class HibernationOptions(TypedDict):
    Configured: bool | None


class LicenseSpecification(TypedDict):
    LicenseConfigurationArn: str | None


class CpuOptions(TypedDict):
    CoreCount: int | None
    ThreadsPerCore: int | None


class PrivateDnsNameOptions(TypedDict):
    EnableResourceNameDnsAAAARecord: bool | None
    EnableResourceNameDnsARecord: bool | None
    HostnameType: str | None


class AssociationParameter(TypedDict):
    Key: str | None
    Value: list[str] | None


class SsmAssociation(TypedDict):
    DocumentName: str | None
    AssociationParameters: list[AssociationParameter] | None


class CreditSpecification(TypedDict):
    CPUCredits: str | None


REPEATED_INVOCATION = "repeated_invocation"


class EC2InstanceProvider(ResourceProvider[EC2InstanceProperties]):
    TYPE = "AWS::EC2::Instance"  # Autogenerated. Don't change
    SCHEMA = util.get_schema_path(Path(__file__))  # Autogenerated. Don't change

    def create(
        self,
        request: ResourceRequest[EC2InstanceProperties],
    ) -> ProgressEvent[EC2InstanceProperties]:
        """
        Create a new resource.

        Primary identifier fields:
          - /properties/Id



        Create-only properties:
          - /properties/ElasticGpuSpecifications
          - /properties/Ipv6Addresses
          - /properties/PlacementGroupName
          - /properties/HostResourceGroupArn
          - /properties/ImageId
          - /properties/CpuOptions
          - /properties/PrivateIpAddress
          - /properties/ElasticInferenceAccelerators
          - /properties/EnclaveOptions
          - /properties/HibernationOptions
          - /properties/KeyName
          - /properties/LicenseSpecifications
          - /properties/NetworkInterfaces
          - /properties/AvailabilityZone
          - /properties/SubnetId
          - /properties/LaunchTemplate
          - /properties/SecurityGroups
          - /properties/Ipv6AddressCount

        Read-only properties:
          - /properties/PublicIp
          - /properties/Id
          - /properties/PublicDnsName
          - /properties/PrivateDnsName
          - /properties/PrivateIp



        """
        model = request.desired_state
        ec2 = request.aws_client_factory.ec2
        # TODO: validations

        if not request.custom_context.get(REPEATED_INVOCATION):
            # this is the first time this callback is invoked
            # TODO: idempotency
            params = util.select_attributes(
                model,
                ["InstanceType", "SecurityGroups", "KeyName", "ImageId", "MaxCount", "MinCount"],
            )

            # This Parameters are not defined in the schema but are required by the API
            params["MaxCount"] = 1
            params["MinCount"] = 1

            if model.get("UserData"):
                params["UserData"] = to_str(base64.b64decode(model["UserData"]))

            response = ec2.run_instances(**params)
            model["Id"] = response["Instances"][0]["InstanceId"]
            request.custom_context[REPEATED_INVOCATION] = True
            return ProgressEvent(
                status=OperationStatus.IN_PROGRESS,
                resource_model=model,
                custom_context=request.custom_context,
            )

        response = ec2.describe_instances(InstanceIds=[model["Id"]])
        instance = response["Reservations"][0]["Instances"][0]
        if instance["State"]["Name"] != "running":
            return ProgressEvent(
                status=OperationStatus.IN_PROGRESS,
                resource_model=model,
                custom_context=request.custom_context,
            )

        model["PrivateIp"] = instance["PrivateIpAddress"]
        model["PrivateDnsName"] = instance["PrivateDnsName"]
        model["AvailabilityZone"] = instance["Placement"]["AvailabilityZone"]

        # PublicIp is not guaranteed to be returned by the request:
        # https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.aws_ec2.Instance.html#instancepublicip
        # it says it is supposed to return an empty string, but trying to add an output with the value will result in
        # an error: `Attribute 'PublicIp' does not exist`
        if public_ip := instance.get("PublicIpAddress"):
            model["PublicIp"] = public_ip

        if public_dns_name := instance.get("PublicDnsName"):
            model["PublicDnsName"] = public_dns_name

        return ProgressEvent(
            status=OperationStatus.SUCCESS,
            resource_model=model,
            custom_context=request.custom_context,
        )

    def read(
        self,
        request: ResourceRequest[EC2InstanceProperties],
    ) -> ProgressEvent[EC2InstanceProperties]:
        """
        Fetch resource information


        """
        raise NotImplementedError

    def delete(
        self,
        request: ResourceRequest[EC2InstanceProperties],
    ) -> ProgressEvent[EC2InstanceProperties]:
        """
        Delete a resource


        """
        model = request.desired_state
        ec2 = request.aws_client_factory.ec2
        ec2.terminate_instances(InstanceIds=[model["Id"]])
        # TODO add checking of ec2 instance state
        return ProgressEvent(
            status=OperationStatus.SUCCESS,
            resource_model=model,
            custom_context=request.custom_context,
        )

    def update(
        self,
        request: ResourceRequest[EC2InstanceProperties],
    ) -> ProgressEvent[EC2InstanceProperties]:
        """
        Update a resource


        """
        desired_state = request.desired_state
        ec2 = request.aws_client_factory.ec2

        groups = desired_state.get("SecurityGroups", desired_state.get("SecurityGroupIds"))

        kwargs = {}
        if groups:
            kwargs["Groups"] = groups
        ec2.modify_instance_attribute(
            InstanceId=desired_state["Id"],
            InstanceType={"Value": desired_state["InstanceType"]},
            **kwargs,
        )

        response = ec2.describe_instances(InstanceIds=[desired_state["Id"]])
        instance = response["Reservations"][0]["Instances"][0]
        if instance["State"]["Name"] != "running":
            return ProgressEvent(
                status=OperationStatus.PENDING,
                resource_model=desired_state,
                custom_context=request.custom_context,
            )

        return ProgressEvent(
            status=OperationStatus.SUCCESS,
            resource_model=desired_state,
            custom_context=request.custom_context,
        )
