"""
Stores credentials in ~/.acc/credentials.toml.
- First run  : prompts user, saves to file
- Later runs : reads from file
- Expired    : prompts for new token, updates file

File location: ~/.acc/credentials.toml
"""

import tomllib
import typer
from pathlib import Path
from acc_cli.utils.validators import (
    validate_aws_access_key_id,
    validate_aws_secret_access_key,
    validate_aws_region,
    validate_databricks_host,
    validate_databricks_token,
)
CREDENTIALS_FILE = Path.home() / ".acc" / "credentials.toml"
#GET CREDENTIALS
def get_s3_credentials() -> dict:
    """
    Returns S3 credentials from [s3] section of ~/.acc/credentials.toml.

    Returns:
        dict with keys: access_key_id, secret_access_key, region
        optional key  : session_token
    """
    config = _read()
    if "s3" not in config:
        typer.echo("\nFirst time using S3 as storage,Enter your credentials.\n")
        config = _collect_s3(config)
    return config["s3"]

def get_databricks_uc_credentials() -> dict:
    """
    Returns Databricks Unity Catalog credentials from [databricks_uc] section.

    Returns:
        dict with keys: host, token
    """
    config = _read()
    if "databricks_uc" not in config:
        typer.echo("\nFirst time using Unity Catalog as storage,Enter your credentials.\n")
        config = _collect_databricks(config, section="databricks_uc")
    return config["databricks_uc"]

def get_azure_credentials() -> dict:
    """
    Returns Azure ADLS credentials from [azure] section of ~/.acc/credentials.toml.
    Returns:
        dict with keys: account_name, account_key
    """
    config = _read()
    if "azure" not in config:
        typer.echo("\nFirst time using Azure ADLS — let's save your credentials.\n")
        config = _collect_azure(config)
    return config["azure"]


#INVALIDATE
def invalidate_s3():
    """
    Called when S3 upload fails due to expired/invalid access key or secret.

    Clears the entire [s3] section and re-prompts for all credentials
    """
    typer.echo("\n[!] Your AWS credentials are expired or invalid.\n")
    update = typer.confirm("Would you like to enter new credentials now?", default=True)

    if not update:
        typer.echo(f"\nManually update them in:\n  {CREDENTIALS_FILE}")
        raise typer.Exit(code=1)

    config = _read()
    config.pop("s3", None)
    _collect_s3(config)

    typer.echo("\nNew credentials saved. Please run 'acc build' again.\n")
    raise typer.Exit(code=0)


def invalidate_s3_session_token():
    """
    Called when S3 upload fails specifically due to an expired STS session token.

    Only clears and re-prompts for session_token — the access_key_id and
    secret_access_key are kept since they are still valid.

    """
    typer.echo("\n[!] Your AWS session token has expired.\n")

    update = typer.confirm("Would you like to enter a new session token now?", default=True)

    if not update:
        typer.echo(f"\nManually update session_token in:\n  {CREDENTIALS_FILE}")
        raise typer.Exit(code=1)

    config = _read()
    new_token = typer.prompt("New AWS Session Token", hide_input=True)

    # update only the session_token, keep everything else
    config.setdefault("s3", {})["session_token"] = new_token
    _write(config)

    typer.echo(f"\n  Session token updated → {CREDENTIALS_FILE}\n")
    typer.echo("New session token saved. Please run 'acc build' again.\n")
    raise typer.Exit(code=0)

def _invalidate_databricks(section: str):
    """
    Shared invalidation logic for any Databricks section.

    Prefills the existing host so the user only needs to paste a new token.
    """
    typer.echo("\n[!] Your Databricks PAT token is expired or invalid.")
    typer.echo("    Generate a new one at:")
    typer.echo("    Workspace → Settings → Developer → Access tokens\n")

    update = typer.confirm("Would you like to enter a new PAT token now?", default=True)

    if not update:
        typer.echo(f"\nManually update it in:\n  {CREDENTIALS_FILE}")
        raise typer.Exit(code=1)

    config        = _read()
    existing_host = config.get(section, {}).get("host", "")
    config.pop(section, None)
    _collect_databricks(config, section=section, prefill_host=existing_host)

    typer.echo("\nNew PAT token saved. Please run 'acc build' again.\n")
    raise typer.Exit(code=0)

