"""Snapshot 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 a snapshot from another AWS resource."""

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


@dataclass
class SnapshotInfo:
    """Basic snapshot info from AWS."""

    id: str
    size_gb: int
    created: datetime
    description: str
    volume_id: str


@dataclass
class SnapshotReport:
    """Snapshot with its references."""

    id: str
    size_gb: int
    created: datetime
    description: str
    volume_id: str
    references: list[Reference] = field(default_factory=list)

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

    @property
    def is_ami_created(self) -> bool:
        return "CreateImage" in self.description

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


def get_snapshots(
    session: boto3.Session,
) -> tuple[dict[str, SnapshotInfo], dict[str, list[Reference]]]:
    """Fetch all snapshots owned by this account.

    Returns (snapshots dict, backup references dict).
    """
    ec2 = session.client("ec2")
    sts = session.client("sts")

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

    paginator = ec2.get_paginator("describe_snapshots")

    snapshots = {}
    backup_refs: dict[str, list[Reference]] = {}

    for page in paginator.paginate(OwnerIds=[account_id]):
        for snap in page["Snapshots"]:
            created = snap["StartTime"]
            if isinstance(created, str):
                created = datetime.fromisoformat(created.replace("Z", "+00:00"))

            snap_id = snap["SnapshotId"]

            # Check if managed by AWS Backup
            tags = {t["Key"]: t["Value"] for t in snap.get("Tags", [])}
            if "aws:backup:source-resource" in tags:
                # Use Name tag if available, otherwise just mark as AWS Backup
                name = tags.get("Name", "AWS Backup")
                backup_refs[snap_id] = [Reference(type="Backup", name=name)]

            snapshots[snap_id] = SnapshotInfo(
                id=snap_id,
                size_gb=snap["VolumeSize"],
                created=created,
                description=snap.get("Description", ""),
                volume_id=snap.get("VolumeId", ""),
            )

    console.print(
        f"Found {len(snapshots)} snapshots ({len(backup_refs)} AWS Backup managed)"
    )
    return snapshots, backup_refs


def get_ami_references(session: boto3.Session) -> dict[str, list[Reference]]:
    """Get snapshot references from AMIs."""
    ec2 = session.client("ec2")
    sts = session.client("sts")

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

    paginator = ec2.get_paginator("describe_images")

    refs: dict[str, list[Reference]] = {}
    for page in paginator.paginate(Owners=[account_id]):
        for image in page["Images"]:
            ami_name = image.get("Name", image["ImageId"])
            for bdm in image.get("BlockDeviceMappings", []):
                ebs = bdm.get("Ebs", {})
                if snapshot_id := ebs.get("SnapshotId"):
                    refs.setdefault(snapshot_id, []).append(
                        Reference(type="AMI", name=ami_name)
                    )

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


def get_lt_references(session: boto3.Session) -> dict[str, list[Reference]]:
    """Get snapshot references from Launch Template block device mappings."""
    ec2 = session.client("ec2")

    # 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", {})

                # Check block device mappings for snapshot IDs
                for bdm in data.get("BlockDeviceMappings", []):
                    ebs = bdm.get("Ebs", {})
                    if snapshot_id := ebs.get("SnapshotId"):
                        refs.setdefault(snapshot_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 snapshot references from Launch Configuration block device mappings."""
    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"]:
            lc_name = lc["LaunchConfigurationName"]
            for bdm in lc.get("BlockDeviceMappings", []):
                ebs = bdm.get("Ebs", {})
                if snapshot_id := ebs.get("SnapshotId"):
                    refs.setdefault(snapshot_id, []).append(
                        Reference(type="LaunchConfig", name=lc_name)
                    )

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


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

    # Fetch all data
    snapshots, backup_refs = get_snapshots(session)
    ami_refs = get_ami_references(session)
    lt_refs = get_lt_references(session)
    lc_refs = get_lc_references(session)

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

    # Build report
    report = []
    for snap_id, snap_info in snapshots.items():
        report.append(
            SnapshotReport(
                id=snap_id,
                size_gb=snap_info.size_gb,
                created=snap_info.created,
                description=snap_info.description,
                volume_id=snap_info.volume_id,
                references=all_refs.get(snap_id, []),
            )
        )

    return report


def delete_snapshot(session: boto3.Session, snapshot_id: str) -> None:
    """Delete a snapshot."""
    ec2 = session.client("ec2")
    ec2.delete_snapshot(SnapshotId=snapshot_id)
    console.print(f"  Deleted snapshot {snapshot_id}")
