"""
SchemaDBConnection - Unified SQLAlchemy connector for multiple database types.
Supports 17 databases with persistent authentication methods only.

Changes:
- Added MariaDB support
- Renamed ssl -> ssl_verify for clarity
- Added ssl_verify auth method for PostgreSQL
- Added ldap auth for Vertica, Db2, SparkSQL
- Commented out Kerberos (deferred to future release)
- Standardized auth method names across databases
"""

import tempfile
import os
import atexit
import logging
from typing import Optional, Dict, Any, List, Tuple, Set

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine
from sqlalchemy.pool import QueuePool, StaticPool, NullPool

from .utils.error_handling import normalize_error, to_api_response

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


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


# ============================================================
# Auth Method Validation
# ============================================================

DISALLOWED_SHORT_LIVED_AUTH = {
    "snowflake": {
        "externalbrowser": "Browser-based SSO provides short-lived tokens",
        "oauth_browser": "OAuth browser flow provides short-lived tokens",
        "okta": "Okta SSO provides short-lived tokens",
        "oauth": "OAuth provides short-lived tokens unless using refresh tokens",
    },
    "redshift": {
        "iam_temporary": "IAM GetClusterCredentials provides 15min-1hr tokens",
        "iam_identity_center": "Identity Center provides short-lived tokens",
    },
    "sqlserver": {
        "azure_ad_interactive": "Interactive Azure AD provides short-lived tokens",
        "azure_ad_token": "Azure AD tokens are typically short-lived",
    },
    "databricks": {
        "oauth_u2m": "User-to-machine OAuth provides short-lived tokens",
        "oauth_browser": "Browser OAuth provides short-lived tokens",
    },
    "athena": {
        "iam_session": "Session tokens are short-lived",
    },
    "bigquery": {
        "oauth_browser": "Browser OAuth provides short-lived tokens",
        "adc_user": "Application Default Credentials (user) may be short-lived",
    },
    "trino": {
        "oauth_browser": "Browser OAuth provides short-lived tokens",
    },
}

# Allowed persistent auth methods per database
# NOTE: Kerberos commented out - deferred to future release
ALLOWED_PERSISTENT_AUTH = {
    "postgresql": ["password", "scram", "ssl_verify", "ssl_cert"],  # kerberos deferred
    "mysql": ["password", "ssl_verify", "ssl_cert"],  # kerberos deferred
    "mariadb": ["password", "ssl_verify", "ssl_cert"],  # kerberos deferred
    "clickhouse": ["password", "ssl_verify", "ssl_cert"],
    "sqlserver": ["password"],  # windows_auth, kerberos deferred
    "trino": ["none", "password", "jwt", "certificate"],  # kerberos deferred
    "sparksql": ["none", "password", "ldap"],  # kerberos deferred
    "vertica": ["password", "ldap", "ssl_verify"],  # kerberos deferred
    "oracle": ["password", "wallet"],  # kerberos deferred
    "teradata": ["password", "ldap"],  # kerberos deferred
    "db2": ["password", "ldap"],  # kerberos deferred
    "snowflake": ["password", "key_pair"],
    "bigquery": ["service_account"],
    "redshift": ["password", "iam_role"],
    "duckdb": ["local_file", "motherduck"],
    "databricks": ["token", "oauth_m2m"],
    "athena": ["iam_credentials"],
}


def validate_auth_method(database_type: str, auth_method: str) -> None:
    """Validate that the auth method provides persistent credentials."""
    disallowed = DISALLOWED_SHORT_LIVED_AUTH.get(database_type, {})
    if auth_method in disallowed:
        raise ValueError(
            f"Auth method '{auth_method}' is not allowed for persistent connections: "
            f"{disallowed[auth_method]}. "
            f"Allowed methods for {database_type}: {ALLOWED_PERSISTENT_AUTH.get(database_type, [])}"
        )
    
    allowed = ALLOWED_PERSISTENT_AUTH.get(database_type, [])
    if allowed and auth_method not in allowed:
        raise ValueError(
            f"Auth method '{auth_method}' is not recognized for {database_type}. "
            f"Allowed persistent auth methods: {allowed}"
        )


# ============================================================
# Standard Pool Configuration
# ============================================================

def get_pool_config(database_type: str) -> Dict[str, Any]:
    """Get standardized pool configuration for a database type."""
    NO_POOL_DATABASES = {"bigquery"}
    
    if database_type in NO_POOL_DATABASES:
        return {"poolclass": NullPool}
    
    return {
        "poolclass": QueuePool,
        "pool_size": 5,
        "max_overflow": 10,
        "pool_timeout": 30,
        "pool_recycle": 1800,
        "pool_pre_ping": True,
    }