def invalidate_databricks_uc():
    """
    Called when UC upload fails due to expired/invalid PAT token.
    """
    _invalidate_databricks(section="databricks_uc")

#FIRST-TIME COLLECTION

def _collect_s3(config: dict) -> dict:
    """
    Prompts for AWS credentials and saves under [s3].
    """
    typer.echo("AWS Access Key ID and Secret are your permanent IAM credentials.")

    access_key_id     = typer.prompt("AWS Access Key ID", value_proc=validate_aws_access_key_id)
    secret_access_key = typer.prompt("AWS Secret Access Key", hide_input=True, value_proc=validate_aws_secret_access_key)
    region            = typer.prompt("AWS Region", default="us-east-1", value_proc=validate_aws_region)
    session_token     = typer.prompt(
                            "AWS Session Token (optional — press Enter to skip)",
                            default="",
                            hide_input=True,
                        )
    creds: dict = {
        "access_key_id"     : access_key_id,
        "secret_access_key" : secret_access_key,
        "region"            : region,
    }

    # only store session_token if user actually provided one
    if session_token.strip():
        creds["session_token"] = session_token.strip()

    config["s3"] = creds
    _write(config)
    typer.echo(f"\n  Credentials saved → {CREDENTIALS_FILE}\n")
    return config


def _collect_databricks(config: dict, section: str, prefill_host: str = "") -> dict:
    """
    Prompts for Databricks host + PAT token and saves under the given section.
    """
    typer.echo("Generate a PAT token at:")
    typer.echo("  Workspace → Settings → Developer → Access tokens\n")

    config[section] = {
        "host"  : typer.prompt(
                      "Databricks Host",
                      default=prefill_host or "https://dbc-xxx.cloud.databricks.com",
                      value_proc=validate_databricks_host
                  ),
        "token" : typer.prompt("Databricks PAT Token", hide_input=True, value_proc=validate_databricks_token),
    }
    _write(config)
    typer.echo(f"\n  Credentials saved → {CREDENTIALS_FILE}\n")
    return config


def invalidate_azure():
    """
    Called when Azure upload fails due to expired/invalid credentials.
    """
    typer.echo("\n[!] Your Azure Storage credentials are expired or invalid.\n")
    update = typer.confirm("Would you like to enter new credentials now?", default=True)
    if not update:
        typer.echo(f"\nManually update them in:\n  {CREDENTIALS_FILE}")
        raise typer.Exit(code=1)
    config = _read()
    config.pop("azure", None)
    _collect_azure(config)
    typer.echo("\nNew credentials saved. Please run 'acc build' again.\n")
    raise typer.Exit(code=0)


def _collect_azure(config: dict) -> dict:
    """
    Prompts for Azure Storage Account credentials and saves under [azure].
    """
    typer.echo("Find your Storage Account Key at:")
    typer.echo("  Azure Portal → Storage Account → Security + networking → Access keys\n")

    account_name = typer.prompt("Storage Account Name")
    account_key  = typer.prompt("Storage Account Key", hide_input=True)

    config["azure"] = {
        "account_name": account_name,
        "account_key":  account_key,
    }
    _write(config)
    typer.echo(f"\n  Credentials saved → {CREDENTIALS_FILE}\n")
    return config

#FILE I/O
def _read() -> dict:
    """
    Reads ~/.acc/credentials.toml.
    Returns empty dict on first run (file doesn't exist yet).
    """
    if not CREDENTIALS_FILE.exists():
        return {}
    with open(CREDENTIALS_FILE, "rb") as f:
        return tomllib.load(f)

def _write(config: dict):
    """
    Writes the full config back to ~/.acc/credentials.toml.
    Creates ~/.acc/ directory if it doesn't exist.
    """
    CREDENTIALS_FILE.parent.mkdir(parents=True, exist_ok=True)

    lines = []
    for section, values in config.items():
        lines.append(f"[{section}]")
        for key, value in values.items():
            lines.append(f'{key} = "{value}"')
        lines.append("")  # blank line between sections

    CREDENTIALS_FILE.write_text("\n".join(lines), encoding="utf-8")
