import json
import logging
import os
import shlex
import sys
import urllib
from collections.abc import Generator
from datetime import UTC
from datetime import datetime as dt
from enum import StrEnum
from importlib.metadata import version
from pathlib import Path
from textwrap import dedent
from typing import Annotated

import humanize
import requests
import typer
from rich import print as rich_print
from rich.progress import Progress, SpinnerColumn, TextColumn
from tzlocal import get_localzone

from ._consts import RH_SAML_URL, AwsConsoleService, AwsRegion
from ._core import (
    assume_role,
    assume_role_with_saml,
    get_aws_account,
    get_aws_accounts,
    get_saml_auth,
    is_kerberos_ticket_valid,
    kinit,
)
from ._models import AwsAccount, AwsCredentials
from ._utils import blend_text, bye, enable_requests_logging, run

app = typer.Typer(rich_markup_mode="rich")
BANNER = r"""
         __                                                         __      __            _
   _____/ /_        ____ __      _______      _________ _____ ___  / /     / /___  ____ _(_)___
  / ___/ __ \______/ __ `/ | /| / / ___/_____/ ___/ __ `/ __ `__ \/ /_____/ / __ \/ __ `/ / __ \
 / /  / / / /_____/ /_/ /| |/ |/ (__  )_____(__  ) /_/ / / / / / / /_____/ / /_/ / /_/ / / / / /
/_/  /_/ /_/      \__,_/ |__/|__/____/     /____/\__,_/_/ /_/ /_/_/     /_/\____/\__, /_/_/ /_/
                                                                                /____/
"""
APP_NAME = "rh-aws-saml-login"
APP_DIR = Path(typer.get_app_dir(APP_NAME))
APP_DIR.mkdir(exist_ok=True)
ACCOUNT_CACHE = APP_DIR / "account_cache.json"

SCRIPT_START_TIME = dt.now(UTC)

logger = logging.getLogger(__name__)


class OutputFormat(StrEnum):
    """Supported output formats"""

    JSON = "json"
    ENV = "env"


def get_export_environment_variables(
    account: AwsAccount,
    credentials: AwsCredentials,
    region: str,
) -> dict:
    """Get the environment variables to export AWS credentials."""
    return {
        "AWS_ACCOUNT_NAME": account.name,
        "AWS_ACCOUNT_UID": account.uid,
        "AWS_ROLE_NAME": account.role_name,
        "AWS_ROLE_ARN": account.role_arn,
        "AWS_ACCESS_KEY_ID": credentials.access_key,
        "AWS_SECRET_ACCESS_KEY": credentials.secret_key,
        "AWS_SESSION_TOKEN": credentials.session_token,
        "AWS_REGION": region,
    }


def open_aws_shell(
    account: AwsAccount,
    credentials: AwsCredentials,
    region: str,
    command: list[str] | str | None = None,
    *,
    quiet: bool = False,
) -> None:
    if not quiet and not command:
        rich_print(
            dedent(f"""
            Spawning a new shell. Use exit or CTRL+d to leave it!

            :nerd_face: {account.name}
            :rocket: {account.role_name}
            :hourglass: {humanize.naturaltime(credentials.expiration, when=SCRIPT_START_TIME)} ({credentials.expiration.astimezone(tz=get_localzone())})
        """)
        )
    if not command:
        command = os.environ.get("SHELL", "/bin/bash")
    run(
        command,
        check=False,
        capture_output=False,
        env=get_export_environment_variables(account, credentials, region),
    )


def open_aws_console(
    open_command: str, credentials: AwsCredentials, console_service: str | None = None
) -> None:
    """Open the AWS console in a browser.

    See https://docs.aws.amazon.com/IAM/latest/UserGuide/id_roles_providers_enable-console-custom-url.html
    """
    aws_federated_signin_endpoint = "https://signin.aws.amazon.com/federation"
    # Get a sign-in token from the AWS sign-in federation endpoint.
    response = requests.get(
        aws_federated_signin_endpoint,
        params={
            "Action": "getSigninToken",
            "SessionDuration": str(credentials.session_timeout_seconds),
            "Session": json.dumps({
                "sessionId": credentials.access_key,
                "sessionKey": credentials.secret_key,
                "sessionToken": credentials.session_token,
            }),
        },
        timeout=10,
    )
    if response.status_code == requests.codes.BAD_REQUEST:
        logger.error(
            "Failed to get a sign-in token. Try lowering the session timeout value via --session-timeout."
        )
        sys.exit(1)

    signin_token = json.loads(response.text)
    # Make a federated URL that can be used to sign into the AWS Management Console.
    query_string = urllib.parse.urlencode({
        "Action": "login",
        "Issuer": "redhat.com",
        "Destination": f"https://{credentials.region}.console.aws.amazon.com/{console_service or ''}",
        "SigninToken": signin_token["SigninToken"],
    })
    federated_url = f"{aws_federated_signin_endpoint}?{query_string}"
    run([*shlex.split(open_command), federated_url], check=False, capture_output=False)


