import logging
import sqlalchemy as sa
from sqlalchemy import inspect, func, select, Table, MetaData
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_tables_list(
    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 Tables List"
    tables_data = None
    options = options or {}
    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 ""
    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 list of tables from {db_type} database."
    )
    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)
        table_names = inspector.get_table_names()
        data = []
        metadata = MetaData()
        with engine.connect() as connection:
            for table in table_names:
                table_obj = Table(table, metadata)
                count_query = select(func.count()).select_from(table_obj)
                row_count = connection.scalar(count_query)
                data.append({"table_name": table, "row_count": row_count})
        if out_type.lower() == "pandas":
            tables_data = pd.DataFrame(data)
        elif out_type.lower() == "polars":
            tables_data = pl.DataFrame(data)
        else:
            tables_data = data
        engine.dispose()
        verbose and logger.info(
            f"[{brick_display_name}] Tables list retrieved successfully. Tables found: {(len(tables_data) if tables_data is not None else 0)}"
        )
    except Exception as e:
        verbose and logger.error(
            f"[{brick_display_name}] An error occurred while retrieving tables list: {str(e)}"
        )
        raise
    return tables_data