"""Regression tests for ProxyExtrasDBManager v2 migration resolver.

The v2 resolver is opt-in via `--use_v2_migration_resolver` / the
`use_v2_resolver=True` kwarg. These tests exercise the v2 path; the v1
(default) behavior is unchanged from pre-fix.
"""

import subprocess
from unittest.mock import patch

import pytest

from litellm_proxy_extras.utils import (
    ProxyExtrasDBManager,
    _max_migration_timestamp,
    _migration_timestamp,
)


def _fake_migrate_deploy_failure(returncode: int, stderr: str):
    def _run(*args, **kwargs):
        raise subprocess.CalledProcessError(
            returncode=returncode,
            cmd=args[0],
            stderr=stderr,
            output="",
        )

    return _run


def test_v2_p3018_permission_error_raises_runtime_error(monkeypatch, tmp_path):
    """v2: a permission failure during migrate deploy raises RuntimeError."""
    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
    monkeypatch.setattr(
        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")

    stderr = (
        "Error: P3018\nMigration name: 20250326162113_baseline\n"
        "Database error code: 42501\npermission denied for schema public"
    )
    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
        with pytest.raises(RuntimeError, match="permission"):
            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)


def test_v2_non_idempotent_p3009_raises_runtime_error(monkeypatch, tmp_path):
    """v2: a non-idempotent migration failure raises (no silent recovery)."""
    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
    monkeypatch.setattr(
        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")

    stderr = (
        "Error: P3009\nMigration `20260101000000_genuinely_broken` failed\n"
        'Reason: syntax error at or near "BRKN" LINE 42'
    )
    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
        with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)


def test_strip_prisma_query_params_removes_connection_limit():
    """DATABASE_URLs with Prisma-specific params should be parseable by psycopg."""
    url = "postgresql://u:p@h:5432/db?connection_limit=100&pool_timeout=60&sslmode=require"
    stripped = ProxyExtrasDBManager._strip_prisma_query_params(url)
    assert "connection_limit" not in stripped
    assert "pool_timeout" not in stripped
    assert "sslmode=require" in stripped


def test_strip_prisma_query_params_passthrough_no_query():
    """URLs without query strings are returned unchanged."""
    url = "postgresql://u:p@h:5432/db"
    assert ProxyExtrasDBManager._strip_prisma_query_params(url) == url


def test_migration_timestamp_extracts_leading_digits():
    assert _migration_timestamp("20260101000000_add_foo") == 20260101000000
    assert _migration_timestamp("20250326162113_baseline") == 20250326162113


def test_migration_timestamp_returns_zero_on_malformed():
    assert _migration_timestamp("0_init") == 0
    assert _migration_timestamp("not_a_migration") == 0


def test_max_migration_timestamp():
    names = {"20250326000000_a", "20260415000000_b", "20251115000000_c"}
    assert _max_migration_timestamp(names) == 20260415000000


def test_max_migration_timestamp_empty_set():
    assert _max_migration_timestamp(set()) == 0


def test_v1_default_still_calls_resolve_all_migrations(monkeypatch, tmp_path):
    """v1 (default) continues to call _resolve_all_migrations on the happy path.

    This is the existing buggy behavior — we're not fixing it in v1, only
    offering v2 as opt-in. This test pins the default so that a future
    inadvertent default flip is caught.
    """
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")

    # Stub `prisma migrate deploy` to claim success with pending migrations
    # applied, which is the code path that triggers the legacy post-migration
    # sanity check (a call to _resolve_all_migrations).
    class FakeResult:
        stdout = "Applied migration.\n"
        stderr = ""

    def fake_run(cmd, *args, **kwargs):
        return FakeResult()

    resolve_called = {"n": 0}

    def fake_resolve(*args, **kwargs):
        resolve_called["n"] += 1

    monkeypatch.setattr("subprocess.run", fake_run)
    monkeypatch.setattr(ProxyExtrasDBManager, "_resolve_all_migrations", fake_resolve)

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True)  # v2 flag NOT set
    assert ok is True
    assert resolve_called["n"] == 1, "v1 default should still invoke the legacy path"


def test_v2_db_push_wraps_subprocess_error_as_runtime_error(monkeypatch, tmp_path):
    """v2: a failing `prisma db push` must raise RuntimeError, not leak
    CalledProcessError past proxy_cli.py's `except RuntimeError`."""
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")

    stderr = "db push error"
    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
        with pytest.raises(RuntimeError, match="prisma db push failed"):
            ProxyExtrasDBManager.setup_database(use_migrate=False, use_v2_resolver=True)