def display_credentials(
    account: AwsAccount, credentials: AwsCredentials, region: str, output: OutputFormat
) -> None:
    """Display AWS credentials in the specified output format."""
    env_vars = get_export_environment_variables(account, credentials, region)
    match output:
        case OutputFormat.JSON:
            print(json.dumps(env_vars, indent=2))  # noqa: T201
        case OutputFormat.ENV:
            for key, value in env_vars.items():
                print(f"{key}={value}")  # noqa: T201


def write_accounts_cache(accounts: list[str]) -> None:
    """Write accounts cache to disk."""
    json.dump(accounts, ACCOUNT_CACHE.open("w"))


def read_accounts_cache() -> list[str]:
    """Read accounts cache from disk."""
    if ACCOUNT_CACHE.exists():
        return json.load(ACCOUNT_CACHE.open())
    return []


def complete_account(ctx: typer.Context, incomplete: str) -> Generator[str, None, None]:  # noqa: ARG001
    for name in read_accounts_cache():
        if name.startswith(incomplete):
            yield name


def version_callback(*, value: bool) -> None:
    if value:
        rich_print(f"Version: {version(APP_NAME)}")
        raise typer.Exit


@app.command(epilog="Made with [red]:heart:[/] by [blue]https://github.com/app-sre[/]")
def cli(  # noqa: PLR0917
    account_name: Annotated[
        str | None,
        typer.Argument(
            help="AWS account name. Supports: account name (e.g., 'my-account'), account/role format (e.g., 'my-account/PowerUserAccess'), or '.' for $AWS_ACCOUNT_NAME environment variable.",
            autocompletion=complete_account,
        ),
    ] = None,
    command: Annotated[list[str] | None, typer.Argument(help="Command to run")] = None,
    region: Annotated[
        AwsRegion,
        typer.Option(
            help="AWS region",
            envvar="RH_AWS_REGION",
        ),
    ] = AwsRegion.US_EAST_1,
    saml_url: Annotated[
        str,
        typer.Option(
            help="SAML URL",
        ),
    ] = RH_SAML_URL,
    session_timeout: Annotated[
        int,
        typer.Option(
            help="Session timeout in minutes. Default: 60 minutes. Max value depends on the AWS IAM role.",
            envvar="RH_AWS_SESSION_TIMEOUT",
        ),
    ] = 60,
    open_command: Annotated[
        str,
        typer.Option(
            help="Command to open the browser (e.g. 'xdg-open' on Linux)",
            envvar="RH_AWS_SAML_LOGIN_OPEN_COMMAND",
        ),
    ] = "open",
    console_service: Annotated[
        AwsConsoleService | None,
        typer.Option(
            help="Directly open this AWS console service",
            envvar="RH_AWS_CONSOLE_SERVICE",
        ),
    ] = None,
    assume_uid: Annotated[
        str | None,
        typer.Option(
            help="Define the target AWS account UID to assume",
        ),
    ] = None,
    assume_role: Annotated[
        str,
        typer.Option(
            help="Define the role name to assume",
        ),
    ] = "role/OrganizationAccountAccessRole",
    *,
    debug: Annotated[bool, typer.Option(help="Enable debug mode")] = False,
    console: Annotated[
        bool,
        typer.Option(help="Open the AWS console in browser instead of a local shell"),
    ] = False,
    kerberos_keytab: Annotated[
        str | None,
        typer.Option(
            help="Base64 encoded Kerberos keytab file. If set, it will be used to acquire a Kerberos ticket.",
            envvar="RH_KERBEROS_KEYTAB",
        ),
    ] = None,
    kerberos_principal: Annotated[
        str,
        typer.Option(
            help="Kerberos principal to use for acquiring a Kerberos ticket.",
            envvar="RH_KERBEROS_PRINCIPAL",
        ),
    ] = "",
    quiet: Annotated[
        bool,
        typer.Option(
            help="Do not display any logging output or other messages!",
            envvar="RH_QUIET",
        ),
    ] = False,
    output: Annotated[
        OutputFormat | None,
        typer.Option(
            help="Instead of opening the AWS console or shell, output the credentials in the specified format. Also implies --quiet.",
            envvar="RH_OUTPUT",
            case_sensitive=False,
        ),
    ] = None,
    display_banner: Annotated[
        bool,
        typer.Option(
            help="Display a shiny banner",
            envvar="RH_DISPLAY_BANNER",
        ),
    ] = True,
    version: Annotated[  # noqa: ARG001
        bool | None, typer.Option("--version", callback=version_callback)
    ] = None,
) -> None:
    """Login to AWS using SAML."""
    log_level = logging.INFO
    if output:
        quiet = True
    if debug:
        log_level = logging.DEBUG
    if quiet:
        log_level = logging.ERROR
    logging.basicConfig(level=log_level, format="%(message)s")

    if display_banner and not quiet:
        rich_print(blend_text(BANNER, (32, 32, 255), (255, 32, 255)))
    if debug:
        enable_requests_logging()

    role = None
    if account_name and "/" in account_name:
        # when account_name contains a '/', split it into account_name and role
        account_name, role = account_name.split("/", 1)

    accounts = _main(
        account_name=account_name,
        role=role,
        region=region,
        console=console,
        saml_url=saml_url,
        session_timeout_seconds=session_timeout * 60,
        command=command,
        open_command=open_command,
        console_service=console_service,
        assume_uid=assume_uid,
        assume_role_name=assume_role,
        kerberos_keytab=kerberos_keytab,
        kerberos_principal=kerberos_principal,
        output=output,
        quiet=quiet,
    )
    write_accounts_cache(accounts)


