"""MySQL database setup (SQLAlchemy)."""
from __future__ import annotations

import logging
import time
from urllib.parse import quote_plus

from sqlalchemy import create_engine, event, text
from sqlalchemy.orm import declarative_base, sessionmaker

from ..config import (
    MYSQL_DATABASE,
    MYSQL_HOST,
    MYSQL_PASSWORD,
    MYSQL_POOL_RECYCLE,
    MYSQL_POOL_SIZE,
    MYSQL_PORT,
    MYSQL_SLOW_QUERY_MS,
    MYSQL_USER,
    ensure_dirs,
    mysql_url,
)

ensure_dirs()

logger = logging.getLogger(__name__)


def _ensure_database() -> None:
    """Create MYSQL_DATABASE if it does not exist yet."""
    user = quote_plus(MYSQL_USER)
    password = quote_plus(MYSQL_PASSWORD)
    server_url = (
        f"mysql+pymysql://{user}:{password}"
        f"@{MYSQL_HOST}:{MYSQL_PORT}/?charset=utf8mb4"
    )
    server_engine = create_engine(server_url, isolation_level="AUTOCOMMIT")
    try:
        with server_engine.connect() as conn:
            conn.execute(
                text(
                    f"CREATE DATABASE IF NOT EXISTS `{MYSQL_DATABASE}` "
                    "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
                )
            )
    finally:
        server_engine.dispose()


_ensure_database()

ENGINE = create_engine(
    mysql_url(),
    pool_size=MYSQL_POOL_SIZE,
    pool_recycle=MYSQL_POOL_RECYCLE,
    pool_pre_ping=True,
    future=True,
)

if MYSQL_SLOW_QUERY_MS > 0:
    @event.listens_for(ENGINE, "before_cursor_execute")
    def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
        conn.info["query_start_time"] = time.perf_counter()

    @event.listens_for(ENGINE, "after_cursor_execute")
    def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
        start = conn.info.pop("query_start_time", None)
        if start is None:
            return
        elapsed_ms = (time.perf_counter() - start) * 1000
        if elapsed_ms >= MYSQL_SLOW_QUERY_MS:
            logger.warning(
                "Slow query (%.0f ms): %s",
                elapsed_ms,
                statement[:500],
            )

SessionLocal = sessionmaker(bind=ENGINE, autoflush=False, expire_on_commit=False, future=True)
Base = declarative_base()


def init_db() -> None:
    from . import models  # noqa: F401  (register models)
    Base.metadata.create_all(ENGINE)
    _migrate()


def _migrate() -> None:
    """Add columns introduced after a DB was first created (create_all won't alter)."""
    added = {
        "sessions": [
            ("progress_rows_done", "INTEGER DEFAULT 0"),
            ("progress_rows_total", "INTEGER DEFAULT 0"),
            ("eta_seconds", "INTEGER DEFAULT 0"),
            ("opening_mode", "VARCHAR(255) DEFAULT 'zero'"),
            ("opening_source_session_id", "INTEGER"),
        ],
        "reconciliation": [
            ("received_payouts", "FLOAT DEFAULT 0"),
            ("all_payouts", "FLOAT DEFAULT 0"),
        ],
        "finance_control": [
            ("tolerance", "FLOAT DEFAULT 1.0"),
        ],
        "exports": [
            ("kind", "VARCHAR(255) DEFAULT 'full'"),
        ],
        "transactions": [
            ("txn_type_en", "VARCHAR(255)"),
            ("storage_flag", "BOOLEAN DEFAULT 0"),
        ],
    }
    with ENGINE.begin() as conn:
        db_name = conn.execute(text("SELECT DATABASE()")).scalar()
        for table, cols in added.items():
            existing = {
                r[0]
                for r in conn.execute(
                    text(
                        "SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS "
                        "WHERE TABLE_SCHEMA = :schema AND TABLE_NAME = :table"
                    ),
                    {"schema": db_name, "table": table},
                )
            }
            for name, decl in cols:
                if name not in existing:
                    conn.execute(text(f"ALTER TABLE `{table}` ADD COLUMN `{name}` {decl}"))


def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()