def test_v2_warn_ahead_of_head_swallows_db_errors(monkeypatch, tmp_path):
    """_warn_if_db_ahead_of_head must never raise — it's informational.

    Non-connection DB errors (e.g. InsufficientPrivilege from a user
    without SELECT on _prisma_migrations) must be caught, not propagated.
    """
    import psycopg

    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")

    class _FakeConn:
        def __enter__(self):
            return self

        def __exit__(self, *a):
            return False

        def execute(self, *a, **kw):
            # Simulate an InsufficientPrivilege (subclass of DatabaseError).
            raise psycopg.errors.InsufficientPrivilege("permission denied")

    def _fake_connect(*a, **kw):
        return _FakeConn()

    monkeypatch.setattr("psycopg.connect", _fake_connect)

    # Must not raise.
    ProxyExtrasDBManager._warn_if_db_ahead_of_head(str(tmp_path))


def test_v2_resolve_specific_migration_failure_raises_runtime_error(
    monkeypatch, tmp_path
):
    """If marking a migration as applied fails inside P3009 idempotent
    recovery, the subprocess error must be re-raised as RuntimeError so
    proxy_cli.py catches it cleanly (instead of leaking CalledProcessError)."""
    monkeypatch.setattr(
        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")
    monkeypatch.setattr(
        ProxyExtrasDBManager, "_roll_back_migration", lambda *a, **kw: None
    )

    # First call: migrate deploy -> P3009 idempotent error.
    # Recovery path tries _resolve_specific_migration; that also raises.
    def _failing_resolve(*a, **kw):
        raise subprocess.CalledProcessError(
            returncode=1,
            cmd="prisma migrate resolve --applied",
            stderr="resolve failed",
            output="",
        )

    monkeypatch.setattr(
        ProxyExtrasDBManager, "_resolve_specific_migration", _failing_resolve
    )

    stderr = (
        "Error: P3009\nMigration `20260101000000_some_migration` failed\n"
        "relation already exists"
    )
    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
        with pytest.raises(
            RuntimeError, match="Failed to mark migration .* as applied"
        ):
            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)


def test_v2_does_not_call_resolve_all_migrations(monkeypatch, tmp_path):
    """v2 must never call _resolve_all_migrations — that's the bug it fixes."""
    monkeypatch.setattr(
        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")

    class FakeResult:
        stdout = "Applied migration.\n"
        stderr = ""

    monkeypatch.setattr("subprocess.run", lambda *a, **kw: FakeResult())

    resolve_called = {"n": 0}
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_resolve_all_migrations",
        lambda *a, **kw: resolve_called.__setitem__("n", resolve_called["n"] + 1),
    )

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
    assert ok is True
    assert resolve_called["n"] == 0, "v2 must not invoke the diff-and-force recovery"


_DEADLOCK_P3018_STDERR = (
    "Error: P3018\n"
    "Migration name: 20260415120000_health_check_latest_per_model_index\n"
    "Database error code: 40P01\n"
    "deadlock detected"
)


def _stub_v2_env(monkeypatch, tmp_path):
    monkeypatch.setenv("DATABASE_URL", "postgresql://u:p@localhost:9/x")
    monkeypatch.setattr(
        ProxyExtrasDBManager, "_warn_if_db_ahead_of_head", lambda _: None
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_get_prisma_dir", lambda: str(tmp_path))
    (tmp_path / "schema.prisma").write_text("// stub")
    monkeypatch.setattr("time.sleep", lambda _: None)


def _succeed_after(failures: int, stderr: str):
    calls = {"n": 0}

    class _OkResult:
        stdout = "Applied migration.\n"
        stderr = ""

    def _run(*args, **kwargs):
        if "deploy" not in args[0]:
            return _OkResult()
        calls["n"] += 1
        if calls["n"] <= failures:
            raise subprocess.CalledProcessError(
                returncode=1, cmd=args[0], stderr=stderr, output=""
            )
        return _OkResult()

    return _run


def test_v2_p3018_deadlock_rolls_back_and_retries(monkeypatch, tmp_path):
    """v2: losing the migrate deploy deadlock race against a concurrent
    instance rolls the ledger row back and retries instead of dying."""
    _stub_v2_env(monkeypatch, tmp_path)

    rolled_back = []
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_roll_back_migration",
        lambda name: rolled_back.append(name),
    )
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_resolve_specific_migration",
        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
    )
    monkeypatch.setattr("subprocess.run", _succeed_after(1, _DEADLOCK_P3018_STDERR))

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
    assert ok is True
    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]


def test_v2_p3018_persistent_deadlock_exhausts_attempts(monkeypatch, tmp_path):
    """v2: a deadlock on every attempt still fails after the retry budget."""
    _stub_v2_env(monkeypatch, tmp_path)
    monkeypatch.setattr(ProxyExtrasDBManager, "_roll_back_migration", lambda name: None)

    with patch(
        "subprocess.run",
        side_effect=_fake_migrate_deploy_failure(1, _DEADLOCK_P3018_STDERR),
    ):
        with pytest.raises(RuntimeError, match="after 4 attempts"):
            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)


