"""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 SnapshotInfo:
    """Snapshot info."""

    id: str
    size_gb: int


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

    id: str
    name: str
    created: datetime
    snapshots: list[SnapshotInfo]


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

    id: str
    name: str
    created: datetime
    snapshots: list[SnapshotInfo]
    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

    @property
    def snapshot_ids(self) -> list[str]:
        return [s.id for s in self.snapshots]

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


def get_snapshot_sizes(
    session: boto3.Session, snapshot_ids: list[str]
) -> dict[str, int]:
    """Fetch snapshot sizes. Returns {snapshot_id: size_gb}."""
    if not snapshot_ids:
        return {}

    ec2 = session.client("ec2")
    sizes: dict[str, int] = {}

    # describe_snapshots accepts max 1000 IDs per call
    for i in range(0, len(snapshot_ids), 1000):
        batch = snapshot_ids[i : i + 1000]
        paginator = ec2.get_paginator("describe_snapshots")
        for page in paginator.paginate(SnapshotIds=batch):
            for snap in page["Snapshots"]:
                sizes[snap["SnapshotId"]] = snap["VolumeSize"]

    return sizes


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")

    # First pass: collect AMIs and snapshot IDs
    ami_data: dict[str, dict] = {}
    all_snapshot_ids: list[str] = []

    for page in paginator.paginate(Owners=[account_id]):
        for image in page["Images"]:
            snapshot_ids = []
            for bdm in image.get("BlockDeviceMappings", []):
                ebs = bdm.get("Ebs", {})
                if snapshot_id := ebs.get("SnapshotId"):
                    snapshot_ids.append(snapshot_id)
                    all_snapshot_ids.append(snapshot_id)

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

            ami_data[image["ImageId"]] = {
                "name": image.get("Name", ""),
                "created": created,
                "snapshot_ids": snapshot_ids,
            }

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

    # Fetch snapshot sizes
    snapshot_sizes = get_snapshot_sizes(session, all_snapshot_ids)
    console.print(f"Found {len(snapshot_sizes)} snapshots")

    # Build AmiInfo with snapshot details
    amis = {}
    for ami_id, data in ami_data.items():
        snapshots = [
            SnapshotInfo(id=sid, size_gb=snapshot_sizes.get(sid, 0))
            for sid in data["snapshot_ids"]
        ]
        amis[ami_id] = AmiInfo(
            id=ami_id,
            name=data["name"],
            created=data["created"],
            snapshots=snapshots,
        )

    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,
                snapshots=ami_info.snapshots,
                instance_count=instance_counts.get(ami_id, 0),
                references=all_refs.get(ami_id, []),
            )
        )

    return report


def delete_ami(session: boto3.Session, ami_id: str, snapshot_ids: list[str]) -> None:
    """Deregister AMI and delete associated snapshots."""
    ec2 = session.client("ec2")

    # Deregister AMI first
    ec2.deregister_image(ImageId=ami_id)
    console.print(f"  Deregistered AMI {ami_id}")

    # Delete snapshots
    for snap_id in snapshot_ids:
        ec2.delete_snapshot(SnapshotId=snap_id)
        console.print(f"  Deleted snapshot {snap_id}")
