"""SQLAlchemy ORM table generation from schema definitions.

Converts :class:`TableDef` metadata into SQLAlchemy :class:`Table` objects
attached to a :class:`MetaData` instance.  These generated tables feed
Alembic's ``--autogenerate`` so that schema changes are detected directly
from Pydantic :class:`DomainModel` classes — no hand-written DDL, no
SQLAlchemy ORM models to maintain.

Usage::

    from sqlalchemy import MetaData
    from db_core.schema import TableDef, ColumnDef
    from db_core.schema.sqlalchemy_gen import table_to_sqlalchemy

    metadata = MetaData()
    table_def = TableDef(
        name="users",
        columns=[ColumnDef("id", "UUID", primary_key=True)],
    )
    table = table_to_sqlalchemy(table_def, metadata)
    # metadata.tables["users"] is now a SQLAlchemy Table
"""

from __future__ import annotations

import re
from typing import Any

from sqlalchemy import (
    BigInteger,
    Boolean,
    Column,
    Date,
    DateTime,
    Float,
    ForeignKeyConstraint,
    Index,
    Integer,
    MetaData,
    String,
    Table,
    Text,
    UniqueConstraint,
)
from sqlalchemy.dialects.postgresql import JSONB, UUID

try:
    from sqlalchemy.dialects.mysql import JSON as MySQLJSON, TINYINT
except ImportError:  # pragma: no cover
    MySQLJSON = None  # type: ignore[assignment]
    TINYINT = None  # type: ignore[assignment]

from db_core.schema.builder import ColumnDef, ForeignKeyDef, IndexDef, TableDef


def _map_col_type(col_type: str) -> Any:
    """Map a SQL type string from :class:`ColumnDef` to a SQLAlchemy type instance.

    Supported types: ``BIGINT``, ``BIGSERIAL``, ``INTEGER``, ``SERIAL``,
    ``UUID``, ``TEXT``, ``BOOLEAN``, ``DOUBLE PRECISION``, ``DATE``,
    ``TIMESTAMPTZ``, ``JSONB``, ``VARCHAR(n)``.

    Args:
        col_type: The ``col_type`` string from a :class:`ColumnDef`.

    Returns:
        A SQLAlchemy type instance (e.g. ``BigInteger()``, ``String(255)``).

    Raises:
        ValueError: If *col_type* is not recognised.
    """
    upper = col_type.upper().strip()

    # VARCHAR(n) — parameterised
    match = re.match(r"VARCHAR\((\d+)\)", upper)
    if match:
        return String(int(match.group(1)))

    _MAP: dict[str, Any] = {
        "SERIAL": Integer,
        "BIGSERIAL": BigInteger,
        "BIGINT": BigInteger,
        "INTEGER": Integer,
        "UUID": UUID,
        "TEXT": Text,
        "BOOLEAN": Boolean,
        "DOUBLE PRECISION": Float,
        "DATE": Date,
        "TIMESTAMPTZ": DateTime(timezone=True),
        "JSONB": JSONB,
    }

    factory = _MAP.get(upper)
    if factory is not None:
        return factory() if callable(factory) else factory

    raise ValueError(
        f"Unknown col_type '{col_type}' — cannot map to SQLAlchemy type. "
        f"Supported types: {', '.join(sorted(_MAP))}, VARCHAR(n)"
    )


