from __future__ import annotations

import atexit
import json
import logging
import os
import tempfile
from pathlib import Path
from typing import Dict, Optional, Set, Union
from urllib.parse import quote_plus

from mnemon.db_connect import (
    ALLOWED_PERSISTENT_AUTH,
    DISALLOWED_SHORT_LIVED_AUTH,
    validate_auth_method,
)

from .exceptions import AuthMaterializationError, ConfigurationError, SecretProviderError
from .secrets import SecretProvider, fetch_secret as _fetch


logger = logging.getLogger(__name__)

ConnectionOutput = Union[str, Dict[str, object]]


# ============================================================
# Temp File Management
# ============================================================

_temp_files: Set[str] = set()


def _write_temp_file(file_data: Dict) -> str:
    """Write file content to temp file, track for cleanup, and return path.

    Args:
        file_data: Dict with "content" key and optional "filename" key.

    Returns:
        Path to the temporary file.
    """
    content = file_data["content"]
    suffix = f".{file_data['filename'].split('.')[-1]}" if "filename" in file_data else ".tmp"

    fd, path = tempfile.mkstemp(suffix=suffix)
    with os.fdopen(fd, 'w') as f:
        f.write(content)

    _temp_files.add(path)
    logger.debug(f"Created temp file: {path}")
    return path


def _cleanup_temp_files():
    """Remove all tracked temp files."""
    global _temp_files
    for path in list(_temp_files):
        try:
            if os.path.exists(path):
                os.remove(path)
                logger.debug(f"Cleaned up temp file: {path}")
        except OSError as e:
            logger.warning(f"Failed to clean up temp file {path}: {e}")
    _temp_files.clear()


atexit.register(_cleanup_temp_files)


# ============================================================
# Helper Functions
# ============================================================

def _default_port(dialect: Optional[str], fallback: str) -> str:
    if not dialect:
        return fallback
    lower = dialect.lower()
    defaults = {"postgresql": "5432", "mysql": "3306", "mssql": "1433"}
    return defaults.get(lower, fallback)


def _fetch_file(provider: SecretProvider, mapping: Dict[str, str], key: str, required: bool = True) -> Optional[str]:
    """Fetch a file path or content, returning a usable file path.

    Handles dual format:
    - If value is a dict with "content" key, writes to temp file and returns path
    - If value is a string, treats it as a file path
    - If value is JSON string containing {"content": ...}, parses and writes to temp
    """
    value = _fetch(provider, mapping, key, required=required)
    if value is None:
        return None

    # Check if it's a JSON string representing file content
    if isinstance(value, str):
        try:
            parsed = json.loads(value)
            if isinstance(parsed, dict) and "content" in parsed:
                return _write_temp_file(parsed)
        except (json.JSONDecodeError, TypeError):
            pass
        # It's a plain path string
        return value

    # It's already a dict with content
    if isinstance(value, dict) and "content" in value:
        return _write_temp_file(value)

    return str(value)


def _load_private_key_bytes(path_or_content: str, passphrase: Optional[str]) -> bytes:
    """Load a private key from path or content string, return serialized bytes."""
    try:
        from cryptography.hazmat.primitives import serialization
    except ImportError as exc:
        raise AuthMaterializationError("cryptography library required for key_pair auth") from exc

    # Check if it's content (starts with PEM header) or a path
    if path_or_content.strip().startswith("-----BEGIN"):
        key_data = path_or_content.encode()
    else:
        key_data = Path(path_or_content).read_bytes()

    password = passphrase.encode() if passphrase else None
    try:
        private_key = serialization.load_pem_private_key(key_data, password=password)
    except Exception as exc:
        raise AuthMaterializationError(f"Unable to load private key: {exc}") from exc
    try:
        return private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption(),
        )
    except Exception as exc:
        raise AuthMaterializationError(f"Unable to serialize private key: {exc}") from exc


# ============================================================
# Auth Method Builders
# ============================================================