# ============================================================
# Main Connection Class
# ============================================================

class SchemaDBConnection:
    """Manages database connections via SQLAlchemy for schema analysis."""

    def __init__(self, config: Dict[str, Any]):
        self.database_type = config["database_type"]
        self.auth_method = config["auth_method"]
        self.connection_config = config["connection"]
        self._engine: Optional[Engine] = None
        
        validate_auth_method(self.database_type, self.auth_method)
        logger.info(f"Initialized SchemaDBConnection for {self.database_type} ({self.auth_method})")

    def connect(self) -> Engine:
        """Create and return SQLAlchemy engine."""
        if self._engine is not None:
            return self._engine

        builder = ENGINE_BUILDERS.get(self.database_type)
        if not builder:
            raise ValueError(f"Unsupported database type: {self.database_type}")

        try:
            self._engine = builder(self.auth_method, self.connection_config)
            logger.info(f"Engine created for {self.database_type}")
            return self._engine
        except Exception as e:
            logger.error(f"Failed to create engine: {e}")
            raise

    def safe_connect(self) -> Tuple[bool, Optional[Engine], Optional[dict]]:
        """Attempt connection with normalized error handling."""
        try:
            engine = self.connect()
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            return (True, engine, None)
        except Exception as e:
            normalized = normalize_error(self.database_type, e)
            return (False, None, to_api_response(normalized))

    def execute_query(self, query: str, params: Optional[Dict] = None) -> List[Dict]:
        """Execute a query and return results as list of dicts."""
        engine = self.connect()
        with engine.connect() as conn:
            result = conn.execute(text(query), params or {})
            if result.returns_rows:
                columns = result.keys()
                return [dict(zip(columns, row)) for row in result.fetchall()]
            else:
                conn.commit()
                return []

    def close_all(self):
        """Dispose of the engine and clean up temp files."""
        if self._engine:
            self._engine.dispose()
            self._engine = None
        _cleanup_temp_files()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close_all()


# ============================================================
# Engine Builders
# ============================================================

def _build_postgresql_engine(auth_method: str, config: Dict) -> Engine:
    """
    PostgreSQL engine.
    Auth: password, scram, ssl_verify, ssl_cert
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config["user"])
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    connect_args = {}
    
    if auth_method == "ssl_cert":
        # mTLS - client certificate authentication
        base_url = f"postgresql+psycopg2://{user}@{host}:{port}/{database}"
        connect_args = {
            "sslmode": config.get("sslmode", "verify-full"),
            "sslcert": _write_temp_file(config["sslcert"]),
            "sslkey": _write_temp_file(config["sslkey"]),
        }
        if config.get("sslrootcert"):
            connect_args["sslrootcert"] = _write_temp_file(config["sslrootcert"])
    elif auth_method == "ssl_verify":
        # Encrypted connection with server cert verification
        password = quote_plus(config["password"])
        base_url = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
        connect_args = {
            "sslmode": config.get("sslmode", "verify-full"),
        }
        if config.get("sslrootcert"):
            connect_args["sslrootcert"] = _write_temp_file(config["sslrootcert"])
    else:  # password, scram (server decides auth mechanism)
        password = quote_plus(config["password"])
        base_url = f"postgresql+psycopg2://{user}:{password}@{host}:{port}/{database}"
        if config.get("sslmode"):
            connect_args["sslmode"] = config["sslmode"]

    pool_config = get_pool_config("postgresql")
    return create_engine(base_url, connect_args=connect_args, **pool_config)


def _build_mysql_engine(auth_method: str, config: Dict) -> Engine:
    """
    MySQL engine.
    Auth: password, ssl_verify, ssl_cert
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config["user"])
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    password = quote_plus(config["password"])
    connect_args = {}
    
    base_url = f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}"
    
    if auth_method == "ssl_cert":
        # mTLS - client certificate required
        ssl_config = {
            "cert": _write_temp_file(config["ssl_cert"]),
            "key": _write_temp_file(config["ssl_key"]),
        }
        if config.get("ssl_ca"):
            ssl_config["ca"] = _write_temp_file(config["ssl_ca"])
        connect_args["ssl"] = ssl_config
    elif auth_method == "ssl_verify":
        # Encrypted connection with server cert verification
        ssl_config = {}
        if config.get("ssl_ca"):
            ssl_config["ca"] = _write_temp_file(config["ssl_ca"])
        ssl_config["check_hostname"] = config.get("check_hostname", True)
        connect_args["ssl"] = ssl_config

    pool_config = get_pool_config("mysql")
    return create_engine(base_url, connect_args=connect_args, **pool_config)


