"""MigrationBuilder — fluent DSL for defining database schema migrations.

Works like ``QueryBuilder`` but for schema: chain column-type methods to build
a table definition, then call :meth:`create_table`, :meth:`drop_table`, or
:meth:`alter_table` to produce an executable :class:`Op`.

Usage::

    from repo_core.migrations.builder import MigrationBuilder

    # Create a new table
    op = (
        MigrationBuilder("articles")
        .bigint("id", primary_key=True, nullable=False)
        .text("title", nullable=False)
        .text("body", nullable=True)
        .boolean("published", nullable=False, default=False)
        .timestamp("created_at", nullable=True, default="now()")
        .timestamp("updated_at", nullable=True, default="now()")
        .create_table()
    )
    print(op.sql())              # → list[str] of SQL DDL
    await op.execute(pg_client)  # → runs DDL against the DB

    # Drop a table
    drop_op = MigrationBuilder("articles").drop_table()

    # Alter an existing table
    alter_op = (
        MigrationBuilder("articles")
        .add_column("subtitle", "TEXT", nullable=True)
        .drop_column("old_field")
        .alter_table()
    )

    # Doris dialect
    doris_op = (
        MigrationBuilder("telegram_messages", dialect="doris")
        .uuid("id")
        .bigint("telegram_id", nullable=False)
        .doris_unique_key(["id"])
        .doris_distributed_by("id", buckets=10)
        .create_table()
    )
"""

from __future__ import annotations

from typing import Any

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

# ---------------------------------------------------------------------------
# Dialect type mappings
# ---------------------------------------------------------------------------

_PG_TYPES: dict[str, str] = {
    "uuid": "UUID",
    "text": "TEXT",
    "bigint": "BIGINT",
    "integer": "INTEGER",
    "boolean": "BOOLEAN",
    "float_": "DOUBLE PRECISION",
    "float_short": "REAL",
    "jsonb": "JSONB",
    "timestamp": "TIMESTAMPTZ",
    "date": "DATE",
}

_DORIS_TYPES: dict[str, str] = {
    "uuid": "VARCHAR(36)",
    "text": "TEXT",
    "bigint": "BIGINT",
    "integer": "INT",
    "boolean": "TINYINT",
    "float_": "DOUBLE",
    "float_short": "FLOAT",
    "jsonb": "JSON",
    "timestamp": "DATETIME",
    "date": "DATE",
}


def _sql_type(logical: str, dialect: str, length: int | None = None) -> str:
    """Map a logical column type name to a dialect-specific SQL type string."""
    if logical == "varchar":
        return f"VARCHAR({length or 255})"
    if logical == "numeric":
        raise ValueError("Use _col() with logical='numeric' and precision/scale arguments.")
    types = _DORIS_TYPES if dialect == "doris" else _PG_TYPES
    return types[logical]


def _resolve_default(default: Any, dialect: str) -> str | None:
    """Translate a Python default value to a dialect-appropriate SQL literal."""
    if default is None:
        return None
    if isinstance(default, str):
        low = default.lower().strip()
        if low in ("now()", "current_timestamp"):
            return "NOW()" if dialect == "postgresql" else "CURRENT_TIMESTAMP"
        # Already SQL-quoted literal/expression (e.g. "'active'", "'{}'::jsonb")
        if default.startswith("'"):
            return default
        escaped = default.replace("'", "''")
        return f"'{escaped}'"
    if isinstance(default, bool):
        if dialect == "doris":
            return "1" if default else "0"
        return "TRUE" if default else "FALSE"
    if isinstance(default, (int, float)):
        return str(default)
    return str(default)


# ---------------------------------------------------------------------------
# Op base and concrete types
# ---------------------------------------------------------------------------


class Op:
    """Abstract base for a migration operation that produces and executes SQL."""

    def sql(self) -> list[str]:
        """Return ordered SQL statements implementing this operation."""
        raise NotImplementedError

    async def execute(self, client: Any) -> None:
        """Execute all SQL statements against the given DB client."""
        for stmt in self.sql():
            await client.execute_raw(stmt)


