"""Database connection management using asyncpg."""

from __future__ import annotations

import asyncio
import logging
from contextlib import asynccontextmanager
from pathlib import Path
from typing import TYPE_CHECKING, Any, AsyncGenerator

import asyncpg

if TYPE_CHECKING:
    from reviewiq_pipeline.config import Config

logger = logging.getLogger(__name__)


class DatabasePool:
    """Manages an asyncpg connection pool."""

    def __init__(self, config: Config):
        self.config = config
        self._pool: asyncpg.Pool | None = None
        self._lock = asyncio.Lock()

    async def initialize(self) -> None:
        """Initialize the connection pool."""
        async with self._lock:
            if self._pool is not None:
                return

            logger.info("Creating database connection pool...")
            self._pool = await asyncpg.create_pool(
                self.config.database_url,
                min_size=self.config.db_pool_min_size,
                max_size=self.config.db_pool_max_size,
                command_timeout=60,
            )
            logger.info("Database pool created successfully")

    async def close(self) -> None:
        """Close the connection pool."""
        async with self._lock:
            if self._pool is not None:
                await self._pool.close()
                self._pool = None
                logger.info("Database pool closed")

    @property
    def pool(self) -> asyncpg.Pool:
        """Get the connection pool, raising if not initialized."""
        if self._pool is None:
            raise RuntimeError("Database pool not initialized. Call initialize() first.")
        return self._pool

    @asynccontextmanager
    async def acquire(self) -> AsyncGenerator[asyncpg.Connection, None]:
        """Acquire a connection from the pool."""
        async with self.pool.acquire() as conn:
            yield conn

    @asynccontextmanager
    async def transaction(self) -> AsyncGenerator[asyncpg.Connection, None]:
        """Acquire a connection and start a transaction."""
        async with self.pool.acquire() as conn:
            async with conn.transaction():
                yield conn

    async def execute(self, query: str, *args: Any) -> str:
        """Execute a query and return the status string."""
        async with self.acquire() as conn:
            return await conn.execute(query, *args)

    async def executemany(self, query: str, args: list[tuple]) -> None:
        """Execute a query with multiple argument sets."""
        async with self.acquire() as conn:
            await conn.executemany(query, args)

    async def fetch(self, query: str, *args: Any) -> list[asyncpg.Record]:
        """Fetch multiple rows."""
        async with self.acquire() as conn:
            return await conn.fetch(query, *args)

    async def fetchrow(self, query: str, *args: Any) -> asyncpg.Record | None:
        """Fetch a single row."""
        async with self.acquire() as conn:
            return await conn.fetchrow(query, *args)

    async def fetchval(self, query: str, *args: Any) -> Any:
        """Fetch a single value."""
        async with self.acquire() as conn:
            return await conn.fetchval(query, *args)

    async def run_migrations(self, migrations_path: str | None = None) -> int:
        """Run all pending migrations.

        Args:
            migrations_path: Path to migrations directory. Uses config default if None.

        Returns:
            Number of migrations run.
        """
        path = Path(migrations_path or self.config.effective_migrations_path)
        if not path.exists():
            logger.warning(f"Migrations path does not exist: {path}")
            return 0

        async with self.transaction() as conn:
            # Create migrations tracking table
            await conn.execute("""
                CREATE TABLE IF NOT EXISTS _migrations (
                    id SERIAL PRIMARY KEY,
                    filename VARCHAR(255) UNIQUE NOT NULL,
                    applied_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
                )
            """)

            # Get already applied migrations
            applied = await conn.fetch("SELECT filename FROM _migrations")
            applied_set = {r["filename"] for r in applied}

            # Find and run pending migrations
            migration_files = sorted(path.glob("*.sql"))
            migrations_run = 0

            for migration_file in migration_files:
                filename = migration_file.name
                if filename in applied_set:
                    continue

                logger.info(f"Running migration: {filename}")
                sql = migration_file.read_text()

                try:
                    await conn.execute(sql)
                    await conn.execute(
                        "INSERT INTO _migrations (filename) VALUES ($1)",
                        filename,
                    )
                    migrations_run += 1
                    logger.info(f"Migration {filename} applied successfully")
                except Exception as e:
                    logger.error(f"Migration {filename} failed: {e}")
                    raise

            logger.info(f"Ran {migrations_run} migrations")
            return migrations_run

    async def check_connection(self) -> bool:
        """Check if the database connection is working."""
        try:
            result = await self.fetchval("SELECT 1")
            return result == 1
        except Exception as e:
            logger.error(f"Database connection check failed: {e}")
            return False
