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

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 query_db(
    sql_text: Str,
    database_name: Str = None,
    host: Str = None,
    port: Int = None,
    username: Str = None,
    password: SecretStr = None,
    options=None,
) -> DataFrame:
    brick_display_name = "Query Database"
    df = None
    options = options or {}
    verbose = options.get("verbose", True)
    db_type = options.get("db_type", "SQLite")
    df_type = options.get("df_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}")
    if verbose:
        logger.info(f"[{brick_display_name}] Starting query to {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.")
        with engine.connect() as conn:
            if df_type.lower() == "pandas":
                df = pd.read_sql(text(sql_text), conn)
            elif df_type.lower() == "polars":
                df = pl.read_database(text(sql_text), connection=conn)
            else:
                raise ValueError("df_type must be 'pandas' or 'polars'")
        engine.dispose()
        verbose and logger.info(
            f"[{brick_display_name}] Query executed successfully. Rows returned: {(len(df) if df is not None else 0)}"
        )
    except Exception as e:
        verbose and logger.error(f"[{brick_display_name}] Error during query: {str(e)}")
        raise
    return df