"""Connection pool management for local PostgreSQL and remote databases."""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
from typing import Any

from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy import text

from config import DatabaseSettings


class ConnectionManager:
    """Manages async connection pool to local PostgreSQL.

    Remote database connections (SQL Server, HANA) are handled via pyodbc
    in thread pool — see remote_executor.py.
    """

    def __init__(self, db_settings: DatabaseSettings) -> None:
        self._settings = db_settings
        self._engine: AsyncEngine | None = None
        self._session_factory: sessionmaker | None = None

    async def initialize(self) -> None:
        self._engine = create_async_engine(
            self._settings.dsn,
            pool_size=self._settings.pool_size,
            max_overflow=self._settings.max_overflow,
            pool_pre_ping=True,
            pool_recycle=3600,
        )
        self._session_factory = sessionmaker(
            self._engine, class_=AsyncSession, expire_on_commit=False
        )

    @property
    def engine(self) -> AsyncEngine:
        if self._engine is None:
            raise RuntimeError("ConnectionManager not initialized. Call initialize() first.")
        return self._engine

    def session(self) -> AsyncSession:
        if self._session_factory is None:
            raise RuntimeError("ConnectionManager not initialized. Call initialize() first.")
        return self._session_factory()

    async def execute(self, query: str, params: dict[str, Any] | None = None) -> list[dict]:
        """Execute a read-only query against local PostgreSQL and return rows as dicts."""
        async with self.session() as session:
            result = await session.execute(text(query), params or {})
            columns = list(result.keys())
            return [dict(zip(columns, row)) for row in result.fetchall()]

    async def execute_write(self, query: str, params: dict[str, Any] | None = None) -> None:
        """Execute a write query (INSERT/UPDATE/DELETE) against local PostgreSQL."""
        async with self.session() as session:
            await session.execute(text(query), params or {})
            await session.commit()

    async def execute_many(self, query: str, params_list: list[dict[str, Any]]) -> None:
        """Execute a parameterized query with multiple parameter sets."""
        async with self.session() as session:
            for params in params_list:
                await session.execute(text(query), params)
            await session.commit()

    async def fetch_one(self, query: str, params: dict[str, Any] | None = None) -> dict | None:
        """Execute query and return a single row as dict, or None.

        Automatically commits if the query is a write operation (INSERT/UPDATE/DELETE).
        """
        async with self.session() as session:
            result = await session.execute(text(query), params or {})
            row = result.fetchone()
            # Commit if this is a write query (INSERT/UPDATE/DELETE with RETURNING)
            stripped = query.strip().upper()
            if stripped.startswith(("INSERT", "UPDATE", "DELETE")):
                await session.commit()
            if row is None:
                return None
            columns = list(result.keys())
            return dict(zip(columns, row))

    async def table_exists(self, schema: str, table: str) -> bool:
        row = await self.fetch_one(
            "SELECT 1 FROM information_schema.tables WHERE table_schema = :schema AND table_name = :table",
            {"schema": schema, "table": table},
        )
        return row is not None

    @asynccontextmanager
    async def raw_asyncpg(self):
        """Yield a raw asyncpg connection from the pool for COPY operations."""
        async with self.engine.connect() as conn:
            raw = await conn.get_raw_connection()
            yield raw.driver_connection

    async def close(self) -> None:
        if self._engine:
            await self._engine.dispose()