def _build_mariadb_engine(auth_method: str, config: Dict) -> Engine:
    """
    MariaDB engine - uses same driver as MySQL.
    Auth: password, ssl_verify, ssl_cert
    """
    return _build_mysql_engine(auth_method, config)


def _build_clickhouse_engine(auth_method: str, config: Dict) -> Engine:
    """
    ClickHouse engine.
    Auth: password, ssl_verify, ssl_cert
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config.get("user", "default"))
    password = quote_plus(config.get("password", ""))
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    
    base_url = f"clickhouse+native://{user}:{password}@{host}:{port}/{database}"
    connect_args = {}
    
    if auth_method == "ssl_cert":
        # mTLS - X.509 certificate authentication
        connect_args = {
            "secure": True,
            "verify": config.get("verify", True),
            "client_cert": _write_temp_file(config["ssl_cert"]),
            "client_key": _write_temp_file(config["ssl_key"]),
        }
        if config.get("ssl_ca"):
            connect_args["ca_certs"] = _write_temp_file(config["ssl_ca"])
    elif auth_method == "ssl_verify":
        # Encrypted connection with server cert verification
        connect_args = {
            "secure": True,
            "verify": True,
        }
        if config.get("ssl_ca"):
            connect_args["ca_certs"] = _write_temp_file(config["ssl_ca"])

    pool_config = get_pool_config("clickhouse")
    return create_engine(base_url, connect_args=connect_args, **pool_config)


def _build_sqlserver_engine(auth_method: str, config: Dict) -> Engine:
    """
    SQL Server engine.
    Auth: password
    Note: windows_auth and kerberos deferred (requires AD)
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config["user"])
    password = quote_plus(config["password"])
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    driver = quote_plus(config.get("driver", "ODBC Driver 18 for SQL Server"))
    trust = "yes" if config.get("trust_server_certificate", True) else "no"
    
    url = (
        f"mssql+pyodbc://{user}:{password}@{host}:{port}/{database}"
        f"?driver={driver}&TrustServerCertificate={trust}"
    )

    pool_config = get_pool_config("sqlserver")
    return create_engine(url, **pool_config)


def _build_trino_engine(auth_method: str, config: Dict) -> Engine:
    """
    Trino engine.
    Auth: none, password, jwt, certificate
    Note: kerberos deferred
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config["user"])
    host = config["host"]
    port = config["port"]
    catalog = quote_plus(config["catalog"])
    schema = quote_plus(config.get("schema", ""))

    base_url = f"trino://{user}@{host}:{port}/{catalog}/{schema}"
    connect_args = {}

    if auth_method == "none":
        pass
    elif auth_method == "password":
        from trino.auth import BasicAuthentication
        connect_args = {
            "http_scheme": "https",
            "auth": BasicAuthentication(config["user"], config["password"]),
        }
    elif auth_method == "jwt":
        from trino.auth import JWTAuthentication
        connect_args = {
            "http_scheme": "https",
            "auth": JWTAuthentication(config["jwt_token"]),
        }
    elif auth_method == "certificate":
        from trino.auth import CertificateAuthentication
        cert_path = _write_temp_file(config["certificate"])
        key_path = _write_temp_file(config["private_key"])
        connect_args = {
            "http_scheme": "https",
            "auth": CertificateAuthentication(cert_path, key_path),
        }

    pool_config = get_pool_config("trino")
    return create_engine(base_url, connect_args=connect_args, **pool_config)


def _build_sparksql_engine(auth_method: str, config: Dict) -> Engine:
    """
    SparkSQL / Hive ThriftServer engine.
    Auth: none, password, ldap
    Note: kerberos deferred
    """
    from urllib.parse import quote_plus
    
    host = config["host"]
    port = config["port"]
    database = quote_plus(config.get("database", "default"))
    connect_args = {}

    if auth_method in ("password", "ldap"):
        user = quote_plus(config["user"])
        password = quote_plus(config["password"])
        url = f"hive://{user}:{password}@{host}:{port}/{database}"
        if auth_method == "ldap":
            connect_args["auth"] = "LDAP"
    else:  # none
        user = quote_plus(config.get("user", "hive"))
        url = f"hive://{user}@{host}:{port}/{database}"

    pool_config = get_pool_config("sparksql")
    return create_engine(url, connect_args=connect_args, **pool_config)


def _build_vertica_engine(auth_method: str, config: Dict) -> Engine:
    """
    Vertica engine.
    Auth: password, ldap, ssl_verify
    Note: kerberos deferred
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config["user"])
    password = quote_plus(config["password"])
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    connect_args = {}
    
    base_url = f"vertica+vertica_python://{user}:{password}@{host}:{port}/{database}"
    
    if auth_method == "ssl_verify":
        connect_args = {
            "ssl": True,
        }
        if config.get("ssl_ca"):
            connect_args["ssl_ca_file"] = _write_temp_file(config["ssl_ca"])
    # ldap and password use same connection string, server handles auth

    pool_config = get_pool_config("vertica")
    return create_engine(base_url, connect_args=connect_args, **pool_config)