class CreateTableOp(Op):
    """Creates a new table, plus its indexes."""

    def __init__(
        self,
        table_name: str,
        columns: list[ColumnDef],
        indexes: list[IndexDef],
        foreign_keys: list[ForeignKeyDef],
        dialect: str,
        if_not_exists: bool,
        doris_key_type: str | None,
        doris_key_columns: list[str],
        doris_distributed_column: str | None,
        doris_buckets: int,
        doris_properties: dict[str, str] | None = None,
        doris_partition_column: str | None = None,
    ) -> None:
        self._table_name = table_name
        self._columns = columns
        self._indexes = indexes
        self._foreign_keys = foreign_keys
        self._dialect = dialect
        self._if_not_exists = if_not_exists
        self._doris_key_type = doris_key_type
        self._doris_key_columns = doris_key_columns
        self._doris_distributed_column = doris_distributed_column
        self._doris_buckets = doris_buckets
        self._doris_properties = doris_properties or {}
        self._doris_partition_column = doris_partition_column

    def sql(self) -> list[str]:
        if self._dialect == "doris":
            return self._sql_doris()
        return self._sql_pg()

    def _sql_pg(self) -> list[str]:
        table_def = TableDef(
            name=self._table_name,
            columns=self._columns,
            indexes=self._indexes,
            foreign_keys=self._foreign_keys,
            if_not_exists=self._if_not_exists,
        )
        stmts = [SchemaBuilder.create_table(table_def)]
        for idx in self._indexes:
            stmts.append(SchemaBuilder.create_index(self._table_name, idx))
        return stmts

    def _sql_doris(self) -> list[str]:
        ie = " IF NOT EXISTS" if self._if_not_exists else ""
        col_lines: list[str] = []
        for col in self._columns:
            line = f"  `{col.name}` {col.col_type}"
            if not col.nullable:
                line += " NOT NULL"
            if col.default is not None:
                line += f" DEFAULT {col.default}"
            col_lines.append(line)
        cols_sql = ",\n".join(col_lines)
        ddl = f"CREATE TABLE{ie} `{self._table_name}` (\n{cols_sql}\n) ENGINE=OLAP"

        if self._doris_key_type and self._doris_key_columns:
            key_cols = ", ".join(f"`{c}`" for c in self._doris_key_columns)
            ddl += f"\n{self._doris_key_type.upper()} KEY({key_cols})"

        if self._doris_partition_column:
            ddl += f"\nPARTITION BY RANGE(`{self._doris_partition_column}`) ()"

        dist_col = self._doris_distributed_column or (
            self._doris_key_columns[0] if self._doris_key_columns else self._columns[0].name
        )
        ddl += f"\nDISTRIBUTED BY HASH(`{dist_col}`) BUCKETS {self._doris_buckets}"

        if self._doris_properties:
            prop_lines = ",\n".join(
                f'    "{k}" = "{v}"' for k, v in self._doris_properties.items()
            )
            ddl += f"\nPROPERTIES (\n{prop_lines}\n)"

        return [ddl]


class DropTableOp(Op):
    """Drops a table."""

    def __init__(self, table_name: str, dialect: str, if_exists: bool = True) -> None:
        self._table_name = table_name
        self._dialect = dialect
        self._if_exists = if_exists

    def sql(self) -> list[str]:
        if self._dialect == "doris":
            ie = " IF EXISTS" if self._if_exists else ""
            return [f"DROP TABLE{ie} `{self._table_name}`"]
        return [SchemaBuilder.drop_table(self._table_name, self._if_exists)]


class RenameTableOp(Op):
    """Renames an existing table."""

    def __init__(
        self, old_name: str, new_name: str, dialect: str
    ) -> None:
        self._old_name = old_name
        self._new_name = new_name
        self._dialect = dialect

    def sql(self) -> list[str]:
        if self._dialect == "doris":
            return [f"ALTER TABLE `{self._old_name}` RENAME `{self._new_name}`"]
        return [f"ALTER TABLE {self._old_name} RENAME TO {self._new_name}"]


