import difflib
import logging
import re
from pathlib import Path
from typing import NamedTuple

import pytest
from sqlalchemy import create_engine, inspect
from sqlalchemy.schema import CreateTable, MetaData, UniqueConstraint

_logger = logging.getLogger(__name__)

_DIALECT_REFLECTED_UNIQUE_CONSTRAINTS = {
    "mysql": {
        "uq_experiments_workspace_name",
        "uq_endpoints_workspace_name",
        "uq_secrets_workspace_secret_name",
        "uq_model_definitions_workspace_name",
    },
    "mssql": {
        "uq_experiments_workspace_name",
        "uq_endpoints_workspace_name",
        "uq_secrets_workspace_secret_name",
        "uq_model_definitions_workspace_name",
    },
}

import mlflow
from mlflow.environment_variables import MLFLOW_TRACKING_URI

pytestmark = pytest.mark.notrackingurimock


def get_database_dialect(uri):
    return create_engine(uri).dialect.name


def get_tracking_uri():
    return MLFLOW_TRACKING_URI.get()


def dump_schema(db_uri):
    engine = create_engine(db_uri)
    created_tables_metadata = MetaData()
    created_tables_metadata.reflect(bind=engine)
    _reattach_missing_unique_constraints(engine, created_tables_metadata)
    # Write out table schema as described in
    # https://docs.sqlalchemy.org/en/13/faq/metadata_schema.html#how-can-i-get-the-create-table-drop-table-output-as-a-string
    lines = []
    for table in created_tables_metadata.sorted_tables:
        # Apply `str.rstrip` to remove trailing whitespaces
        lines += map(str.rstrip, str(CreateTable(table)).splitlines())
    return "\n".join(lines)


def _reattach_missing_unique_constraints(engine, metadata):
    constraint_names = _DIALECT_REFLECTED_UNIQUE_CONSTRAINTS.get(engine.dialect.name)
    if not constraint_names:
        return
    inspector = inspect(engine)
    for table in metadata.sorted_tables:
        existing_unique_columns = {
            tuple(constraint.columns.keys())
            for constraint in table.constraints
            if isinstance(constraint, UniqueConstraint)
        }
        # Not all dialects reflect `UniqueConstraint` objects the same way. MySQL reports
        # them as indexes; MSSQL doesn't implement `get_unique_constraints` at all. We
        # normalize the reflection results via `_get_unique_constraints` so the same code
        # path can reattach missing `UniqueConstraint`s across dialects.
        for unique in _get_unique_constraints(inspector, engine.dialect.name, table.name):
            name = unique.get("name")
            columns = tuple(unique.get("column_names") or ())
            duplicates_index = unique.get("duplicates_index")
            if not columns or name not in constraint_names:
                continue
            if engine.dialect.name == "mysql" and not duplicates_index:
                # MySQL exposes unique constraints as unique indexes. SQLAlchemy treats those as
                # indexes during reflection, so the reflected metadata lacks the original
                # `UniqueConstraint`. Only recreate constraints that are backed by an actual
                # unique index reported via `duplicates_index`.
                continue
            if columns in existing_unique_columns:
                continue
            if missing_columns := tuple(column for column in columns if column not in table.c):
                _logger.warning(
                    "Skipping recreation of unique constraint '%s' on table '%s' due to "
                    "missing columns: %s",
                    name,
                    table.name,
                    ", ".join(missing_columns),
                )
                continue
            constraint = UniqueConstraint(*[table.c[column] for column in columns], name=name)
            table.append_constraint(constraint)
            existing_unique_columns.add(columns)


def _get_unique_constraints(inspector, dialect, table_name):
    try:
        unique_constraints = inspector.get_unique_constraints(table_name)
    except NotImplementedError:
        unique_constraints = None
    if unique_constraints is None:
        unique_constraints = []
    if not unique_constraints:
        unique_constraints = [
            {
                "name": index.get("name"),
                "column_names": index.get("column_names"),
                "duplicates_index": index.get("unique"),
            }
            for index in inspector.get_indexes(table_name)
            if index.get("unique")
        ]
    return unique_constraints


