"""
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
import sys
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"
import logging
logging.getLogger("azure").setLevel(logging.WARNING)
logging.getLogger("urllib3").setLevel(logging.WARNING)
def _verify_s3(creds: dict) -> tuple[bool, str]:
    """
    Tries a cheap STS get-caller-identity call to verify AWS creds.
    Returns (True, "") on success or (False, reason) on failure.
    """
    try:
        import boto3
        from botocore.exceptions import ClientError, NoCredentialsError

        session = boto3.Session(
            aws_access_key_id=creds["access_key_id"],
            aws_secret_access_key=creds["secret_access_key"],
            aws_session_token=creds.get("session_token"),
            region_name=creds.get("region", "us-east-1"),
        )
        session.client("sts").get_caller_identity()
        return True, ""
    except ClientError as e:
        code = e.response["Error"]["Code"]
        if code in ("ExpiredTokenException", "ExpiredToken"):
            return False, "Session token is expired."
        if code in ("InvalidClientTokenId", "AuthFailure", "SignatureDoesNotMatch"):
            return False, "Access Key ID or Secret is invalid."
        if code == "NoSuchBucket":
            return False, "Bucket does not exist."
        return False, str(e)
    except NoCredentialsError:
        return False, "No credentials provided."
    except Exception as e:
        return False, str(e)

def _verify_databricks(creds: dict) -> tuple[bool, str]:
    """
    Hits the Databricks /api/2.0/token/list endpoint to verify the PAT token.
    Returns (True, "") on success or (False, reason) on failure.
    """
    try:
        import requests
        host  = creds["host"].rstrip("/")
        token = creds["token"]
        resp  = requests.get(
            f"{host}/api/2.0/workspace/list?path=/",
            headers={"Authorization": f"Bearer {token}"},
            timeout=10,
        )
        if resp.status_code == 200:
            return True, ""
        if resp.status_code == 401:
            return False, "Token is invalid or expired."
        if resp.status_code == 403:
            return False, "Token does not have required permissions."
        return False, f"Unexpected response: {resp.status_code}"
    except Exception as e:
        return False, str(e)


def _verify_azure(creds: dict) -> tuple[bool, str]:
    """
    Lists containers in the Azure Storage Account to verify credentials.
    Returns (True, "") on success or (False, reason) on failure.
    """
    try:
        from azure.storage.blob import BlobServiceClient
        from azure.core.exceptions import ClientAuthenticationError, HttpResponseError

        client = BlobServiceClient(
            account_url=f"https://{creds['account_name']}.blob.core.windows.net",
            credential=creds["account_key"],
        )
        # Use next() without max_results — not supported in all SDK versions
        next(iter(client.list_containers()), None)
        return True, ""
    except ClientAuthenticationError:
        return False, "Account name or key is invalid."
    except HttpResponseError as e:
        return False, str(e)
    except Exception as e:
        return False, str(e)
#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"]

    # Credentials exist — re-verify before handing them back
    creds = config["s3"]
    typer.echo("  Verifying saved AWS credentials...")
    ok, reason = _verify_s3(creds)

    if not ok:
        if "session" in reason.lower() and "expired" in reason.lower():
            invalidate_s3_session_token()
        else:
            invalidate_s3()

    return creds

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"]

    # Credentials exist — re-verify before handing them back
    creds = config["databricks_uc"]
    typer.echo("  Verifying saved Databricks credentials...")
    ok, reason = _verify_databricks(creds)

    if not ok:
        _invalidate_databricks(section="databricks_uc")
    return creds

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"]
    creds = config["azure"]
    typer.echo("  Verifying saved Azure credentials...")
    ok, reason = _verify_azure(creds)

    if not ok:
        invalidate_azure()
    return creds


#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")
    sys.exit(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")
    sys.exit(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")
    sys.exit(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.")
    while True:
        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=False)
        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()
        typer.echo("\n  Verifying credentials with AWS...")
        ok, reason = _verify_s3(creds)
        if ok:
            typer.echo("Credentials verified.\n")
            config["s3"] = creds
            _write(config)
            typer.echo(f"  Credentials saved → {CREDENTIALS_FILE}\n")
            return config
        typer.echo(f"{reason}")
        retry = typer.confirm("  Would you like to re-enter your credentials?", default=True)
        if not retry:
            typer.echo("\nExiting. Credentials were not saved.")
            sys.exit(1)
        typer.echo()

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")
    while True:
        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=False, value_proc=validate_databricks_token)
        creds = {"host": host, "token": token}

        typer.echo("\n  Verifying token with Databricks...")
        ok, reason = _verify_databricks(creds)
        if ok:
            typer.echo(" Token verified.\n")
            config[section] = creds
            _write(config)
            typer.echo(f"  Credentials saved → {CREDENTIALS_FILE}\n")
            return config

        typer.echo(f" {reason}")
        retry = typer.confirm("  Would you like to re-enter your credentials?", default=True)
        if not retry:
            typer.echo("\nExiting. Credentials were not saved.")
            sys.exit(1)
        typer.echo()


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")
    sys.exit(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")
    while True:
        account_name = typer.prompt("Storage Account Name")
        account_key  = typer.prompt("Storage Account Key")

        creds = {"account_name": account_name, "account_key": account_key}

        typer.echo("\n  Verifying credentials with Azure...")
        ok, reason = _verify_azure(creds)
        if ok:
            typer.echo(" Credentials verified.\n")
            config["azure"] = creds
            _write(config)
            typer.echo(f"  Credentials saved → {CREDENTIALS_FILE}\n")
            return config

        typer.echo(f"{reason}")
        retry = typer.confirm("  Would you like to re-enter your credentials?", default=True)
        if not retry:
            typer.echo("\nExiting. Credentials were not saved.")
            sys.exit(1)
        typer.echo()

#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")
