"""PostgreSQL checkpointer utilities for LangGraph memory."""

import logging
from typing import Optional

from langgraph.checkpoint.postgres import PostgresSaver
from psycopg_pool import AsyncConnectionPool
from modules.shared.configuration import APP_CONFIG

logger = logging.getLogger(__name__)

# Global checkpointer instance
_checkpointer_instance: Optional[PostgresSaver] = None
_connection_pool: Optional[AsyncConnectionPool] = None


async def initialize_checkpointer() -> None:
    """Initialize the PostgreSQL checkpointer for LangGraph.

    This should be called during application startup.
    Creates a connection pool and PostgresSaver instance.
    """
    global _checkpointer_instance, _connection_pool

    if _checkpointer_instance is not None:
        logger.info("Checkpointer already initialized")
        return

    try:
        # Get database configuration from environment
        host = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_HOST", "localhost")
        database = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_DATABASE", "poweron_chat")
        user = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_USER", "poweron_dev")
        password = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_PASSWORD_SECRET")
        port = APP_CONFIG.get("LANGGRAPH_CHECKPOINT_DB_PORT", "5432")

        # Build connection string
        connection_string = f"postgresql://{user}:{password}@{host}:{port}/{database}"

        # Create async connection pool
        _connection_pool = AsyncConnectionPool(
            conninfo=connection_string,
            min_size=2,
            max_size=10,
        )

        # Initialize the connection pool
        await _connection_pool.open()

        # Create PostgresSaver with the pool
        _checkpointer_instance = PostgresSaver(_connection_pool)

        # Setup the checkpointer (creates tables if needed)
        await _checkpointer_instance.setup()

        logger.info("PostgreSQL checkpointer initialized successfully")

    except Exception as e:
        logger.error(f"Failed to initialize PostgreSQL checkpointer: {str(e)}")
        raise


async def close_checkpointer() -> None:
    """Close the checkpointer and connection pool.

    This should be called during application shutdown.
    """
    global _checkpointer_instance, _connection_pool

    if _connection_pool is not None:
        try:
            await _connection_pool.close()
            logger.info("PostgreSQL checkpointer connection pool closed")
        except Exception as e:
            logger.error(f"Error closing checkpointer connection pool: {str(e)}")

    _checkpointer_instance = None
    _connection_pool = None


def get_checkpointer() -> PostgresSaver:
    """Get the global PostgreSQL checkpointer instance.

    Returns:
        The initialized PostgresSaver instance

    Raises:
        RuntimeError: If checkpointer is not initialized
    """
    if _checkpointer_instance is None:
        raise RuntimeError(
            "PostgreSQL checkpointer not initialized. "
            "Call initialize_checkpointer() during application startup."
        )
    return _checkpointer_instance