"""Tests for database setup utilities."""

from unittest.mock import MagicMock, patch

import pytest

from preloop.models.db.setup import reset_database, setup_database


@pytest.fixture
def mock_engine():
    """Mock SQLAlchemy engine."""
    mock = MagicMock()
    mock.connect.return_value.__enter__.return_value = MagicMock()
    return mock


@pytest.fixture
def mock_subprocess_run():
    """Mock subprocess.run for Alembic commands."""
    with patch("subprocess.run") as mock_run:
        mock_run.return_value = MagicMock(returncode=0, stdout="", stderr="")
        yield mock_run


def test_setup_database_success(mock_engine, mock_subprocess_run):
    """Test successful database setup using Alembic."""
    with patch("preloop.models.db.setup.get_engine", return_value=mock_engine):
        setup_database()

        # Verify that Alembic upgrade was called
        assert mock_subprocess_run.called
        call_args = mock_subprocess_run.call_args
        assert call_args[0][0] == ["alembic", "upgrade", "head"]


def test_setup_database_with_pgvector(mock_engine, mock_subprocess_run):
    """Test database setup with pgvector extension."""
    with patch("preloop.models.db.setup.get_engine", return_value=mock_engine):
        setup_database(database_url="postgresql://localhost/test")

        # Verify that Alembic upgrade was called
        assert mock_subprocess_run.called

        # Verify that the pgvector extension command was executed
        mock_conn = mock_engine.connect.return_value.__enter__.return_value
        mock_conn.execute.assert_called_once()
        mock_conn.commit.assert_called_once()


def test_setup_database_error(mock_engine, mock_subprocess_run):
    """Test database setup error handling when Alembic fails."""
    mock_subprocess_run.return_value = MagicMock(
        returncode=1, stdout="", stderr="Alembic error"
    )

    with patch("preloop.models.db.setup.get_engine", return_value=mock_engine):
        with pytest.raises(RuntimeError, match="Alembic migration failed"):
            setup_database()


def test_reset_database_success(mock_engine, mock_subprocess_run):
    """Test successful database reset using Alembic."""
    with (
        patch("preloop.models.db.setup.get_engine", return_value=mock_engine),
        patch("preloop.models.db.setup.setup_database") as mock_setup,
    ):
        reset_database()

        # Verify that Alembic downgrade was called
        assert mock_subprocess_run.called
        call_args = mock_subprocess_run.call_args
        assert call_args[0][0] == ["alembic", "downgrade", "base"]

        # Verify that setup_database was called
        mock_setup.assert_called_once_with(None)


def test_reset_database_error(mock_engine, mock_subprocess_run):
    """Test database reset error handling when Alembic downgrade fails."""
    mock_subprocess_run.return_value = MagicMock(
        returncode=1, stdout="", stderr="Alembic error"
    )

    with patch("preloop.models.db.setup.get_engine", return_value=mock_engine):
        with pytest.raises(RuntimeError, match="Alembic downgrade failed"):
            reset_database()