def _map_col_type_doris(col_type: str) -> Any:
    """Map a SQL type string from :class:`ColumnDef` to a MySQL-compatible
    SQLAlchemy type instance for use with Apache Doris.

    Doris uses the MySQL protocol, so types are mapped to MySQL-dialect
    SQLAlchemy types.  Key differences from :func:`_map_col_type`:

    * ``UUID`` → ``String(36)`` (Doris has no native UUID type)
    * ``JSONB`` → ``JSON`` (MySQL ``JSON`` type via ``sqlalchemy.dialects.mysql``)
    * ``TIMESTAMPTZ`` → ``DateTime`` (Doris ``DATETIME`` has no timezone)
    * ``BOOLEAN`` → ``TINYINT`` (Doris does not accept the ``BOOL`` keyword)
    * ``DOUBLE PRECISION`` → ``Float`` (same SA type, different SQL spelling)

    Supported types: ``BIGINT``, ``BIGSERIAL``, ``INTEGER``, ``SERIAL``,
    ``UUID``, ``TEXT``, ``BOOLEAN``, ``DOUBLE PRECISION``, ``DOUBLE``,
    ``DATE``, ``TIMESTAMPTZ``, ``DATETIME``, ``JSONB``, ``JSON``,
    ``VARCHAR(n)``.

    Args:
        col_type: The ``col_type`` string from a :class:`ColumnDef`.

    Returns:
        A MySQL-compatible SQLAlchemy type instance.

    Raises:
        ValueError: If *col_type* is not recognised.
    """
    upper = col_type.upper().strip()

    # VARCHAR(n) — parameterised
    match = re.match(r"VARCHAR\((\d+)\)", upper)
    if match:
        return String(int(match.group(1)))

    _DORIS_MAP: dict[str, Any] = {
        "SERIAL": Integer,
        "BIGSERIAL": BigInteger,
        "BIGINT": BigInteger,
        "INTEGER": Integer,
        "UUID": String(36),
        "TEXT": Text,
        "BOOLEAN": TINYINT() if TINYINT is not None else Integer,
        "TINYINT": TINYINT() if TINYINT is not None else Integer,
        "DOUBLE PRECISION": Float,
        "DOUBLE": Float,
        "DATE": Date,
        "TIMESTAMPTZ": DateTime,
        "DATETIME": DateTime,
        "JSONB": MySQLJSON() if MySQLJSON is not None else Text,
        "JSON": MySQLJSON() if MySQLJSON is not None else Text,
    }

    factory = _DORIS_MAP.get(upper)
    if factory is not None:
        return factory() if callable(factory) else factory

    raise ValueError(
        f"Unknown col_type '{col_type}' — cannot map to Doris/MySQL SQLAlchemy type. "
        f"Supported types: {', '.join(sorted(_DORIS_MAP))}, VARCHAR(n)"
    )


def _should_include_field(
    col: ColumnDef,
    *,
    db_include: set[str] | None,
) -> bool:
    """Check whether a column should be included based on the model's
    ``_db_include`` restriction.

    Args:
        col: The column definition.
        db_include: The set of field names allowed in the DB table,
            or ``None`` to include all fields.

    Returns:
        ``True`` if the column should appear in the generated table.
    """
    if db_include is None:
        return True
    return col.name in db_include


def table_to_sqlalchemy(
    table: TableDef,
    metadata: MetaData,
    *,
    db_include: set[str] | None = None,
) -> Table:
    """Convert a :class:`TableDef` into a SQLAlchemy :class:`Table` attached
    to *metadata*.

    Columns, indexes, unique constraints, and foreign keys are all derived
    from the structured definition.  This is the bridge between Pydantic-
    derived schema metadata and Alembic's ``--autogenerate``.

    Args:
        table: The structured table definition (from :func:`model_to_table_def`).
        metadata: The SQLAlchemy :class:`MetaData` to attach the table to.
        db_include: Optional set of field names to include.  When ``None``
            (the default), all columns from *table* are included.  Pass the
            model's ``_db_include`` attribute to restrict which fields appear
            in the generated table.

    Returns:
        A fully-configured SQLAlchemy :class:`Table` registered on *metadata*.

    Raises:
        ValueError: If a ``ColumnDef.col_type`` cannot be mapped.
    """
    columns: list[Column] = []
    indexes: list[Index] = []

    for col_def in table.columns:
        if not _should_include_field(col_def, db_include=db_include):
            continue

        sa_type = _map_col_type(col_def.col_type)

        # Column constraints
        kwargs: dict[str, Any] = {}
        if col_def.primary_key:
            kwargs["primary_key"] = True
        elif not col_def.nullable:
            kwargs["nullable"] = False

        columns.append(Column(col_def.name, sa_type, **kwargs))

    # Indexes
    for idx_def in table.indexes:
        idx_name = idx_def.name
        if idx_name is None:
            prefix = "uq" if idx_def.unique else "ix"
            cols_part = "_".join(idx_def.columns)
            idx_name = f"{prefix}_{table.name}_{cols_part}"

        idx_kwargs: dict[str, Any] = {"unique": idx_def.unique}
        if idx_def.using != "btree":
            idx_kwargs["postgresql_using"] = idx_def.using

        indexes.append(Index(idx_name, *idx_def.columns, **idx_kwargs))

    # Foreign keys
    fks: list[ForeignKeyConstraint] = []
    for fk_def in table.foreign_keys:
        fk_name = f"fk_{table.name}_{fk_def.column}_{fk_def.ref_table}"
        fks.append(
            ForeignKeyConstraint(
                [fk_def.column],
                [f"{fk_def.ref_table}.{fk_def.ref_column}"],
                ondelete=fk_def.on_delete,
                name=fk_name,
            )
        )

    return Table(
        table.name,
        metadata,
        *columns,
        *indexes,
        *fks,
    )