class AlterTableOp(Op):
    """Alters an existing table: add/drop columns, rename columns, change
    column types, add/drop indexes."""

    def __init__(
        self,
        table_name: str,
        add_columns: list[ColumnDef],
        drop_columns: list[str],
        rename_columns: list[tuple[str, str]],
        change_column_types: list[tuple[str, str, str | None]],
        add_indexes: list[IndexDef],
        drop_indexes: list[str],
        dialect: str,
    ) -> None:
        self._table_name = table_name
        self._add_columns = add_columns
        self._drop_columns = drop_columns
        self._rename_columns = rename_columns
        self._change_column_types = change_column_types
        self._add_indexes = add_indexes
        self._drop_indexes = drop_indexes
        self._dialect = dialect

    def sql(self) -> list[str]:
        if self._dialect == "doris":
            return self._sql_doris()
        return self._sql_pg()

    def _sql_pg(self) -> list[str]:
        stmts: list[str] = []
        for col_name, new_type, using_expr in self._change_column_types:
            using_clause = f" USING {using_expr}" if using_expr else ""
            stmts.append(
                f"ALTER TABLE {self._table_name} "
                f"ALTER COLUMN {col_name} TYPE {new_type}{using_clause}"
            )
        for col in self._add_columns:
            stmts.append(SchemaBuilder.add_column(self._table_name, col))
        for col_name in self._drop_columns:
            stmts.append(SchemaBuilder.drop_column(self._table_name, col_name))
        for old, new in self._rename_columns:
            stmts.append(f"ALTER TABLE {self._table_name} RENAME COLUMN {old} TO {new}")
        for idx in self._add_indexes:
            stmts.append(SchemaBuilder.create_index(self._table_name, idx))
        for idx_name in self._drop_indexes:
            stmts.append(SchemaBuilder.drop_index(idx_name))
        return stmts

    def _sql_doris(self) -> list[str]:
        import warnings

        stmts: list[str] = []
        if self._change_column_types:
            warnings.warn(
                "ALTER COLUMN TYPE is not supported in Doris; skipping type-change operations.",
                stacklevel=2,
            )
        for col in self._add_columns:
            line = f"ALTER TABLE `{self._table_name}` ADD COLUMN `{col.name}` {col.col_type}"
            if col.default is not None:
                line += f" DEFAULT {col.default}"
            stmts.append(line)
        for col_name in self._drop_columns:
            stmts.append(f"ALTER TABLE `{self._table_name}` DROP COLUMN `{col_name}`")
        if self._rename_columns:
            for old, new in self._rename_columns:
                stmts.append(f"ALTER TABLE `{self._table_name}` RENAME COLUMN `{old}` `{new}`")
        for idx in self._add_indexes:
            cols = ", ".join(f"`{c}`" for c in idx.columns)
            idx_name = idx.name or f"ix_{self._table_name}_{'_'.join(idx.columns)}"
            stmts.append(f"CREATE INDEX `{idx_name}` ON `{self._table_name}` ({cols})")
        for idx_name in self._drop_indexes:
            stmts.append(f"DROP INDEX `{idx_name}` ON `{self._table_name}`")
        return stmts


# ---------------------------------------------------------------------------
# MigrationBuilder
# ---------------------------------------------------------------------------


