from __future__ import annotations

from datetime import datetime, timezone
from typing import Any

from awspipe.aws import AwsSessionContext
from awspipe.providers.base import Column, TableProvider


COMMON_COLUMNS = (
    Column("profile_name", "VARCHAR", "The AWS CLI profile used for the fetch."),
    Column("account_id", "VARCHAR", "The AWS account identifier."),
    Column("region", "VARCHAR", "The AWS region for the row."),
    Column("arn", "VARCHAR", "The ARN for the resource when available."),
    Column("title", "VARCHAR", "A display-friendly title for the resource."),
)


def _tag_name(tags: list[dict[str, str]] | None) -> str | None:
    if not tags:
        return None
    for tag in tags:
        if tag.get("Key") == "Name":
            return tag.get("Value")
    return None


def _ts(value: datetime | None) -> str | None:
    if value is None:
        return None
    return value.astimezone(timezone.utc).isoformat()


def fetch_aws_account(context: AwsSessionContext) -> list[dict[str, Any]]:
    identity = context.caller_identity
    return [
        {
            "profile_name": context.profile or "default",
            "account_id": context.account_id,
            "region": context.region,
            "arn": identity["Arn"],
            "title": context.account_id,
            "user_id": identity["UserId"],
        }
    ]


def fetch_aws_region(context: AwsSessionContext) -> list[dict[str, Any]]:
    client = context.session.client("ec2", region_name=context.region)
    regions = client.describe_regions(AllRegions=False)["Regions"]
    return [
        {
            "profile_name": context.profile or "default",
            "account_id": context.account_id,
            "region": item["RegionName"],
            "arn": None,
            "title": item["RegionName"],
            "opt_in_status": item.get("OptInStatus"),
            "endpoint": item.get("Endpoint"),
        }
        for item in regions
    ]


def fetch_aws_s3_bucket(context: AwsSessionContext) -> list[dict[str, Any]]:
    client = context.session.client("s3", region_name=context.region)
    response = client.list_buckets()
    rows: list[dict[str, Any]] = []
    for bucket in response.get("Buckets", []):
        bucket_name = bucket["Name"]
        location = client.get_bucket_location(Bucket=bucket_name).get("LocationConstraint") or "us-east-1"
        rows.append(
            {
                "profile_name": context.profile or "default",
                "account_id": context.account_id,
                "region": location,
                "arn": f"arn:aws:s3:::{bucket_name}",
                "title": bucket_name,
                "name": bucket_name,
                "creation_date": _ts(bucket.get("CreationDate")),
            }
        )
    return rows


def fetch_aws_ec2_instance(context: AwsSessionContext) -> list[dict[str, Any]]:
    client = context.session.client("ec2", region_name=context.region)
    rows: list[dict[str, Any]] = []
    next_token: str | None = None
    while True:
        request = {"NextToken": next_token} if next_token else {}
        page = client.describe_instances(**request)
        for reservation in page.get("Reservations", []):
            for instance in reservation.get("Instances", []):
                instance_id = instance["InstanceId"]
                rows.append(
                    {
                        "profile_name": context.profile or "default",
                        "account_id": context.account_id,
                        "region": context.region,
                        "arn": f"arn:aws:ec2:{context.region}:{context.account_id}:instance/{instance_id}",
                        "title": _tag_name(instance.get("Tags")) or instance_id,
                        "instance_id": instance_id,
                        "instance_type": instance.get("InstanceType"),
                        "state": instance.get("State", {}).get("Name"),
                        "vpc_id": instance.get("VpcId"),
                        "subnet_id": instance.get("SubnetId"),
                        "private_ip_address": instance.get("PrivateIpAddress"),
                        "public_ip_address": instance.get("PublicIpAddress"),
                        "launch_time": _ts(instance.get("LaunchTime")),
                    }
                )
        next_token = page.get("NextToken")
        if not next_token:
            return rows


def fetch_aws_vpc(context: AwsSessionContext) -> list[dict[str, Any]]:
    client = context.session.client("ec2", region_name=context.region)
    response = client.describe_vpcs()
    rows: list[dict[str, Any]] = []
    for vpc in response.get("Vpcs", []):
        vpc_id = vpc["VpcId"]
        rows.append(
            {
                "profile_name": context.profile or "default",
                "account_id": context.account_id,
                "region": context.region,
                "arn": f"arn:aws:ec2:{context.region}:{context.account_id}:vpc/{vpc_id}",
                "title": _tag_name(vpc.get("Tags")) or vpc_id,
                "vpc_id": vpc_id,
                "cidr_block": vpc.get("CidrBlock"),
                "state": vpc.get("State"),
                "is_default": vpc.get("IsDefault"),
            }
        )
    return rows