def table_to_sqlalchemy_doris(
    table: TableDef,
    metadata: MetaData,
    *,
    warehouse_include: set[str] | None = None,
    duplicate_key: list[str] | None = None,
    unique_key: list[str] | None = None,
    buckets: int = 10,
) -> Table:
    """Convert a :class:`TableDef` into a MySQL-compatible SQLAlchemy
    :class:`Table` for use with Apache Doris.

    Unlike :func:`table_to_sqlalchemy`, this function:

    * Uses :func:`_map_col_type_doris` for MySQL-compatible types.
    * Skips foreign key constraints (Doris does not support them).
    * Skips ``postgresql_using`` on indexes (GIN/GiST unsupported).
    * Stores Doris-specific metadata (``duplicate_key`` / ``unique_key``,
      ``buckets``) as table info for the Alembic migration template to
      consume.

    Exactly one of *duplicate_key* or *unique_key* may be set.  If neither
    is provided, ``DUPLICATE KEY(["id"])`` is assumed (append-only default).

    Args:
        table: The structured table definition.
        metadata: The SQLAlchemy :class:`MetaData` to attach the table to.
        warehouse_include: Optional set of field names to include.
            Pass the model's ``_warehouse_include`` attribute.
        duplicate_key: List of column names for the ``DUPLICATE KEY``
            clause.  Mutually exclusive with *unique_key*.
        unique_key: List of column names for the ``UNIQUE KEY`` clause
            (enables upsert semantics).  Mutually exclusive with
            *duplicate_key*.
        buckets: Number of hash buckets for ``DISTRIBUTED BY HASH``.
            Defaults to 10.

    Returns:
        A MySQL-compatible SQLAlchemy :class:`Table` registered on *metadata*.

    Raises:
        ValueError: If both *duplicate_key* and *unique_key* are set, or
            if a ``ColumnDef.col_type`` cannot be mapped.
    """
    if duplicate_key and unique_key:
        raise ValueError(
            f"Table '{table.name}': duplicate_key and unique_key are mutually exclusive. "
            f"Set one or neither (defaults to DUPLICATE KEY([\"id\"]))."
        )
    columns: list[Column] = []
    indexes: list[Index] = []

    for col_def in table.columns:
        if not _should_include_field(col_def, db_include=warehouse_include):
            continue

        sa_type = _map_col_type_doris(col_def.col_type)

        # Column constraints — no PRIMARY KEY for Doris key models;
        # the DUPLICATE KEY / UNIQUE KEY clause serves that role.
        kwargs: dict[str, Any] = {}
        if not col_def.nullable and not col_def.primary_key:
            kwargs["nullable"] = False

        columns.append(Column(col_def.name, sa_type, **kwargs))

    # Indexes — skip postgresql_using, Doris only supports btree
    for idx_def in table.indexes:
        idx_name = idx_def.name
        if idx_name is None:
            prefix = "uq" if idx_def.unique else "ix"
            cols_part = "_".join(idx_def.columns)
            idx_name = f"{prefix}_{table.name}_{cols_part}"

        idx_kwargs: dict[str, Any] = {"unique": idx_def.unique}
        indexes.append(Index(idx_name, *idx_def.columns, **idx_kwargs))

    # Store Doris-specific attributes for the Alembic template
    sa_table = Table(
        table.name,
        metadata,
        *columns,
        *indexes,
        # No foreign keys — Doris doesn't support them
    )

    if unique_key:
        sa_table.info["doris_key_type"] = "unique"
        sa_table.info["doris_unique_key"] = unique_key
    else:
        sa_table.info["doris_key_type"] = "duplicate"
        sa_table.info["doris_duplicate_key"] = duplicate_key or ["id"]

    sa_table.info["doris_buckets"] = buckets

    return sa_table


def table_defs_to_metadata(
    table_defs: list[TableDef],
    *,
    db_include_map: dict[str, set[str] | None] | None = None,
) -> MetaData:
    """Convert a list of :class:`TableDef` instances into a populated
    :class:`MetaData`.

    Convenience wrapper around :func:`table_to_sqlalchemy` for bulk
    conversion — useful when generating Alembic metadata from a registry
    of domain models.

    Args:
        table_defs: List of table definitions to convert.
        db_include_map: Optional mapping of ``table_name → set of field names``
            for per-table field filtering.  Tables not in the map include
            all fields.

    Returns:
        A SQLAlchemy :class:`MetaData` with all tables registered.
    """
    metadata = MetaData()
    db_include_map = db_include_map or {}

    for td in table_defs:
        table_to_sqlalchemy(
            td,
            metadata,
            db_include=db_include_map.get(td.name),
        )

    return metadata
