"""AMI usage report."""

from dataclasses import dataclass, field
from datetime import datetime

import boto3
from rich.console import Console

console = Console()


@dataclass
class Reference:
    """A reference to an AMI from another AWS resource."""

    type: str  # LaunchTemplate, LaunchConfig, Instance
    name: str  # human-readable identifier


@dataclass
class AmiInfo:
    """Basic AMI info from AWS."""

    id: str
    name: str
    created: datetime
    snapshot_ids: list[str]


@dataclass
class AmiReport:
    """AMI with its references."""

    id: str
    name: str
    created: datetime
    snapshot_ids: list[str]
    instance_count: int = 0
    references: list[Reference] = field(default_factory=list)

    @property
    def is_orphaned(self) -> bool:
        return len(self.references) == 0 and self.instance_count == 0

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "name": self.name,
            "created": self.created.isoformat(),
            "snapshot_ids": self.snapshot_ids,
            "instance_count": self.instance_count,
            "references": [{"type": r.type, "name": r.name} for r in self.references],
            "is_orphaned": self.is_orphaned,
        }


def get_amis(session: boto3.Session) -> dict[str, AmiInfo]:
    """Fetch all AMIs owned by this account. Returns {ami_id: AmiInfo}."""
    ec2 = session.client("ec2")
    sts = session.client("sts")

    account_id = sts.get_caller_identity()["Account"]

    paginator = ec2.get_paginator("describe_images")

    amis = {}
    for page in paginator.paginate(Owners=[account_id]):
        for image in page["Images"]:
            # Extract snapshot IDs from block device mappings
            snapshot_ids = []
            for bdm in image.get("BlockDeviceMappings", []):
                ebs = bdm.get("Ebs", {})
                if snapshot_id := ebs.get("SnapshotId"):
                    snapshot_ids.append(snapshot_id)

            # Parse creation date
            created = datetime.fromisoformat(
                image["CreationDate"].replace("Z", "+00:00")
            )

            amis[image["ImageId"]] = AmiInfo(
                id=image["ImageId"],
                name=image.get("Name", ""),
                created=created,
                snapshot_ids=snapshot_ids,
            )

    console.print(f"Found {len(amis)} AMIs owned by account")
    return amis


def get_lt_references(session: boto3.Session) -> dict[str, list[Reference]]:
    """Get AMI references from Launch Template versions.

    Checks ALL versions of each template, not just the default.
    """
    ec2 = session.client("ec2")

    # First get all launch templates
    lt_paginator = ec2.get_paginator("describe_launch_templates")
    templates = []
    for page in lt_paginator.paginate():
        for lt in page["LaunchTemplates"]:
            templates.append(
                {"id": lt["LaunchTemplateId"], "name": lt["LaunchTemplateName"]}
            )

    refs: dict[str, list[Reference]] = {}

    # For each template, get all versions
    for lt in templates:
        version_paginator = ec2.get_paginator("describe_launch_template_versions")
        for page in version_paginator.paginate(LaunchTemplateId=lt["id"]):
            for version in page["LaunchTemplateVersions"]:
                version_num = version["VersionNumber"]
                data = version.get("LaunchTemplateData", {})
                ami_id = data.get("ImageId")

                if ami_id:
                    refs.setdefault(ami_id, []).append(
                        Reference(
                            type="LaunchTemplate",
                            name=f"{lt['name']}:v{version_num}",
                        )
                    )

    console.print(
        f"Found {sum(len(v) for v in refs.values())} Launch Template references"
    )
    return refs


def get_lc_references(session: boto3.Session) -> dict[str, list[Reference]]:
    """Get AMI references from Launch Configurations (deprecated but still exist)."""
    autoscaling = session.client("autoscaling")
    paginator = autoscaling.get_paginator("describe_launch_configurations")

    refs: dict[str, list[Reference]] = {}
    for page in paginator.paginate():
        for lc in page["LaunchConfigurations"]:
            ami_id = lc.get("ImageId")
            if ami_id:
                refs.setdefault(ami_id, []).append(
                    Reference(type="LaunchConfig", name=lc["LaunchConfigurationName"])
                )

    console.print(
        f"Found {sum(len(v) for v in refs.values())} Launch Configuration references"
    )
    return refs


def get_ec2_instance_counts(session: boto3.Session) -> dict[str, int]:
    """Get EC2 instance counts per AMI.

    Includes both running and stopped instances - stopped instances
    still reference their AMI.
    """
    ec2 = session.client("ec2")
    paginator = ec2.get_paginator("describe_instances")

    counts: dict[str, int] = {}
    for page in paginator.paginate(
        Filters=[{"Name": "instance-state-name", "Values": ["running", "stopped"]}]
    ):
        for reservation in page["Reservations"]:
            for instance in reservation["Instances"]:
                ami_id = instance.get("ImageId")
                if ami_id:
                    counts[ami_id] = counts.get(ami_id, 0) + 1

    console.print(f"Found {sum(counts.values())} EC2 instances (running + stopped)")
    return counts


def build_report(
    session: boto3.Session | None = None,
) -> list[AmiReport]:
    """Build AMI usage report."""
    if session is None:
        session = boto3.Session()

    # Fetch all data
    amis = get_amis(session)
    lt_refs = get_lt_references(session)
    lc_refs = get_lc_references(session)
    instance_counts = get_ec2_instance_counts(session)

    # Merge references
    all_refs: dict[str, list[Reference]] = {}
    for refs in [lt_refs, lc_refs]:
        for ami_id, ref_list in refs.items():
            all_refs.setdefault(ami_id, []).extend(ref_list)

    # Build report
    report = []
    for ami_id, ami_info in amis.items():
        report.append(
            AmiReport(
                id=ami_id,
                name=ami_info.name,
                created=ami_info.created,
                snapshot_ids=ami_info.snapshot_ids,
                instance_count=instance_counts.get(ami_id, 0),
                references=all_refs.get(ami_id, []),
            )
        )

    return report
