import boto3
from pathlib import Path
from subprocess import run, CalledProcessError
from dotenv import set_key
import os
import typer
from typing import Optional

app = typer.Typer(help="AWS SSO Credentials Manager")


@app.command()
def load_config(
    profile_name: str = typer.Option(
        "genai", "--profile", "-p", help="AWS SSO profile name"
    ),
    sso_start_url: str = typer.Option(
        "https://iqviaportal.awsapps.com/start/#", "--sso-url", help="SSO start URL"
    ),
    sso_account_id: str = typer.Option(
        "387588690219", "--account-id", help="AWS account ID"
    ),
    sso_role_name: str = typer.Option(
        "TS-SSO-bedrock+sagemaker", "--role-name", help="SSO role name"
    ),
    region: str = typer.Option("eu-central-1", "--region", "-r", help="AWS region"),
    sso_region: str = typer.Option("us-east-1", "--sso-region", help="SSO region"),
    sso_session_name: Optional[str] = typer.Option(
        None, "--session-name", help="SSO session name (defaults to profile name)"
    ),
    env_file: str = typer.Option(".env", "--env-file", help="Path to .env file"),
    auto_login: bool = typer.Option(
        True, "--auto-login/--no-auto-login", help="Automatically run SSO login"
    ),
    update_config: bool = typer.Option(
        True, "--update-config/--no-update-config", help="Update AWS config file"
    ),
) -> None:
    """
    Retrieve AWS SSO credentials and save them to environment variables and .env file.
    """

    # Use profile name as session name if not provided
    if sso_session_name is None:
        sso_session_name = profile_name

    # Create configuration string
    config_str = f"""[profile {profile_name}]
sso_session = {sso_session_name}
sso_account_id = {sso_account_id}
sso_role_name = {sso_role_name}
region = {region}
[sso-session {sso_session_name}]
sso_start_url = {sso_start_url}
sso_region = {sso_region}
sso_registration_scopes = sso:account:access
"""

    if update_config:
        typer.echo(f"Updating AWS config for profile '{profile_name}'...")
        config_file = Path.home() / ".aws" / "config"
        if not config_file.exists():
            config_file.parent.mkdir(parents=True, exist_ok=True)
        config_file.write_text(config_str)
        typer.echo(f"✓ Config updated at {config_file}")

    def login(profile_name: str) -> bool:
        """Login to AWS SSO and return success status"""
        try:
            typer.echo(f"Logging in to AWS SSO for profile '{profile_name}'...")
            result = run(
                ["aws", "sso", "login", "--profile", profile_name],
                capture_output=True,
                text=True,
                check=True,
            )
            if result.stdout:
                typer.echo(result.stdout)
            return True
        except CalledProcessError as e:
            typer.echo(f"Error during SSO login: {e.stderr}", err=True)
            return False
        except Exception as e:
            typer.echo(f"Unexpected error during login: {e}", err=True)
            return False

    # Perform SSO login if requested
    if auto_login:
        if not login(profile_name):
            typer.echo("Failed to login. Exiting.", err=True)
            raise typer.Exit(1)

    # Get credentials
    try:
        typer.echo(f"Retrieving credentials for profile '{profile_name}'...")
        session = boto3.Session(profile_name=profile_name)
        credentials = session.get_credentials()

        if not credentials:
            typer.echo(
                f"No credentials found for profile '{profile_name}'. Please run SSO login first.",
                err=True,
            )
            raise typer.Exit(1)

        # Print the credentials
        typer.echo("✓ Successfully retrieved AWS SSO credentials.")
        typer.echo(f"AWS Access Key ID: {credentials.access_key}")
        typer.echo(
            f"AWS Secret Access Key: {credentials.secret_key[:10]}..."
            if credentials.secret_key
            else "None"
        )
        typer.echo(
            f"AWS Session Token: {credentials.token[:20]}..."
            if credentials.token
            else "None"
        )

        aws_keys = {
            "AWS_ACCESS_KEY_ID": credentials.access_key,
            "AWS_SECRET_ACCESS_KEY": credentials.secret_key,
            "AWS_SESSION_TOKEN": credentials.token,
            "AWS_REGION": region,
            "AWS_PROFILE": profile_name,
        }

        # Set environment variables
        for key, val in aws_keys.items():
            if val:
                os.environ[key] = val

        # Update .env file
        typer.echo(f"Updating {env_file} file...")
        for key, val in aws_keys.items():
            if val:
                set_key(env_file, key, val)

        typer.echo(f"✓ Credentials saved to {env_file} and environment variables")
        typer.echo(f"✓ AWS Region: {region}")
        typer.echo(f"✓ Profile: {profile_name}")

    except Exception as e:
        typer.echo(f"Error retrieving credentials: {e}", err=True)
        raise typer.Exit(1)


@app.command()
def show_config(
    profile_name: str = typer.Option(
        "genai", "--profile", "-p", help="AWS SSO profile name"
    ),
) -> None:
    """Show current AWS configuration for the specified profile."""
    config_file = Path.home() / ".aws" / "config"

    if not config_file.exists():
        typer.echo("AWS config file not found.", err=True)
        raise typer.Exit(1)

    typer.echo(f"AWS Config file location: {config_file}")
    typer.echo("\nCurrent configuration:")
    typer.echo(config_file.read_text())


if __name__ == "__main__":
    app()