def test_v2_p3009_deadlocked_ledger_row_rolls_back_and_retries(monkeypatch, tmp_path):
    """v2: the surviving instance sees the victim's failed ledger row as P3009.
    When that row's logs show a deadlock, roll it back and retry."""
    _stub_v2_env(monkeypatch, tmp_path)

    stderr = (
        "Error: P3009\n"
        "migrate found failed migrations in the target database\n"
        "The `20260415120000_health_check_latest_per_model_index` migration "
        "started at 2026-09-01 18:46:13 UTC failed"
    )
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_failed_migration_logs",
        lambda name: "ERROR: deadlock detected\nDETAIL: Process 72 waits for ShareLock",
    )
    rolled_back = []
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_roll_back_migration",
        lambda name: rolled_back.append(name),
    )
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_resolve_specific_migration",
        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
    )
    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
    assert ok is True
    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]


def test_v2_p3009_empty_ledger_logs_rolls_back_and_retries(monkeypatch, tmp_path):
    """v2: empty failed ledger logs mean a concurrent deploy moved it on."""
    _stub_v2_env(monkeypatch, tmp_path)

    stderr = (
        "Error: P3009\n"
        "migrate found failed migrations in the target database\n"
        "The `20260415120000_health_check_latest_per_model_index` migration "
        "started at 2026-09-01 18:46:13 UTC failed"
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: "")
    rolled_back = []
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_roll_back_migration",
        lambda name: rolled_back.append(name),
    )
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_resolve_specific_migration",
        lambda name: pytest.fail("a deadlocked migration must never be marked applied"),
    )
    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
    assert ok is True
    assert rolled_back == ["20260415120000_health_check_latest_per_model_index"]


def test_v2_p3009_unreadable_ledger_still_raises(monkeypatch, tmp_path):
    """v2: an unreadable ledger cannot establish that P3009 was a deadlock."""
    _stub_v2_env(monkeypatch, tmp_path)

    stderr = (
        "Error: P3009\n"
        "migrate found failed migrations in the target database\n"
        "The `20260415120000_health_check_latest_per_model_index` migration "
        "started at 2026-09-01 18:46:13 UTC failed"
    )
    monkeypatch.setattr(ProxyExtrasDBManager, "_failed_migration_logs", lambda name: None)
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_roll_back_migration",
        lambda name: pytest.fail("an unreadable ledger must not trigger a retry"),
    )
    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))

    with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
        ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)


def test_v2_p3009_non_deadlock_ledger_row_still_raises(monkeypatch, tmp_path):
    """v2: a failed ledger row whose logs show a real SQL error stays fatal."""
    _stub_v2_env(monkeypatch, tmp_path)

    stderr = (
        "Error: P3009\n"
        "migrate found failed migrations in the target database\n"
        "The `20260101000000_genuinely_broken` migration started at "
        "2026-09-01 18:46:13 UTC failed"
    )
    monkeypatch.setattr(
        ProxyExtrasDBManager,
        "_failed_migration_logs",
        lambda name: 'ERROR: syntax error at or near "BRKN"',
    )

    with patch("subprocess.run", side_effect=_fake_migrate_deploy_failure(1, stderr)):
        with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
            ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)


def test_v2_bare_deadlock_stderr_retries(monkeypatch, tmp_path):
    """v2: a deadlock reported without a Prisma error code (the advisory-lock
    waiter as victim) is retried, not fatal."""
    _stub_v2_env(monkeypatch, tmp_path)
    monkeypatch.setattr(
        "subprocess.run", _succeed_after(1, "Database error: deadlock detected")
    )

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
    assert ok is True


_P1002_ADVISORY_LOCK_STDERR = (
    "Error: P1002\n\n"
    "The database server at `127.0.0.1`:`45743` was reached but timed out.\n\n"
    "Context: Timed out trying to acquire a postgres advisory lock "
    "(SELECT pg_advisory_lock(72707369)). Elapsed: 10000ms."
)


def test_v2_advisory_lock_timeout_retries(monkeypatch, tmp_path):
    """v2: the advisory-lock waiter that times out while a peer's retry holds
    the lock retries instead of dying."""
    _stub_v2_env(monkeypatch, tmp_path)
    monkeypatch.setattr("subprocess.run", _succeed_after(2, _P1002_ADVISORY_LOCK_STDERR))

    ok = ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)
    assert ok is True


def test_v2_p1002_without_advisory_lock_context_still_raises(monkeypatch, tmp_path):
    """v2: a plain P1002 (database unreachable) stays fatal."""
    _stub_v2_env(monkeypatch, tmp_path)
    stderr = "Error: P1002\n\nThe database server at `db`:`5432` was reached but timed out."
    monkeypatch.setattr("subprocess.run", _succeed_after(1, stderr))

    with pytest.raises(RuntimeError, match="cannot be auto-recovered"):
        ProxyExtrasDBManager.setup_database(use_migrate=True, use_v2_resolver=True)