def _build_oracle_engine(auth_method: str, config: Dict) -> Engine:
    """
    Oracle engine.
    Auth: password, wallet
    Note: kerberos deferred
    """
    from urllib.parse import quote_plus
    
    host = config["host"]
    port = config["port"]
    service = quote_plus(config["service_name"])
    connect_args = {}
    
    if auth_method == "wallet":
        # Wallet contains credentials - no user/pass in URL
        url = f"oracle+oracledb://@{host}:{port}/?service_name={service}"
        connect_args = {
            "config_dir": config["wallet_location"],
            "wallet_location": config["wallet_location"],
        }
        if config.get("wallet_password"):
            connect_args["wallet_password"] = config["wallet_password"]
    else:  # password
        user = quote_plus(config["user"])
        password = quote_plus(config["password"])
        url = f"oracle+oracledb://{user}:{password}@{host}:{port}/?service_name={service}"

    pool_config = get_pool_config("oracle")
    return create_engine(url, connect_args=connect_args, **pool_config)


def _build_teradata_engine(auth_method: str, config: Dict) -> Engine:
    """
    Teradata engine.
    Auth: password, ldap
    Note: kerberos deferred
    """
    from urllib.parse import quote_plus
    
    host = config["host"]
    port = config.get("port", 1025)
    user = quote_plus(config["user"])
    password = quote_plus(config["password"])
    database = quote_plus(config.get("database", ""))
    
    url = f"teradatasql://{user}:{password}@{host}:{port}/{database}"
    connect_args = {
        "logmech": "LDAP" if auth_method == "ldap" else "TD2"
    }

    pool_config = get_pool_config("teradata")
    return create_engine(url, connect_args=connect_args, **pool_config)


def _build_db2_engine(auth_method: str, config: Dict) -> Engine:
    """
    IBM Db2 engine.
    Auth: password, ldap
    Note: kerberos deferred
    """
    from urllib.parse import quote_plus
    
    user = quote_plus(config["user"])
    password = quote_plus(config["password"])
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    
    url = f"db2+ibm_db://{user}:{password}@{host}:{port}/{database}"
    connect_args = {}
    
    # LDAP auth is typically configured server-side; client just sends credentials
    # No special connect_args needed for LDAP vs password

    pool_config = get_pool_config("db2")
    return create_engine(url, connect_args=connect_args, **pool_config)


