from typing import Annotated
import boto3
import boto3.session

from admyral.action import action, ArgumentMetadata
from admyral.context import ctx
from admyral.typings import JsonValue
from admyral.utils.temporalio import serialize_datetime_in_json
from admyral.actions.integrations.shared.aws import AWSSecret


# TODO: add documentation for cross-account role setup
@action(
    display_name="List Vulnerabilities",
    display_namespace="Amazon Inspector2",
    description="List Amazon Inspector2 vulnerabilities",
    secrets_placeholders=["AWS_SECRET"],
)
def list_amazon_inspector2_vulnerabilities(
    start_time: Annotated[
        str | None,
        ArgumentMetadata(
            display_name="Start Time",
            description="The start time of the findings as ISO-8601 formatted string (inclusive)",
        ),
    ] = None,
    end_time: Annotated[
        str | None,
        ArgumentMetadata(
            display_name="End Time",
            description="The end time of the findings as ISO-8601 formatted string (inclusive)",
        ),
    ] = None,
    limit: Annotated[
        int,
        ArgumentMetadata(
            display_name="Limit",
            description="The maximum number of findings to return.",
        ),
    ] = 1000,
) -> list[dict[str, JsonValue]]:
    aws_secret = ctx.get().secrets.get("AWS_SECRET")
    aws_secret = AWSSecret.model_validate(aws_secret)

    session = boto3.session.Session(
        aws_access_key_id=aws_secret.aws_access_key_id,
        aws_secret_access_key=aws_secret.aws_secret_access_key,
    )
    inspector_client = session.client("inspector2")

    filter = {}
    if start_time or end_time:
        date_filter = {}
        if start_time:
            date_filter["startInclusive"] = start_time
        if end_time:
            date_filter["endInclusive"] = end_time
        filter["firstObservedAt"] = [date_filter]

    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/inspector2/client/list_findings.html
    paginator = inspector_client.get_paginator("list_findings")
    findings = []
    for page in paginator.paginate(
        filterCriteria=filter, PaginationConfig={"PageSize": 100}
    ):
        findings.extend(page.get("findings", []))
        if len(findings) >= limit:
            findings = findings[:limit]
            break

    findings = findings[:limit]

    return serialize_datetime_in_json(findings[:limit])
