"""Test the database management functions."""

import unittest

import psycopg

from pum.database import (
    configure_database_connect_access,
    create_database,
    drop_database,
    get_database_connect_access,
)


class TestDatabase(unittest.TestCase):
    """Test the database create/drop functions."""

    def setUp(self):
        """Set up test fixtures."""
        # Use a pg_service name for connection
        self.pg_service = "pum_test"
        self.test_db = "test_database_create_drop"
        self.test_db_template = "test_database_template"
        self.test_db_from_template = "test_database_from_template"
        self.allowed_role = "test_database_connect_allowed"
        self.blocked_role = "test_database_connect_blocked"
        self.intruder_role = "test_database_connect_intruder"

        # Ensure test databases don't exist
        self._cleanup_db(self.test_db)
        self._cleanup_db(self.test_db_template)
        self._cleanup_db(self.test_db_from_template)
        self._cleanup_role(self.allowed_role)
        self._cleanup_role(self.blocked_role)
        self._cleanup_role(self.intruder_role)

    def tearDown(self):
        """Clean up test fixtures."""
        self._cleanup_db(self.test_db)
        self._cleanup_db(self.test_db_template)
        self._cleanup_db(self.test_db_from_template)
        self._cleanup_role(self.allowed_role)
        self._cleanup_role(self.blocked_role)
        self._cleanup_role(self.intruder_role)

    def _cleanup_db(self, dbname: str):
        """Drop a database if it exists."""
        try:
            with psycopg.connect(
                f"service={self.pg_service} dbname=postgres", autocommit=True
            ) as conn:
                with conn.cursor() as cur:
                    # Terminate connections
                    cur.execute(
                        "SELECT pg_terminate_backend(pid) "
                        "FROM pg_stat_activity "
                        "WHERE datname = %s AND pid <> pg_backend_pid()",
                        [dbname],
                    )
                    cur.execute(f"DROP DATABASE IF EXISTS {dbname}")
        except psycopg.OperationalError:
            # If we can't connect, that's fine - test database might not exist
            pass

    def _cleanup_role(self, role_name: str):
        """Drop a role if it exists."""
        try:
            with psycopg.connect(
                f"service={self.pg_service} dbname=postgres", autocommit=True
            ) as conn:
                with conn.cursor() as cur:
                    cur.execute(f"DROP ROLE IF EXISTS {role_name}")
        except psycopg.OperationalError:
            pass

    def test_create_database(self):
        """Test creating a database."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        # Create the database
        create_database(connection_params, self.test_db)

        # Verify it exists
        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT 1 FROM pg_database WHERE datname = %s",
                    [self.test_db],
                )
                result = cur.fetchone()
                self.assertIsNotNone(result)

    def test_drop_database(self):
        """Test dropping a database."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        # First create the database
        with psycopg.connect(f"service={self.pg_service} dbname=postgres", autocommit=True) as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE DATABASE {self.test_db}")

        # Verify it exists
        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT 1 FROM pg_database WHERE datname = %s",
                    [self.test_db],
                )
                result = cur.fetchone()
                self.assertIsNotNone(result)

        # Drop the database
        drop_database(connection_params, self.test_db)

        # Verify it's gone
        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT 1 FROM pg_database WHERE datname = %s",
                    [self.test_db],
                )
                result = cur.fetchone()
                self.assertIsNone(result)

    def test_drop_database_with_active_connections(self):
        """Test dropping a database that has active connections."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        # Create the database
        with psycopg.connect(f"service={self.pg_service} dbname=postgres", autocommit=True) as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE DATABASE {self.test_db}")

        # Open a connection to the test database (active connection)
        active_conn = psycopg.connect(f"service={self.pg_service} dbname={self.test_db}")

        try:
            # Drop should still work (it terminates connections first)
            drop_database(connection_params, self.test_db)

            # Verify it's gone
            with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        "SELECT 1 FROM pg_database WHERE datname = %s",
                        [self.test_db],
                    )
                    result = cur.fetchone()
                    self.assertIsNone(result)
        finally:
            # Close the active connection (it should already be terminated)
            try:
                active_conn.close()
            except Exception:
                pass

    def test_create_database_from_template(self):
        """Test creating a database from a template."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        # Create template database and add some data
        with psycopg.connect(f"service={self.pg_service} dbname=postgres", autocommit=True) as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE DATABASE {self.test_db_template}")

        # Add a table to the template
        with psycopg.connect(f"service={self.pg_service} dbname={self.test_db_template}") as conn:
            with conn.cursor() as cur:
                cur.execute("CREATE TABLE test_table (id serial PRIMARY KEY, name text)")
                cur.execute("INSERT INTO test_table (name) VALUES ('template_data')")
                conn.commit()

        # Create new database from template
        create_database(
            connection_params, self.test_db_from_template, template=self.test_db_template
        )

        # Verify the new database has the template's structure and data
        with psycopg.connect(
            f"service={self.pg_service} dbname={self.test_db_from_template}"
        ) as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT name FROM test_table")
                result = cur.fetchone()
                self.assertIsNotNone(result)
                self.assertEqual(result[0], "template_data")

    def test_create_database_idempotent(self):
        """Test that creating a database twice doesn't fail."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        # Create the database
        create_database(connection_params, self.test_db)

        # Creating it again should raise an error (we don't have IF NOT EXISTS)
        with self.assertRaises(psycopg.errors.DuplicateDatabase):
            create_database(connection_params, self.test_db)

    def test_drop_nonexistent_database(self):
        """Test dropping a database that doesn't exist."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        # Drop a non-existent database - should not raise error due to IF EXISTS
        # Actually, DROP DATABASE IF EXISTS is not used in the implementation
        # So this will fail. Let's check the implementation...
        # Looking at database.py, it doesn't use IF EXISTS, so this will raise an error
        with self.assertRaises(psycopg.errors.InvalidCatalogName):
            drop_database(connection_params, "nonexistent_database_123456")

    def test_configure_database_connect_access(self):
        """Test configuring CONNECT access on an existing database."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        create_database(connection_params, self.test_db)

        with psycopg.connect(f"service={self.pg_service} dbname=postgres", autocommit=True) as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE ROLE {self.allowed_role}")
                cur.execute(f"CREATE ROLE {self.blocked_role}")

        configure_database_connect_access(
            connection_params,
            self.test_db,
            grant_roles=[self.allowed_role],
            public=False,
        )

        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT has_database_privilege(%s, %s, 'CONNECT')",
                    [self.allowed_role, self.test_db],
                )
                self.assertTrue(cur.fetchone()[0])

                cur.execute(
                    "SELECT has_database_privilege(%s, %s, 'CONNECT')",
                    [self.blocked_role, self.test_db],
                )
                self.assertFalse(cur.fetchone()[0])

    def test_configure_database_connect_access_revokes_explicit_role(self):
        """Test revoking explicit CONNECT grants from selected roles."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}

        create_database(connection_params, self.test_db)

        with psycopg.connect(f"service={self.pg_service} dbname=postgres", autocommit=True) as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE ROLE {self.allowed_role}")
                cur.execute(f"CREATE ROLE {self.intruder_role}")
                cur.execute(f"GRANT CONNECT ON DATABASE {self.test_db} TO {self.intruder_role}")

        configure_database_connect_access(
            connection_params,
            self.test_db,
            revoke_roles=[self.intruder_role],
            grant_roles=[self.allowed_role],
            public=False,
        )

        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT has_database_privilege(%s, %s, 'CONNECT')",
                    [self.intruder_role, self.test_db],
                )
                self.assertFalse(cur.fetchone()[0])

                cur.execute(
                    "SELECT has_database_privilege(%s, %s, 'CONNECT')",
                    [self.allowed_role, self.test_db],
                )
                self.assertTrue(cur.fetchone()[0])

    def test_db2_only_role_cannot_connect_to_db1(self):
        """Test CONNECT isolation across databases for a DB-specific role."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}
        db1 = "test_database_connect_db1"
        db2 = "test_database_connect_db2"
        db2_role = "test_database_connect_db2_user"

        self._cleanup_db(db1)
        self._cleanup_db(db2)
        self._cleanup_role(db2_role)

        try:
            create_database(connection_params, db1)
            create_database(connection_params, db2)

            with psycopg.connect(
                f"service={self.pg_service} dbname=postgres", autocommit=True
            ) as conn:
                with conn.cursor() as cur:
                    cur.execute(f"CREATE ROLE {db2_role}")

            configure_database_connect_access(
                connection_params,
                db1,
                grant_roles=[],
                public=False,
            )
            configure_database_connect_access(
                connection_params,
                db2,
                grant_roles=[db2_role],
                public=False,
            )

            with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
                with conn.cursor() as cur:
                    cur.execute(
                        "SELECT has_database_privilege(%s, %s, 'CONNECT')",
                        [db2_role, db1],
                    )
                    self.assertFalse(cur.fetchone()[0])

                    cur.execute(
                        "SELECT has_database_privilege(%s, %s, 'CONNECT')",
                        [db2_role, db2],
                    )
                    self.assertTrue(cur.fetchone()[0])
        finally:
            self._cleanup_db(db1)
            self._cleanup_db(db2)

    def test_get_database_connect_access_default(self):
        """Test get_database_connect_access on a freshly created database (default ACL)."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}
        create_database(connection_params, self.test_db)

        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            public_has_connect, roles = get_database_connect_access(conn, self.test_db)

        self.assertTrue(public_has_connect)
        self.assertEqual(roles, [])

    def test_get_database_connect_access_after_revoke_public(self):
        """Test get_database_connect_access reflects revoked PUBLIC and explicit grants."""
        connection_params = {"service": self.pg_service, "dbname": "postgres"}
        create_database(connection_params, self.test_db)

        with psycopg.connect(f"service={self.pg_service} dbname=postgres", autocommit=True) as conn:
            with conn.cursor() as cur:
                cur.execute(f"CREATE ROLE {self.allowed_role}")

        configure_database_connect_access(
            connection_params,
            self.test_db,
            grant_roles=[self.allowed_role],
            public=False,
        )

        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            public_has_connect, roles = get_database_connect_access(conn, self.test_db)

        self.assertFalse(public_has_connect)
        self.assertIn(self.allowed_role, roles)

    def test_get_database_connect_access_nonexistent_database(self):
        """Test get_database_connect_access raises ValueError for unknown database."""
        with psycopg.connect(f"service={self.pg_service} dbname=postgres") as conn:
            with self.assertRaises(ValueError):
                get_database_connect_access(conn, "nonexistent_database_xyz_123")


if __name__ == "__main__":
    unittest.main()