def fetch_aws_subnet(context: AwsSessionContext) -> list[dict[str, Any]]:
    client = context.session.client("ec2", region_name=context.region)
    response = client.describe_subnets()
    rows: list[dict[str, Any]] = []
    for subnet in response.get("Subnets", []):
        subnet_id = subnet["SubnetId"]
        rows.append(
            {
                "profile_name": context.profile or "default",
                "account_id": context.account_id,
                "region": context.region,
                "arn": f"arn:aws:ec2:{context.region}:{context.account_id}:subnet/{subnet_id}",
                "title": _tag_name(subnet.get("Tags")) or subnet_id,
                "subnet_id": subnet_id,
                "vpc_id": subnet.get("VpcId"),
                "cidr_block": subnet.get("CidrBlock"),
                "availability_zone": subnet.get("AvailabilityZone"),
                "available_ip_address_count": subnet.get("AvailableIpAddressCount"),
            }
        )
    return rows


def fetch_aws_security_group(context: AwsSessionContext) -> list[dict[str, Any]]:
    client = context.session.client("ec2", region_name=context.region)
    response = client.describe_security_groups()
    rows: list[dict[str, Any]] = []
    for group in response.get("SecurityGroups", []):
        group_id = group["GroupId"]
        group_name = group["GroupName"]
        rows.append(
            {
                "profile_name": context.profile or "default",
                "account_id": context.account_id,
                "region": context.region,
                "arn": f"arn:aws:ec2:{context.region}:{context.account_id}:security-group/{group_id}",
                "title": group_name,
                "group_id": group_id,
                "group_name": group_name,
                "vpc_id": group.get("VpcId"),
                "description": group.get("Description"),
                "owner_id": group.get("OwnerId"),
            }
        )
    return rows


PROVIDERS = [
    TableProvider(
        name="aws_account",
        description="AWS account and caller identity details for the active session.",
        columns=COMMON_COLUMNS + (Column("user_id", "VARCHAR", "The caller user identifier."),),
        scope="global",
        fetcher=fetch_aws_account,
    ),
    TableProvider(
        name="aws_region",
        description="Enabled AWS regions visible to the current account.",
        columns=COMMON_COLUMNS + (
            Column("opt_in_status", "VARCHAR", "Region opt-in status."),
            Column("endpoint", "VARCHAR", "Region endpoint."),
        ),
        scope="global",
        fetcher=fetch_aws_region,
    ),
    TableProvider(
        name="aws_s3_bucket",
        description="S3 buckets visible to the current credentials.",
        columns=COMMON_COLUMNS + (
            Column("name", "VARCHAR", "The bucket name."),
            Column("creation_date", "TIMESTAMP", "Bucket creation timestamp."),
        ),
        scope="global",
        fetcher=fetch_aws_s3_bucket,
    ),
    TableProvider(
        name="aws_ec2_instance",
        description="EC2 instances in the selected region.",
        columns=COMMON_COLUMNS + (
            Column("instance_id", "VARCHAR", "The EC2 instance identifier."),
            Column("instance_type", "VARCHAR", "The EC2 instance type."),
            Column("state", "VARCHAR", "The EC2 instance state."),
            Column("vpc_id", "VARCHAR", "The VPC identifier."),
            Column("subnet_id", "VARCHAR", "The subnet identifier."),
            Column("private_ip_address", "VARCHAR", "The private IP address."),
            Column("public_ip_address", "VARCHAR", "The public IP address."),
            Column("launch_time", "TIMESTAMP", "Launch timestamp."),
        ),
        scope="regional",
        fetcher=fetch_aws_ec2_instance,
    ),
    TableProvider(
        name="aws_vpc",
        description="VPCs in the selected region.",
        columns=COMMON_COLUMNS + (
            Column("vpc_id", "VARCHAR", "The VPC identifier."),
            Column("cidr_block", "VARCHAR", "The IPv4 CIDR block."),
            Column("state", "VARCHAR", "The VPC state."),
            Column("is_default", "BOOLEAN", "Whether the VPC is the default VPC."),
        ),
        scope="regional",
        fetcher=fetch_aws_vpc,
    ),
    TableProvider(
        name="aws_subnet",
        description="Subnets in the selected region.",
        columns=COMMON_COLUMNS + (
            Column("subnet_id", "VARCHAR", "The subnet identifier."),
            Column("vpc_id", "VARCHAR", "The VPC identifier."),
            Column("cidr_block", "VARCHAR", "The subnet CIDR block."),
            Column("availability_zone", "VARCHAR", "The availability zone."),
            Column("available_ip_address_count", "BIGINT", "Available IP addresses in the subnet."),
        ),
        scope="regional",
        fetcher=fetch_aws_subnet,
    ),
    TableProvider(
        name="aws_security_group",
        description="Security groups in the selected region.",
        columns=COMMON_COLUMNS + (
            Column("group_id", "VARCHAR", "The security group identifier."),
            Column("group_name", "VARCHAR", "The security group name."),
            Column("vpc_id", "VARCHAR", "The VPC identifier."),
            Column("description", "VARCHAR", "The group description."),
            Column("owner_id", "VARCHAR", "The AWS account owner identifier."),
        ),
        scope="regional",
        fetcher=fetch_aws_security_group,
    ),
]


class StaticProviderRegistry:
    def __init__(self) -> None:
        self._providers = {provider.name: provider for provider in PROVIDERS}

    def get(self, table_name: str) -> TableProvider | None:
        return self._providers.get(table_name)

    def all(self) -> list[TableProvider]:
        return list(self._providers.values())