class _CreateTable(NamedTuple):
    table: str
    columns: str


_CREATE_TABLE_REGEX = re.compile(
    r"""
CREATE TABLE (?P<table>\S+?) \(
(?P<columns>.+?)
\)
""".strip(),
    flags=re.DOTALL,
)


def parse_create_tables(schema):
    return [
        _CreateTable(
            table=m.group("table"),
            columns=set(m.group("columns").splitlines()),
        )
        for m in _CREATE_TABLE_REGEX.finditer(schema)
    ]


def schema_equal(schema_a, schema_b):
    create_tables_a = parse_create_tables(schema_a)
    create_tables_b = parse_create_tables(schema_b)
    assert create_tables_a != []
    assert create_tables_b != []
    return create_tables_a == create_tables_b


def get_schema_path(db_uri):
    return Path(__file__).parent / "schemas" / (get_database_dialect(db_uri) + ".sql")


def iter_parameter_sets():
    a = """
CREATE TABLE table (
    col VARCHAR(10)
)
"""
    b = """
CREATE TABLE table (
    col VARCHAR(10)
)
"""
    yield pytest.param(a, b, True, id="identical schemas")

    a = """
CREATE TABLE table1 (
    col VARCHAR(10)
)
"""
    b = """
CREATE TABLE table2 (
    col VARCHAR(10)
)
"""
    yield pytest.param(a, b, False, id="different table names")

    a = """
CREATE TABLE table (
    col1 VARCHAR(10)
)
"""
    b = """
CREATE TABLE table (
    col2 VARCHAR(10)
)
"""
    yield pytest.param(a, b, False, id="different column names")


@pytest.mark.parametrize(("a", "b", "expected"), iter_parameter_sets())
def test_schema_equal(a, b, expected):
    assert schema_equal(a, b) is expected


def initialize_database():
    with mlflow.start_run():
        pass


def get_schema_update_command(dialect):
    this_script = Path(__file__).relative_to(Path.cwd())
    docker_compose_yml = this_script.parent / "compose.yml"
    return f"docker compose -f {docker_compose_yml} run --rm mlflow-{dialect} python {this_script}"


def test_schema_is_up_to_date():
    initialize_database()
    tracking_uri = get_tracking_uri()
    schema_path = get_schema_path(tracking_uri)
    existing_schema = schema_path.read_text()
    latest_schema = dump_schema(tracking_uri)
    dialect = get_database_dialect(tracking_uri)
    update_command = get_schema_update_command(dialect)
    message = (
        f"{schema_path.relative_to(Path.cwd())} is not up-to-date. "
        f"Please run this command to update it: {update_command}"
    )
    diff = "".join(
        difflib.ndiff(
            existing_schema.splitlines(keepends=True), latest_schema.splitlines(keepends=True)
        )
    )
    rel_path = schema_path.relative_to(Path.cwd())
    message = f"""
=================================== EXPECTED ===================================
{latest_schema}
==================================== ACTUAL ====================================
{existing_schema}
===================================== DIFF =====================================
{diff}
================================== HOW TO FIX ==================================
Manually copy & paste the expected schema in {rel_path} or run the following command:
{update_command}
"""
    assert schema_equal(existing_schema, latest_schema), message


def main():
    tracking_uri = get_tracking_uri()
    assert tracking_uri, f"Environment variable {MLFLOW_TRACKING_URI} must be set"
    get_database_dialect(tracking_uri)  # Ensure `tracking_uri` is a database URI
    mlflow.set_tracking_uri(tracking_uri)
    initialize_database()
    schema_path = get_schema_path(tracking_uri)
    existing_schema = schema_path.read_text()
    latest_schema = dump_schema(tracking_uri)
    if not schema_equal(existing_schema, latest_schema):
        schema_path.write_text(latest_schema)


if __name__ == "__main__":
    main()