from __future__ import annotations

import re
from dataclasses import dataclass, replace
from typing import Any
from urllib.parse import quote

import sqlglot
from sqlglot import exp
from sqlglot.errors import ParseError

from errors import ManifestError
from metadata import Column, Table, TableName
from storage.object_store import ObjectStore, join_key

SNAPSHOT_PREFIX = "snapshot"
INCREMENT_PREFIX = "increment"


SNAPSHOT_SCHEMA_SUFFIX = "-schema.sql"
SNAPSHOT_SCHEMA_CREATE_SUFFIX = "-schema-create.sql"


@dataclass(frozen=True)
class SnapshotTable:
    table: Table
    fivetran_table: str
    snapshot_file_stem: str

    @property
    def name(self) -> TableName:
        return self.table.name

    @property
    def source_schema(self) -> str:
        return self.table.source_schema

    @property
    def source_table(self) -> str:
        return self.table.source_table

    @property
    def primary_key(self) -> tuple[str, ...]:
        return self.table.primary_key

    @property
    def columns(self) -> tuple[Column, ...]:
        return self.table.columns

    @property
    def path_key(self) -> str:
        return f"{quote_object_path_segment(self.source_schema)}/{quote_object_path_segment(self.source_table)}"

    @property
    def quoted(self) -> str:
        return self.table.quoted

    @property
    def column_names(self) -> list[str]:
        return self.table.column_names

    def column(self, name: str) -> Column:
        return self.table.column(name)


def fivetran_table_name(source_schema: str, source_table: str) -> str:
    return re.sub(r"[^0-9A-Za-z_]+", "_", f"{source_schema}_{source_table}").strip("_")


def quote_object_path_segment(value: str) -> str:
    encoded = quote(value, safe="")
    if encoded in {".", ".."}:
        return encoded.replace(".", "%2E")
    return encoded


def parse_create_table_sql(
    sql: str,
    *,
    source_schema: str,
    source_table: str,
    snapshot_file_stem: str | None = None,
) -> SnapshotTable:
    create = _parse_create_table(sql)
    schema = create.this
    if not isinstance(schema, exp.Schema):
        raise ManifestError("CREATE TABLE does not contain a column schema")
    table_name = schema.this.name if schema.this is not None else ""
    if table_name and table_name != source_table:
        raise ManifestError(f"CREATE TABLE name {table_name!r} does not match snapshot file table {source_table!r}")

    columns: list[Column] = []
    primary_key: list[str] = []

    for item in schema.expressions:
        if isinstance(item, exp.ColumnDef):
            column = _column_from_sqlglot(item)
            columns.append(column)
            if _column_has_constraint(item, exp.PrimaryKeyColumnConstraint):
                primary_key.append(column.name)
            continue
        if isinstance(item, exp.PrimaryKey):
            primary_key = [identifier.name for identifier in item.expressions]

    if not columns:
        raise ManifestError(f"CREATE TABLE for {source_schema}.{source_table} has no columns")
    column_names = {column.name for column in columns}
    missing_pk = [name for name in primary_key if name not in column_names]
    if missing_pk:
        raise ManifestError(f"primary key columns missing from schema: {missing_pk}")

    columns = [replace(column, primary_key=column.name in primary_key) for column in columns]

    return SnapshotTable(
        table=Table(
            name=TableName(source_schema=source_schema, source_table=source_table),
            primary_key=tuple(primary_key),
            columns=tuple(columns),
        ),
        fivetran_table=fivetran_table_name(source_schema, source_table),
        snapshot_file_stem=snapshot_file_stem or f"{source_schema}.{source_table}",
    )


def _parse_create_table(sql: str) -> exp.Create:
    try:
        expressions = sqlglot.parse(sql, read="mysql")
    except ParseError as exc:
        raise ManifestError(f"failed to parse snapshot schema SQL: {exc}") from exc
    for expression in expressions:
        if isinstance(expression, exp.Create) and str(expression.args.get("kind")).upper() == "TABLE":
            return expression
    raise ManifestError("snapshot schema file does not contain CREATE TABLE")


def _column_from_sqlglot(column: exp.ColumnDef) -> Column:
    name = column.this.name
    kind = column.args.get("kind")
    if not isinstance(kind, exp.DataType):
        raise ManifestError(f"column {name!r} is missing a type")
    source_type = kind.sql(dialect="mysql")
    precision, scale = _column_precision_scale(kind)
    return Column(
        name=name,
        source_type=source_type,
        nullable=not _column_has_constraint(column, exp.NotNullColumnConstraint),
        precision=precision,
        scale=scale,
    )


def _column_precision_scale(kind: exp.DataType) -> tuple[int | None, int | None]:
    base = kind.this
    if base in {exp.DataType.Type.BIT, exp.DataType.Type.TINYINT}:
        params = _data_type_int_params(kind)
        return (params[0] if params else None), None
    if base not in {exp.DataType.Type.DECIMAL, exp.DataType.Type.DECIMAL32, exp.DataType.Type.DECIMAL64}:
        return None, None
    params = _data_type_int_params(kind)
    precision = params[0] if params else 10
    scale = params[1] if len(params) > 1 else 0
    return precision, scale


def _data_type_int_params(kind: exp.DataType) -> tuple[int, ...]:
    params: list[int] = []
    for expression in kind.expressions:
        value = expression.this
        if isinstance(value, exp.Literal):
            params.append(int(value.this))
    return tuple(params)


def _column_has_constraint(
    column: exp.ColumnDef,
    constraint_type: type[Any],
) -> bool:
    for constraint in column.args.get("constraints") or []:
        if isinstance(constraint, exp.ColumnConstraint) and isinstance(constraint.args.get("kind"), constraint_type):
            return True
    return False


def snapshot_metadata_pos(store: ObjectStore, snapshot_prefix: str) -> int:
    key = join_key(snapshot_prefix, "metadata")
    raw = store.read_text(key)
    match = re.search(r"^\s*Pos:\s*(\d+)\s*$", raw, flags=re.MULTILINE)
    if match is None:
        raise ManifestError(f"snapshot metadata is missing Pos: {key}")
    return int(match.group(1))


__all__ = [
    "SnapshotTable",
    "fivetran_table_name",
    "parse_create_table_sql",
    "snapshot_metadata_pos",
]