from __future__ import annotations

import os
import re
from collections.abc import Callable
from dataclasses import dataclass
from urllib.parse import urlparse

_BASE_TEST_DB_NAME = "helix_test"


def _resolve_test_db_name() -> str:
    """Append xdist worker id so parallel workers get isolated databases.

    pytest-xdist sets ``PYTEST_XDIST_WORKER`` to ``gw0``/``gw1``/... per
    worker, and leaves it unset when running single-process. Keeping the
    base name when the env var is missing preserves the current local-dev
    behaviour (``helix_test``) so scripts and IDE runners don't break.
    """
    worker = os.environ.get("PYTEST_XDIST_WORKER", "")
    if not worker:
        return _BASE_TEST_DB_NAME
    safe = re.sub(r"[^a-zA-Z0-9_]", "_", worker)
    return f"{_BASE_TEST_DB_NAME}_{safe}"


TEST_DB_NAME = _resolve_test_db_name()


@dataclass(frozen=True)
class TestDatabaseConfig:
    backend: str
    async_url: str


def _replace_database_name(url: str, database_name: str) -> str:
    return url.rsplit("/", 1)[0] + f"/{database_name}"


def _parse_mysql_connection(url: str) -> tuple[str, int, str, str]:
    normalized = (
        url.replace("mysql+pymysql://", "mysql://")
        .replace("mysql+aiomysql://", "mysql://")
        .replace("mysql+asyncmy://", "mysql://")
    )
    parsed = urlparse(normalized)

    user = parsed.username or ""
    password = parsed.password or ""
    host = parsed.hostname or "localhost"
    port = parsed.port or 3306
    return host, port, user, password


def _parse_mysql_user(url: str) -> tuple[str, str]:
    normalized = (
        url.replace("mysql+pymysql://", "mysql://")
        .replace("mysql+aiomysql://", "mysql://")
        .replace("mysql+asyncmy://", "mysql://")
    )
    parsed = urlparse(normalized)
    return parsed.username or "", parsed.password or ""


def mysql_database_exists(url: str) -> bool:
    import pymysql

    parsed = urlparse(url.replace("mysql+pymysql://", "mysql://"))
    database = parsed.path.lstrip("/")

    if not database:
        return False

    try:
        conn = pymysql.connect(
            host=parsed.hostname or "localhost",
            port=parsed.port or 3306,
            user=parsed.username or "",
            password=parsed.password or "",
            database=database,
        )
    except pymysql.MySQLError:
        return False

    conn.close()
    return True


def provision_mysql_database(admin_url: str, database_name: str, access_url: str) -> None:
    import pymysql

    host, port, user, password = _parse_mysql_connection(admin_url)
    target_user, target_password = _parse_mysql_user(access_url)
    conn = pymysql.connect(host=host, port=port, user=user, password=password)
    try:
        with conn.cursor() as cur:
            cur.execute(
                f"CREATE DATABASE IF NOT EXISTS `{database_name}` "
                "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
            )
            if target_user:
                try:
                    cur.execute(
                        "CREATE USER IF NOT EXISTS %s@'%%' IDENTIFIED BY %s",
                        (target_user, target_password),
                    )
                except Exception:
                    # The user may already exist with a different auth setup.
                    pass
                cur.execute(
                    f"GRANT ALL PRIVILEGES ON `{database_name}`.* TO %s@'%%'",
                    (target_user,),
                )
        conn.commit()
    finally:
        conn.close()


def resolve_test_database_config(
    *,
    test_async_base: str,
    test_sync_base: str,
    admin_url: str | None,
    database_exists: Callable[[str], bool] = mysql_database_exists,
    provision_database: Callable[[str, str, str], None] = provision_mysql_database,
) -> TestDatabaseConfig:
    if test_async_base.startswith("sqlite") or test_sync_base.startswith("sqlite"):
        raise RuntimeError("Tests must run against MySQL/TiDB. SQLite is not supported.")

    test_async_url = _replace_database_name(test_async_base, TEST_DB_NAME)
    test_sync_url = _replace_database_name(test_sync_base, TEST_DB_NAME)

    if database_exists(test_sync_url):
        return TestDatabaseConfig(backend="mysql", async_url=test_async_url)

    admin_connection_url = admin_url or test_sync_base
    explicit_admin_url = bool(admin_url)

    try:
        provision_database(admin_connection_url, TEST_DB_NAME, test_sync_url)
    except Exception as exc:  # pragma: no cover - exact DB client exceptions vary by backend
        if explicit_admin_url:
            raise RuntimeError(
                "Failed to provision the MySQL test database via TEST_DATABASE_ADMIN_URL."
            ) from exc
        raise RuntimeError(
            "Failed to provision the MySQL/TiDB test database. "
            "Set TEST_DATABASE_ADMIN_URL to a MySQL admin connection, or point "
            "TEST_DATABASE_URL / TEST_DATABASE_URL_SYNC at a pre-created test database."
        ) from exc
    if not database_exists(test_sync_url):
        raise RuntimeError(
            "Provisioned the MySQL/TiDB test database, but the configured test user cannot access it."
        )

    return TestDatabaseConfig(backend="mysql", async_url=test_async_url)