def _build_password(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """Password auth for PostgreSQL, MySQL, MariaDB, etc."""
    username = _fetch(provider, mapping, "username")
    password = _fetch(provider, mapping, "password")
    host = _fetch(provider, mapping, "host")
    database = _fetch(provider, mapping, "database")
    dialect = _fetch(provider, mapping, "dialect", required=False) or "postgresql"
    port = _fetch(provider, mapping, "port", required=False) or _default_port(dialect, "5432")
    return f"{dialect}://{quote_plus(username)}:{quote_plus(password)}@{host}:{port}/{quote_plus(database)}"


def _build_ssl_verify(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """SSL-verified connection with password auth and server certificate verification."""
    username = _fetch(provider, mapping, "username")
    password = _fetch(provider, mapping, "password")
    host = _fetch(provider, mapping, "host")
    database = _fetch(provider, mapping, "database")
    dialect = _fetch(provider, mapping, "dialect", required=False) or "postgresql"
    port = _fetch(provider, mapping, "port", required=False) or _default_port(dialect, "5432")
    sslmode = _fetch(provider, mapping, "sslmode", required=False) or "verify-full"

    url = f"{dialect}://{quote_plus(username)}:{quote_plus(password)}@{host}:{port}/{quote_plus(database)}"
    connect_args = {"sslmode": sslmode}

    # Optional CA certificate for verification
    ssl_ca = _fetch_file(provider, mapping, "ssl_ca", required=False)
    if ssl_ca:
        connect_args["sslrootcert"] = ssl_ca

    return {"url": url, "connect_args": connect_args}


def _build_ssl_cert(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """mTLS - client certificate authentication."""
    connection_string = _fetch(provider, mapping, "connection_string")
    ssl_cert_path = _fetch_file(provider, mapping, "ssl_cert_path")
    ssl_key_path = _fetch_file(provider, mapping, "ssl_key_path")
    ssl_ca_path = _fetch_file(provider, mapping, "ssl_ca_path")
    sslmode = _fetch(provider, mapping, "sslmode", required=False) or "verify-full"

    return {
        "url": connection_string,
        "connect_args": {
            "sslcert": ssl_cert_path,
            "sslkey": ssl_key_path,
            "sslrootcert": ssl_ca_path,
            "sslmode": sslmode,
        },
    }


def _build_kerberos(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Kerberos authentication."""
    connection_string = _fetch(provider, mapping, "connection_string")
    kerberos_service_name = _fetch(provider, mapping, "kerberos_service_name")
    kerberos_principal = _fetch(provider, mapping, "kerberos_principal", required=False)
    connect_args = {"krbsrvname": kerberos_service_name}
    if kerberos_principal:
        connect_args["krbprincipal"] = kerberos_principal
    return {"url": connection_string, "connect_args": connect_args}


def _build_scram(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """SCRAM authentication (server decides mechanism)."""
    username = _fetch(provider, mapping, "username")
    password = _fetch(provider, mapping, "password")
    host = _fetch(provider, mapping, "host")
    database = _fetch(provider, mapping, "database")
    dialect = _fetch(provider, mapping, "dialect", required=False) or "postgresql"
    port = _fetch(provider, mapping, "port", required=False) or _default_port(dialect, "5432")
    return f"{dialect}://{quote_plus(username)}:{quote_plus(password)}@{host}:{port}/{quote_plus(database)}?sslmode=require"


def _build_key_pair(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Snowflake key pair authentication."""
    username = _fetch(provider, mapping, "username")
    private_key_value = _fetch(provider, mapping, "private_key_path")
    private_key_passphrase = _fetch(provider, mapping, "private_key_passphrase", required=False)
    account = _fetch(provider, mapping, "account")
    database = _fetch(provider, mapping, "database")
    warehouse = _fetch(provider, mapping, "warehouse", required=False) or ""
    schema = _fetch(provider, mapping, "schema", required=False) or ""
    role = _fetch(provider, mapping, "role", required=False) or ""

    url = f"snowflake://{quote_plus(username)}@{account}/{quote_plus(database)}"
    params = []
    if warehouse:
        params.append(f"warehouse={quote_plus(warehouse)}")
    if schema:
        params.append(f"schema={quote_plus(schema)}")
    if role:
        params.append(f"role={quote_plus(role)}")
    if params:
        url = f"{url}?{'&'.join(params)}"

    private_key_bytes = _load_private_key_bytes(private_key_value, private_key_passphrase)
    return {"url": url, "connect_args": {"private_key": private_key_bytes}}


def _build_service_account(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """BigQuery service account authentication."""
    project_id = _fetch(provider, mapping, "project_id")
    credentials_value = _fetch(provider, mapping, "credentials_path")

    # Handle credentials as path or content
    if credentials_value.strip().startswith("{"):
        # It's JSON content, write to temp file
        creds_path = _write_temp_file({"content": credentials_value, "filename": "creds.json"})
    else:
        creds_path = credentials_value

    url = f"bigquery://{quote_plus(project_id)}"
    dataset = _fetch(provider, mapping, "dataset", required=False)
    if dataset:
        url = f"{url}/{quote_plus(dataset)}"

    location = _fetch(provider, mapping, "location", required=False) or "US"

    return {"url": url, "connect_args": {"credentials_path": creds_path, "location": location}}


def _build_iam_role(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Redshift IAM role authentication."""
    cluster_identifier = _fetch(provider, mapping, "cluster_identifier")
    database = _fetch(provider, mapping, "database")
    db_user = _fetch(provider, mapping, "db_user")
    region = _fetch(provider, mapping, "region")
    access_key_id = _fetch(provider, mapping, "access_key_id")
    secret_access_key = _fetch(provider, mapping, "secret_access_key")
    profile = _fetch(provider, mapping, "profile", required=False)

    try:
        import boto3
    except ImportError as exc:
        raise AuthMaterializationError("boto3 required for iam_role authentication") from exc

    session_kwargs = {"region_name": region}
    if profile:
        session_kwargs["profile_name"] = profile
    session = boto3.Session(**session_kwargs)
    client = session.client(
        "redshift",
        aws_access_key_id=access_key_id,
        aws_secret_access_key=secret_access_key,
    )
    try:
        creds = client.get_cluster_credentials(
            DbUser=db_user,
            ClusterIdentifier=cluster_identifier,
            DbName=database,
            DurationSeconds=900,
        )
    except Exception as exc:
        raise AuthMaterializationError(f"Failed generating Redshift IAM credentials: {exc}") from exc

    password = creds["DbPassword"]
    user = creds.get("DbUser", db_user)
    host = f"{cluster_identifier}.{region}.redshift.amazonaws.com"
    url = f"redshift+redshift_connector://{quote_plus(user)}:{quote_plus(password)}@{host}:5439/{quote_plus(database)}"
    return {"url": url, "connect_args": {}}


def _build_iam_credentials(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """AWS Athena IAM credentials authentication."""
    region = _fetch(provider, mapping, "region")
    s3_staging_dir = _fetch(provider, mapping, "s3_staging_dir")
    access_key_id = _fetch(provider, mapping, "access_key_id")
    secret_access_key = _fetch(provider, mapping, "secret_access_key")
    database = _fetch(provider, mapping, "database", required=False) or ""
    catalog = _fetch(provider, mapping, "catalog", required=False)
    work_group = _fetch(provider, mapping, "work_group", required=False)

    url = (
        f"awsathena+rest://{quote_plus(access_key_id)}:{quote_plus(secret_access_key)}"
        f"@athena.{region}.amazonaws.com:443/{quote_plus(database)}"
        f"?s3_staging_dir={quote_plus(s3_staging_dir)}"
    )
    if catalog:
        url += f"&catalog_name={quote_plus(catalog)}"
    if work_group:
        url += f"&work_group={quote_plus(work_group)}"

    return url


def _build_ssl(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """Basic SSL encryption (deprecated, prefer ssl_verify)."""
    username = _fetch(provider, mapping, "username")
    password = _fetch(provider, mapping, "password")
    host = _fetch(provider, mapping, "host")
    database = _fetch(provider, mapping, "database")
    dialect = _fetch(provider, mapping, "dialect", required=False) or "postgresql"
    port = _fetch(provider, mapping, "port", required=False) or _default_port(dialect, "")
    port_part = f":{port}" if port else ""
    if dialect.startswith("mssql"):
        suffix = "encrypt=yes"
    else:
        suffix = "ssl=true"
    return f"{dialect}://{quote_plus(username)}:{quote_plus(password)}@{host}{port_part}/{quote_plus(database)}?{suffix}"


def _build_windows_auth(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """SQL Server Windows (Integrated) authentication."""
    host = _fetch(provider, mapping, "host")
    database = _fetch(provider, mapping, "database")
    port = _fetch(provider, mapping, "port", required=False) or "1433"
    driver = _fetch(provider, mapping, "driver", required=False) or "ODBC Driver 18 for SQL Server"
    driver_enc = quote_plus(driver)
    trust = "yes" if _fetch(provider, mapping, "trust_server_certificate", required=False) else "no"
    return f"mssql+pyodbc://@{host}:{port}/{quote_plus(database)}?trusted_connection=yes&driver={driver_enc}&TrustServerCertificate={trust}"


def _build_token(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """Databricks Personal Access Token authentication."""
    host = _fetch(provider, mapping, "host")
    http_path = _fetch(provider, mapping, "http_path")
    token = _fetch(provider, mapping, "token")
    catalog = _fetch(provider, mapping, "catalog", required=False) or ""
    schema = _fetch(provider, mapping, "schema", required=False) or ""
    params = [f"http_path={quote_plus(http_path)}"]
    if catalog:
        params.append(f"catalog={quote_plus(catalog)}")
    if schema:
        params.append(f"schema={quote_plus(schema)}")
    return f"databricks://token:{quote_plus(token)}@{host}?{'&'.join(params)}"


def _build_oauth_m2m(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Databricks OAuth machine-to-machine (service principal) authentication."""
    host = _fetch(provider, mapping, "host")
    http_path = _fetch(provider, mapping, "http_path")
    client_id = _fetch(provider, mapping, "client_id")
    client_secret = _fetch(provider, mapping, "client_secret")
    catalog = _fetch(provider, mapping, "catalog", required=False)
    schema = _fetch(provider, mapping, "schema", required=False)

    # This requires the databricks-sql-connector with OAuth support
    try:
        from mnemon.utils.databricks_engine import build_databricks_connection
    except ImportError as exc:
        raise AuthMaterializationError("mnemon.utils.databricks_engine required for oauth_m2m") from exc

    config = {
        "host": host,
        "http_path": http_path,
        "client_id": client_id,
        "client_secret": client_secret,
    }
    if catalog:
        config["catalog"] = catalog
    if schema:
        config["schema"] = schema

    conn = build_databricks_connection("oauth_m2m", config)

    # Return a special marker that DBClient can recognize
    return {"_databricks_connection": conn, "url": None, "connect_args": {}}


def _build_jwt(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Trino JWT authentication."""
    host = _fetch(provider, mapping, "host")
    jwt_token = _fetch(provider, mapping, "jwt_token")
    port = _fetch(provider, mapping, "port", required=False) or "443"
    catalog = _fetch(provider, mapping, "catalog", required=False) or ""
    schema = _fetch(provider, mapping, "schema", required=False) or ""

    try:
        from trino.auth import JWTAuthentication
    except ImportError as exc:
        raise AuthMaterializationError("trino library required for jwt auth") from exc

    url = f"trino://{host}:{port}/{catalog}/{schema}"
    return {
        "url": url,
        "connect_args": {
            "http_scheme": "https",
            "auth": JWTAuthentication(jwt_token),
        },
    }


def _build_certificate(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Trino certificate authentication."""
    host = _fetch(provider, mapping, "host")
    cert_path = _fetch_file(provider, mapping, "cert_path")
    key_path = _fetch_file(provider, mapping, "key_path")
    port = _fetch(provider, mapping, "port", required=False) or "443"
    catalog = _fetch(provider, mapping, "catalog", required=False) or ""
    schema = _fetch(provider, mapping, "schema", required=False) or ""

    try:
        from trino.auth import CertificateAuthentication
    except ImportError as exc:
        raise AuthMaterializationError("trino library required for certificate auth") from exc

    url = f"trino://{host}:{port}/{catalog}/{schema}"
    return {
        "url": url,
        "connect_args": {
            "http_scheme": "https",
            "auth": CertificateAuthentication(cert_path, key_path),
        },
    }


def _build_wallet(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """Oracle wallet authentication."""
    connection_string = _fetch(provider, mapping, "connection_string")
    wallet_location = _fetch(provider, mapping, "wallet_location")
    wallet_password = _fetch(provider, mapping, "wallet_password", required=False)
    connect_args = {
        "config_dir": wallet_location,
        "wallet_location": wallet_location,
    }
    if wallet_password:
        connect_args["wallet_password"] = wallet_password
    return {"url": f"oracle+oracledb://{connection_string}", "connect_args": connect_args}


def _build_ldap(mapping: Dict[str, str], provider: SecretProvider) -> ConnectionOutput:
    """LDAP authentication for Teradata, Vertica, SparkSQL, Db2."""
    username = _fetch(provider, mapping, "username")
    password = _fetch(provider, mapping, "password")
    host = _fetch(provider, mapping, "host")
    database = _fetch(provider, mapping, "database")
    port = _fetch(provider, mapping, "port", required=False) or ""
    dialect = _fetch(provider, mapping, "dialect", required=False) or "teradatasql"
    port_part = f":{port}" if port else ""
    url = f"{dialect}://{quote_plus(username)}:{quote_plus(password)}@{host}{port_part}/{quote_plus(database)}"

    connect_args = {"logmech": "LDAP"}
    ldap_server = _fetch(provider, mapping, "ldap_server", required=False)
    if ldap_server:
        connect_args["ldap_server"] = ldap_server

    return {"url": url, "connect_args": connect_args}


def _build_local_file(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """DuckDB/SQLite local file authentication."""
    file_path = _fetch(provider, mapping, "file_path")
    dialect = "duckdb" if file_path.endswith(".duckdb") else "sqlite"
    return f"{dialect}:///{file_path}"


def _build_motherduck(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """MotherDuck cloud DuckDB authentication."""
    token = _fetch(provider, mapping, "token")
    database = _fetch(provider, mapping, "database", required=False) or "my_db"
    return f"duckdb:///md:{quote_plus(database)}?motherduck_token={quote_plus(token)}"


def _build_none(mapping: Dict[str, str], provider: SecretProvider) -> str:
    """No authentication (for dev/test or trusted network)."""
    host = _fetch(provider, mapping, "host")
    port = _fetch(provider, mapping, "port", required=False) or ""
    catalog = _fetch(provider, mapping, "catalog", required=False) or ""
    schema = _fetch(provider, mapping, "schema", required=False) or ""
    dialect = _fetch(provider, mapping, "dialect", required=False) or "trino"
    user = _fetch(provider, mapping, "user", required=False) or ""
    port_part = f":{port}" if port else ""
    path = "/".join([p for p in [catalog, schema] if p])
    path = f"/{path}" if path else ""
    user_part = f"{quote_plus(user)}@" if user else ""
    return f"{dialect}://{user_part}{host}{port_part}{path}"


# ============================================================
# Builder Registry
# ============================================================

DATABASE_BUILDERS = {
    "password": _build_password,
    "ssl_verify": _build_ssl_verify,
    "ssl_cert": _build_ssl_cert,
    "kerberos": _build_kerberos,
    "scram": _build_scram,
    "key_pair": _build_key_pair,
    "service_account": _build_service_account,
    "iam_role": _build_iam_role,
    "iam_credentials": _build_iam_credentials,
    "ssl": _build_ssl,
    "windows_auth": _build_windows_auth,
    "token": _build_token,
    "oauth_m2m": _build_oauth_m2m,
    "jwt": _build_jwt,
    "certificate": _build_certificate,
    "wallet": _build_wallet,
    "ldap": _build_ldap,
    "local_file": _build_local_file,
    "motherduck": _build_motherduck,
    "none": _build_none,
}


# ============================================================
# Main Entry Point
# ============================================================

def build_database_connection(
    auth_type: str,
    mapping: Dict[str, str],
    provider: SecretProvider,
    database_type: Optional[str] = None,
) -> ConnectionOutput:
    """Build a database connection from auth configuration.

    Args:
        auth_type: The authentication method (e.g., "password", "ssl_cert", "key_pair")
        mapping: Dictionary mapping semantic keys to secret names
        provider: SecretProvider instance to fetch secret values
        database_type: Optional database type for validation (e.g., "postgresql", "snowflake")

    Returns:
        Either a connection URL string, or a dict with "url" and "connect_args" keys.

    Raises:
        ConfigurationError: If auth_type is not recognized
        AuthMaterializationError: If credentials cannot be built
        SecretProviderError: If secrets cannot be fetched
    """
    # Validate auth method if database_type is provided
    if database_type:
        try:
            validate_auth_method(database_type, auth_type)
        except ValueError as exc:
            raise ConfigurationError(str(exc)) from exc

    if auth_type not in DATABASE_BUILDERS:
        allowed_methods = list(DATABASE_BUILDERS.keys())
        raise ConfigurationError(
            f"Unsupported database auth_type '{auth_type}'. "
            f"Supported methods: {allowed_methods}"
        )

    builder = DATABASE_BUILDERS[auth_type]
    try:
        return builder(mapping, provider)
    except (ConfigurationError, SecretProviderError, AuthMaterializationError):
        raise
    except Exception as exc:
        raise AuthMaterializationError(f"Failed to materialize database auth '{auth_type}': {exc}") from exc