class MigrationBuilder:
    """Fluent builder for database schema migrations.

    Chain column-type methods to describe a table, then call a terminal method
    (:meth:`create_table`, :meth:`drop_table`, :meth:`alter_table`) to get an
    executable :class:`Op`.

    Args:
        table_name: The database table name.
        dialect: ``"postgresql"`` (default) or ``"doris"``.  Controls SQL type
            mapping and DDL generation style.
        if_not_exists: When ``True`` (default), emits ``IF NOT EXISTS`` /
            ``IF EXISTS`` guards on create/drop.
    """

    def __init__(
        self,
        table_name: str,
        *,
        dialect: str = "postgresql",
        if_not_exists: bool = True,
    ) -> None:
        self._table_name = table_name
        self._dialect = dialect
        self._if_not_exists = if_not_exists

        # create_table state
        self._columns: list[ColumnDef] = []
        self._indexes: list[IndexDef] = []
        self._foreign_keys: list[ForeignKeyDef] = []

        # Doris table-level key / distribution settings
        self._doris_key_type: str | None = None
        self._doris_key_columns: list[str] = []
        self._doris_distributed_column: str | None = None
        self._doris_buckets: int = 10
        self._doris_properties: dict[str, str] = {"replication_num": "1"}
        self._doris_partition_column: str | None = None

        # alter_table state
        self._add_columns: list[ColumnDef] = []
        self._drop_columns: list[str] = []
        self._rename_columns: list[tuple[str, str]] = []
        self._change_column_types: list[tuple[str, str, str | None]] = []
        self._add_indexes: list[IndexDef] = []
        self._drop_indexes: list[str] = []

    # ------------------------------------------------------------------
    # Internal helper
    # ------------------------------------------------------------------

    def _col(
        self,
        name: str,
        logical_type: str,
        *,
        primary_key: bool = False,
        nullable: bool = True,
        default: Any = None,
        length: int | None = None,
        precision: int | None = None,
        scale: int | None = None,
    ) -> "MigrationBuilder":
        if logical_type == "numeric":
            p = precision if precision is not None else 10
            s = scale if scale is not None else 2
            sql_type = f"NUMERIC({p},{s})"
        elif logical_type == "varchar":
            sql_type = f"VARCHAR({length or 255})"
        else:
            sql_type = _sql_type(logical_type, self._dialect)

        self._columns.append(
            ColumnDef(
                name=name,
                col_type=sql_type,
                nullable=nullable,
                default=_resolve_default(default, self._dialect),
                primary_key=primary_key,
            )
        )
        return self

    # ------------------------------------------------------------------
    # Column type methods (used with create_table)
    # ------------------------------------------------------------------

    def uuid(
        self,
        name: str,
        *,
        primary_key: bool = False,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """UUID column — maps to ``UUID`` (PostgreSQL) / ``VARCHAR(36)`` (Doris)."""
        return self._col(name, "uuid", primary_key=primary_key, nullable=nullable, default=default)

    def text(
        self,
        name: str,
        *,
        primary_key: bool = False,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """TEXT column."""
        return self._col(name, "text", primary_key=primary_key, nullable=nullable, default=default)

    def bigint(
        self,
        name: str,
        *,
        primary_key: bool = False,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """BIGINT column."""
        return self._col(name, "bigint", primary_key=primary_key, nullable=nullable, default=default)

    def integer(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """INTEGER column — maps to ``INTEGER`` (PostgreSQL) / ``INT`` (Doris)."""
        return self._col(name, "integer", nullable=nullable, default=default)

    def boolean(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """BOOLEAN column — maps to ``BOOLEAN`` (PostgreSQL) / ``TINYINT`` (Doris)."""
        return self._col(name, "boolean", nullable=nullable, default=default)

    def float_(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """Float column — maps to ``DOUBLE PRECISION`` (PostgreSQL) / ``DOUBLE`` (Doris)."""
        return self._col(name, "float_", nullable=nullable, default=default)

    def float_short(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """Single-precision float — maps to ``REAL`` (PostgreSQL) / ``FLOAT`` (Doris)."""
        return self._col(name, "float_short", nullable=nullable, default=default)

    def jsonb(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """JSONB/JSON column — maps to ``JSONB`` (PostgreSQL) / ``JSON`` (Doris)."""
        return self._col(name, "jsonb", nullable=nullable, default=default)

    def timestamp(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """Timestamp column — maps to ``TIMESTAMPTZ`` (PostgreSQL) / ``DATETIME`` (Doris).

        Pass ``default="now()"`` for a ``NOW()`` / ``CURRENT_TIMESTAMP`` default.
        """
        return self._col(name, "timestamp", nullable=nullable, default=default)

    def date(
        self,
        name: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """DATE column."""
        return self._col(name, "date", nullable=nullable, default=default)

    def varchar(
        self,
        name: str,
        length: int,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """VARCHAR(n) column."""
        return self._col(name, "varchar", nullable=nullable, default=default, length=length)

    def numeric(
        self,
        name: str,
        precision: int,
        scale: int,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """NUMERIC(precision, scale) column."""
        return self._col(
            name, "numeric", nullable=nullable, default=default,
            precision=precision, scale=scale,
        )

    # ------------------------------------------------------------------
    # Index & FK methods (used with create_table)
    # ------------------------------------------------------------------

    def index(
        self,
        columns: list[str],
        *,
        unique: bool = False,
        using: str = "btree",
        name: str | None = None,
    ) -> "MigrationBuilder":
        """Add an index definition (emits a separate ``CREATE INDEX`` statement)."""
        self._indexes.append(IndexDef(columns=columns, unique=unique, using=using, name=name))
        return self

    def foreign_key(
        self,
        column: str,
        ref_table: str,
        *,
        ref_column: str = "id",
        on_delete: str = "CASCADE",
    ) -> "MigrationBuilder":
        """Add a foreign key constraint — silently ignored for the Doris dialect."""
        if self._dialect != "doris":
            self._foreign_keys.append(
                ForeignKeyDef(
                    column=column,
                    ref_table=ref_table,
                    ref_column=ref_column,
                    on_delete=on_delete,
                )
            )
        return self

    # ------------------------------------------------------------------
    # Doris-specific methods
    # ------------------------------------------------------------------

    def doris_unique_key(self, columns: list[str]) -> "MigrationBuilder":
        """Set the Doris ``UNIQUE KEY`` for the table (Doris dialect only)."""
        self._doris_key_type = "unique"
        self._doris_key_columns = list(columns)
        return self

    def doris_duplicate_key(self, columns: list[str]) -> "MigrationBuilder":
        """Set the Doris ``DUPLICATE KEY`` for the table (Doris dialect only)."""
        self._doris_key_type = "duplicate"
        self._doris_key_columns = list(columns)
        return self

    def doris_distributed_by(self, column: str, buckets: int = 10) -> "MigrationBuilder":
        """Set ``DISTRIBUTED BY HASH(column) BUCKETS n`` (Doris dialect only)."""
        self._doris_distributed_column = column
        self._doris_buckets = buckets
        return self

    def doris_partition_by_range(self, column: str) -> "MigrationBuilder":
        """Set ``PARTITION BY RANGE(column) ()`` for dynamic partitioning (Doris)."""
        self._doris_partition_column = column
        return self

    def doris_properties(self, props: dict[str, str]) -> "MigrationBuilder":
        """Set ``PROPERTIES (...)`` clause for the table (Doris dialect only).

        Merges *props* on top of defaults (``replication_num = "1"``).

        Example::

            .doris_properties({
                "replication_num": "1",
                "colocate_with": "telegram_group",
            })
        """
        self._doris_properties.update(props)
        return self

    # ------------------------------------------------------------------
    # Alter-table methods (used with alter_table)
    # ------------------------------------------------------------------

    def add_column(
        self,
        name: str,
        col_type: str,
        *,
        nullable: bool = True,
        default: Any = None,
    ) -> "MigrationBuilder":
        """Stage adding a column to an existing table (for :meth:`alter_table`)."""
        self._add_columns.append(
            ColumnDef(
                name=name,
                col_type=col_type,
                nullable=nullable,
                default=_resolve_default(default, self._dialect),
            )
        )
        return self

    def drop_column(self, name: str) -> "MigrationBuilder":
        """Stage dropping a column from an existing table (for :meth:`alter_table`)."""
        self._drop_columns.append(name)
        return self

    def rename_column(self, old_name: str, new_name: str) -> "MigrationBuilder":
        """Stage renaming a column (for :meth:`alter_table`; not supported in Doris)."""
        self._rename_columns.append((old_name, new_name))
        return self

    def change_column_type(
        self,
        name: str,
        new_type: str,
        *,
        using: str | None = None,
    ) -> "MigrationBuilder":
        """Stage changing a column's data type (for :meth:`alter_table`).

        Args:
            name: Column name.
            new_type: Target SQL type, e.g. ``"BIGINT"``.
            using: Optional ``USING`` expression for the cast, e.g.
                ``"channel_id::BIGINT"``.  Required when PostgreSQL cannot
                perform an implicit cast.
        """
        self._change_column_types.append((name, new_type, using))
        return self

    def add_index(
        self,
        columns: list[str],
        *,
        unique: bool = False,
        using: str = "btree",
        name: str | None = None,
    ) -> "MigrationBuilder":
        """Stage adding an index to an existing table (for :meth:`alter_table`)."""
        self._add_indexes.append(IndexDef(columns=columns, unique=unique, using=using, name=name))
        return self

    def drop_index(self, name: str) -> "MigrationBuilder":
        """Stage dropping an index by name (for :meth:`alter_table`)."""
        self._drop_indexes.append(name)
        return self

    # ------------------------------------------------------------------
    # Terminal methods — produce Op objects
    # ------------------------------------------------------------------

    def create_table(self) -> CreateTableOp:
        """Return a :class:`CreateTableOp` for the columns defined on this builder."""
        return CreateTableOp(
            table_name=self._table_name,
            columns=list(self._columns),
            indexes=list(self._indexes),
            foreign_keys=list(self._foreign_keys),
            dialect=self._dialect,
            if_not_exists=self._if_not_exists,
            doris_key_type=self._doris_key_type,
            doris_key_columns=list(self._doris_key_columns),
            doris_distributed_column=self._doris_distributed_column,
            doris_buckets=self._doris_buckets,
            doris_properties=dict(self._doris_properties),
            doris_partition_column=self._doris_partition_column,
        )

    def drop_table(self) -> DropTableOp:
        """Return a :class:`DropTableOp` for this table."""
        return DropTableOp(
            table_name=self._table_name,
            dialect=self._dialect,
            if_exists=self._if_not_exists,
        )

    def rename_table(self, new_name: str) -> RenameTableOp:
        """Return a :class:`RenameTableOp` that renames this table to *new_name*."""
        return RenameTableOp(
            old_name=self._table_name,
            new_name=new_name,
            dialect=self._dialect,
        )

    def alter_table(self) -> AlterTableOp:
        """Return an :class:`AlterTableOp` from the alter operations staged on this builder."""
        return AlterTableOp(
            table_name=self._table_name,
            add_columns=list(self._add_columns),
            drop_columns=list(self._drop_columns),
            rename_columns=list(self._rename_columns),
            change_column_types=list(self._change_column_types),
            add_indexes=list(self._add_indexes),
            drop_indexes=list(self._drop_indexes),
            dialect=self._dialect,
        )