def _main(  # noqa: PLR0917
    account_name: str | None,
    role: str | None,
    region: str,
    saml_url: str,
    session_timeout_seconds: int,
    command: list[str] | None,
    open_command: str,
    console_service: str | None,
    assume_uid: str | None,
    assume_role_name: str,
    kerberos_keytab: str | None = None,
    kerberos_principal: str = "",
    output: OutputFormat | None = None,
    *,
    console: bool,
    quiet: bool,
) -> list[str]:
    with Progress(
        SpinnerColumn(finished_text="✅"),
        TextColumn("[progress.description]{task.description}"),
        disable=quiet,
    ) as progress:
        task = progress.add_task(
            description="Test for a valid Kerberos ticket ...", total=1
        )
        if not is_kerberos_ticket_valid():
            progress.stop()
            logger.info("No valid Kerberos ticket found. Acquiring one ...")
            kinit(kerberos_keytab, kerberos_principal)
            progress.start()
        progress.update(task, completed=1)

        task = progress.add_task(description="Getting SAML token ...", total=1)
        aws_url, saml_token = get_saml_auth(saml_url)
        progress.update(task, completed=1)

        task = progress.add_task(description="Getting AWS accounts ...", total=1)
        aws_accounts = get_aws_accounts(
            aws_url, saml_token, session_timeout_seconds, region
        )

        progress.stop()
        if not (account := get_aws_account(aws_accounts, account_name, role)):
            if role:
                logger.error("Account with role not found: %s/%s", account_name, role)
            else:
                logger.error("Account not found: %s", account_name)
            sys.exit(1)
        progress.start()
        progress.update(task, completed=1)

        task = progress.add_task(
            description="Getting temporary AWS credentials ...", total=1
        )
        credentials = assume_role_with_saml(account, saml_token)
        progress.update(task, completed=1)

        if assume_uid:
            account = AwsAccount(
                name=assume_uid,
                uid=assume_uid,
                role_name=account.name,
                role_arn=f"arn:aws:iam::{assume_uid}:{assume_role_name}",
                session_timeout_seconds=session_timeout_seconds,
                region=region,
            )
            task = progress.add_task(description="Assume role ...", total=1)
            credentials = assume_role(account, credentials)
            progress.update(task, completed=1)

    if output:
        display_credentials(account, credentials, region, output)
    elif console:
        open_aws_console(open_command, credentials, console_service)
    else:
        open_aws_shell(account, credentials, region, command, quiet=quiet)
    if not quiet:
        bye()
    return [acc.name for acc in aws_accounts]
