import logging
import sqlalchemy as sa
from sqlalchemy import inspect
import pandas as pd
import polars as pl
from os import getenv
from coded_flows.types import Str, SecretStr, Int, DataFrame, DataRecords, Union

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


def _coalesce(*values):
    return next((v for v in values if v is not None))


def get_table_schema(
    table_name: Str = None,
    database_name: Str = None,
    host: Str = None,
    port: Int = None,
    username: Str = None,
    password: SecretStr = None,
    options=None,
) -> Union[DataFrame, DataRecords]:
    brick_display_name = "Get Table Schema"
    schema = None
    options = options or {}
    op_table_name = options.get("table_name", "")
    verbose = options.get("verbose", True)
    db_type = options.get("db_type", "SQLite")
    out_type = options.get("out_type", "Pandas")
    env_prefix = options.get("env_prefix", "")
    env_prefix = env_prefix + "_" if env_prefix.strip() else ""
    table_name = _coalesce(table_name, op_table_name)
    try:
        database_name = _coalesce(database_name, getenv(f"{env_prefix}DB_NAME"))
        if db_type != "SQLite":
            host = _coalesce(host, getenv(f"{env_prefix}DB_HOST"))
            port = _coalesce(port, getenv(f"{env_prefix}DB_PORT"))
            username = _coalesce(username, getenv(f"{env_prefix}DB_USERNAME"))
            password = _coalesce(password, getenv(f"{env_prefix}DB_PASSWORD"))
        if not database_name or (
            db_type != "SQLite"
            and (not host or not port or (not username) or (not password))
        ):
            verbose and logger.error(
                f"[{brick_display_name}] Either the env file is not defined, not referenced in the project folder, or one or more DB configuration inputs are missing."
            )
            raise ValueError(
                "Either the env file is not defined, not referenced in the project folder, or one or more DB configuration inputs are missing."
            )
        try:
            if hasattr(password, "get_secret_value"):
                password = password.get_secret_value()
            else:
                password = str(password)
        except Exception:
            password = str(password)
    except Exception as e:
        verbose and logger.error(
            f"[{brick_display_name}] Either the env file is not defined, not referenced in the project folder, or one or more DB configuration inputs are missing."
        )
        raise ValueError(
            "Either the env file is not defined, not referenced in the project folder, or one or more DB configuration inputs are missing."
        )
    try:
        if port:
            port = int(port)
    except (ValueError, TypeError):
        verbose and logger.error(
            f"[{brick_display_name}] Port must be an integer, got {port!r}"
        )
        raise ValueError(f"Port must be an integer, got {port!r}")
    if port and (not 1 <= port <= 65535):
        verbose and logger.error(
            f"[{brick_display_name}] Port must be between 1 and 65535, got {port}"
        )
        raise ValueError(f"Port must be between 1 and 65535, got {port}")
    verbose and logger.info(
        f"[{brick_display_name}] Retrieving schema from {db_type} database for table '{table_name}'."
    )
    driver_map = {
        "Postgres": "postgresql+psycopg",
        "MySQL": "mysql+pymysql",
        "MariaDB": "mariadb+pymysql",
        "Oracle Database": "oracle+oracledb",
        "Microsoft SQL Server": "mssql+pymssql",
        "SQLite": "sqlite",
    }
    if db_type == "SQLite":
        if host or username or password:
            verbose and logger.warning(
                f"[{brick_display_name}] Ignoring host, username, password for SQLite."
            )
        url = f"{driver_map[db_type]}:///{database_name}"
    else:
        if not all([host, str(port), username, password, database_name]):
            verbose and logger.error(
                f"[{brick_display_name}] Missing required connection parameters for {db_type}."
            )
            raise ValueError(
                f"[{brick_display_name}] Missing required connection parameters for {db_type}."
            )
        url = f"{driver_map[db_type]}://{username}:{password}@{host}:{port}/{database_name}"
    verbose and logger.info(f"[{brick_display_name}] Connection URL constructed.")
    try:
        engine = sa.create_engine(url)
        verbose and logger.info(f"[{brick_display_name}] Engine created.")
        inspector = inspect(engine)
        if not inspector.has_table(table_name):
            verbose and logger.error(
                f"[{brick_display_name}] Table '{table_name}' does not exist in the database."
            )
            raise ValueError(f"Table '{table_name}' does not exist in the database.")
        columns = inspector.get_columns(table_name)
        columns = [{**col, "type": str(col["type"])} for col in columns]
        if out_type.lower() == "pandas":
            schema = pd.DataFrame(columns)
        elif out_type.lower() == "polars":
            schema = pl.DataFrame(columns)
        else:
            schema = columns
        engine.dispose()
        verbose and logger.info(
            f"[{brick_display_name}] Schema retrieved successfully. Columns found: {(len(schema) if schema is not None else 0)}"
        )
    except Exception as e:
        verbose and logger.error(
            f"[{brick_display_name}] An error occured while retrieving schema"
        )
        raise
    return schema