def _build_snowflake_engine(auth_method: str, config: Dict) -> Engine:
    """
    Snowflake engine.
    Auth: password, key_pair
    """
    from snowflake.sqlalchemy import URL
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization

    connect_args = {}

    if auth_method == "key_pair":
        private_key_content = config["private_key"]["content"]
        passphrase = config.get("private_key_passphrase")
        
        p_key = serialization.load_pem_private_key(
            private_key_content.encode(),
            password=passphrase.encode() if passphrase else None,
            backend=default_backend()
        )
        private_key_bytes = p_key.private_bytes(
            encoding=serialization.Encoding.DER,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        connect_args["private_key"] = private_key_bytes

        url = URL(
            account=config["account"],
            user=config["user"],
            warehouse=config["warehouse"],
            database=config["database"],
            schema=config.get("schema", "PUBLIC"),
            **({"role": config["role"]} if config.get("role") else {}),
        )
    else:  # password
        url = URL(
            account=config["account"],
            user=config["user"],
            password=config["password"],
            warehouse=config["warehouse"],
            database=config["database"],
            schema=config.get("schema", "PUBLIC"),
            **({"role": config["role"]} if config.get("role") else {}),
        )

    pool_config = get_pool_config("snowflake")
    return create_engine(url, connect_args=connect_args, **pool_config)


def _build_bigquery_engine(auth_method: str, config: Dict) -> Engine:
    """
    BigQuery engine.
    Auth: service_account
    """
    from urllib.parse import quote_plus
    import json

    creds_content = config["credentials_json"]["content"]
    creds_dict = json.loads(creds_content)
    creds_path = _write_temp_file({"content": json.dumps(creds_dict), "filename": "creds.json"})

    project = quote_plus(config["project"])
    dataset = quote_plus(config.get("dataset", ""))
    location = config.get("location", "US")

    url = f"bigquery://{project}/{dataset}" if dataset else f"bigquery://{project}"
    
    pool_config = get_pool_config("bigquery")
    return create_engine(url, credentials_path=creds_path, location=location, **pool_config)


def _build_redshift_engine(auth_method: str, config: Dict) -> Engine:
    """
    Redshift engine.
    Auth: password, iam_role
    """
    from urllib.parse import quote_plus
    
    host = config["host"]
    port = config["port"]
    database = quote_plus(config["database"])
    connect_args = {}
    
    if auth_method == "iam_role":
        url = f"redshift+redshift_connector://@{host}:{port}/{database}"
        connect_args = {
            "iam": True,
            "db_user": config["db_user"],
            "cluster_identifier": config["cluster_identifier"],
            "region": config.get("region", "us-east-1"),
        }
        if config.get("profile"):
            connect_args["profile"] = config["profile"]
    else:  # password
        user = quote_plus(config["user"])
        password = quote_plus(config["password"])
        url = f"redshift+redshift_connector://{user}:{password}@{host}:{port}/{database}"
        if config.get("sslmode"):
            connect_args["sslmode"] = config["sslmode"]

    pool_config = get_pool_config("redshift")
    return create_engine(url, connect_args=connect_args, **pool_config)


def _build_duckdb_engine(auth_method: str, config: Dict) -> Engine:
    """
    DuckDB/MotherDuck engine.
    Auth: local_file, motherduck
    """
    from urllib.parse import quote_plus
    
    if auth_method == "local_file":
        url = f"duckdb:///{config['database_path']}"
        return create_engine(url, poolclass=StaticPool)
    else:  # motherduck
        database = quote_plus(config["database"])
        token = quote_plus(config["token"])
        url = f"duckdb:///md:{database}?motherduck_token={token}"
        pool_config = get_pool_config("duckdb")
        return create_engine(url, **pool_config)


def _build_databricks_engine(auth_method: str, config: Dict) -> Engine:
    """
    Build Databricks engine.
    
    Supported auth methods:
    - token: Personal Access Token (PAT) - long-lived
    - oauth_m2m: OAuth machine-to-machine (service principal) - persistent
    """
    from .utils.databricks_engine import DatabricksEngine, build_databricks_connection
    conn = build_databricks_connection(auth_method, config)
    return DatabricksEngine(conn)


def _build_athena_engine(auth_method: str, config: Dict) -> Engine:
    """
    AWS Athena engine.
    Auth: iam_credentials
    """
    from urllib.parse import quote_plus
    
    access_key = quote_plus(config["aws_access_key_id"])
    secret_key = quote_plus(config["aws_secret_access_key"])
    region = config["region"]
    schema = quote_plus(config["schema"])
    s3_staging = quote_plus(config["s3_staging_dir"])
    
    url = (
        f"awsathena+rest://{access_key}:{secret_key}"
        f"@athena.{region}.amazonaws.com:443/{schema}"
        f"?s3_staging_dir={s3_staging}"
    )
    
    if config.get("catalog"):
        url += f"&catalog_name={quote_plus(config['catalog'])}"
    if config.get("work_group"):
        url += f"&work_group={quote_plus(config['work_group'])}"

    pool_config = get_pool_config("athena")
    return create_engine(url, **pool_config)


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

ENGINE_BUILDERS = {
    "postgresql": _build_postgresql_engine,
    "mysql": _build_mysql_engine,
    "mariadb": _build_mariadb_engine,
    "clickhouse": _build_clickhouse_engine,
    "sqlserver": _build_sqlserver_engine,
    "trino": _build_trino_engine,
    "sparksql": _build_sparksql_engine,
    "vertica": _build_vertica_engine,
    "oracle": _build_oracle_engine,
    "teradata": _build_teradata_engine,
    "db2": _build_db2_engine,
    "snowflake": _build_snowflake_engine,
    "bigquery": _build_bigquery_engine,
    "redshift": _build_redshift_engine,
    "duckdb": _build_duckdb_engine,
    "databricks": _build_databricks_engine,
    "athena": _build_athena_engine,
}
