"""
LangGraph Checkpointer 配置模块

支持多种持久化后端：
1. MemorySaver - 内存存储（仅用于测试）
2. PostgresSaver - PostgreSQL 数据库（生产环境）

使用 message_id 和 chat_id 作为 thread_id 来恢复工作流状态
"""

import os
import asyncio
from typing import Optional
from langgraph.checkpoint.memory import MemorySaver
from config import POSTGRES_CONNECTION_STRING

CHECKPOINTER_TYPE = "postgres"


def get_thread_id(chat_id: str, message_id: str) -> str:
    """
    生成唯一的 thread_id
    
    Args:
        chat_id: 聊天会话 ID
        message_id: 消息 ID
    
    Returns:
        组合后的唯一 thread_id
    """
    return f"{chat_id}:{message_id}"


def parse_thread_id(thread_id: str) -> tuple[str, str]:
    """
    解析 thread_id 获取 chat_id 和 message_id
    
    Args:
        thread_id: 组合的 thread_id
    
    Returns:
        (chat_id, message_id) 元组
    """
    parts = thread_id.split(":", 1)
    if len(parts) == 2:
        return parts[0], parts[1]
    return parts[0], ""


async def verify_checkpointer_tables_exist() -> bool:
    """
    验证 checkpointer 所需的数据库表是否存在且结构正确
    
    Returns:
        True 如果所有必需的表都存在且结构正确，否则 False
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError:
        return False
    
    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    # 检查 checkpoints 表
                    await cur.execute("""
                        SELECT EXISTS (
                            SELECT FROM information_schema.tables 
                            WHERE table_schema = 'public' 
                            AND table_name = 'checkpoints'
                        )
                    """)
                    checkpoints_exists = (await cur.fetchone())[0]
                    
                    # 检查 checkpoint_writes 表
                    await cur.execute("""
                        SELECT EXISTS (
                            SELECT FROM information_schema.tables 
                            WHERE table_schema = 'public' 
                            AND table_name = 'checkpoint_writes'
                        )
                    """)
                    writes_exists = (await cur.fetchone())[0]
                    
                    # 检查 checkpoint_blobs 表
                    await cur.execute("""
                        SELECT EXISTS (
                            SELECT FROM information_schema.tables 
                            WHERE table_schema = 'public' 
                            AND table_name = 'checkpoint_blobs'
                        )
                    """)
                    blobs_exists = (await cur.fetchone())[0]
                    
                    if not (checkpoints_exists and writes_exists and blobs_exists):
                        return False
                    
                    # 检查 checkpoints 表的 checkpoint 字段是否为 JSONB 类型
                    await cur.execute("""
                        SELECT data_type 
                        FROM information_schema.columns 
                        WHERE table_name = 'checkpoints' 
                        AND column_name = 'checkpoint'
                    """)
                    result = await cur.fetchone()
                    if not result or result[0] != 'jsonb':
                        print(f"[Checkpointer] checkpoints.checkpoint column type is {result[0] if result else 'NULL'}, expected 'jsonb'")
                        return False
                    
                    # 检查 checkpoint_writes 表是否有 task_path 列
                    await cur.execute("""
                        SELECT EXISTS (
                            SELECT FROM information_schema.columns 
                            WHERE table_name = 'checkpoint_writes' 
                            AND column_name = 'task_path'
                        )
                    """)
                    task_path_exists = (await cur.fetchone())[0]
                    if not task_path_exists:
                        print(f"[Checkpointer] checkpoint_writes.task_path column does not exist")
                        return False
                    
                    # 检查 checkpoint_writes 表的主键是否正确
                    await cur.execute("""
                        SELECT COUNT(*) 
                        FROM information_schema.table_constraints tc
                        JOIN information_schema.key_column_usage kcu 
                        ON tc.constraint_name = kcu.constraint_name
                        WHERE tc.table_name = 'checkpoint_writes' 
                        AND tc.constraint_type = 'PRIMARY KEY'
                    """)
                    pk_count = (await cur.fetchone())[0]
                    if pk_count != 5:  # 应该有 5 个主键列
                        print(f"[Checkpointer] checkpoint_writes has {pk_count} primary key columns, expected 5")
                        return False
                    
                    return True
    except Exception as e:
        print(f"[Checkpointer] Error verifying tables: {e}")
        return False


async def init_checkpointer_db():
    """
    初始化 PostgreSQL checkpointer 数据库表
    使用 AsyncPostgresSaver 的 setup 方法创建正确的表结构
    """
    if CHECKPOINTER_TYPE != "postgres":
        return
    
    try:
        from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
        from psycopg_pool import AsyncConnectionPool
        import psycopg
    except ImportError as e:
        print(f"[Checkpointer] ERROR: Missing required package: {e}")
        print(f"[Checkpointer] Please install: pip install langgraph-checkpoint-postgres psycopg-pool")
        raise ImportError(
            "langgraph-checkpoint-postgres or psycopg-pool not installed. "
            "Run: pip install langgraph-checkpoint-postgres psycopg-pool"
        ) from e
    
    # 首先检查表是否已存在
    tables_exist = await verify_checkpointer_tables_exist()
    if tables_exist:
        print(f"[Checkpointer] Database tables already exist, skipping initialization")
        return
    
    print(f"[Checkpointer] Initializing PostgreSQL database tables...")
    
    # 尝试使用 AsyncPostgresSaver 的 setup 方法
    try:
        async with AsyncConnectionPool(
            POSTGRES_CONNECTION_STRING,
            kwargs={"autocommit": True}
        ) as pool:
            checkpointer = AsyncPostgresSaver(pool)
            await checkpointer.setup()
            print(f"[Checkpointer] AsyncPostgresSaver.setup() executed")
    except Exception as e:
        print(f"[Checkpointer] AsyncPostgresSaver.setup() failed: {e}")
    
    # 验证表是否真的被创建
    tables_exist_after = await verify_checkpointer_tables_exist()
    
    if not tables_exist_after:
        print(f"[Checkpointer] Tables not created by setup(), using direct SQL...")
        try:
            await create_checkpointer_tables_directly()
        except Exception as e2:
            print(f"[Checkpointer] Direct SQL initialization also failed: {e2}")
            raise
    else:
        print(f"[Checkpointer] PostgreSQL database tables verified successfully")


async def create_checkpointer_tables_directly():
    """
    直接通过 SQL 语句创建 checkpointer 所需的数据库表
    这是备用方案，当 AsyncPostgresSaver.setup() 失败时使用
    
    注意：LangGraph AsyncPostgresSaver 使用 JSONB 类型存储数据
    表名必须是: checkpoints, checkpoint_writes, checkpoint_blobs
    """
    try:
        from psycopg_pool import AsyncConnectionPool
    except ImportError as e:
        print(f"[Checkpointer] ERROR: Missing psycopg-pool: {e}")
        raise
    
    print(f"[Checkpointer] Creating checkpointer tables using direct SQL...")
    
    async with AsyncConnectionPool(
        POSTGRES_CONNECTION_STRING,
        kwargs={"autocommit": True}
    ) as pool:
        async with pool.connection() as conn:
            async with conn.cursor() as cur:
                # 先删除旧表（如果存在错误的表结构）
                await cur.execute("DROP TABLE IF EXISTS checkpoints CASCADE")
                await cur.execute("DROP TABLE IF EXISTS writes CASCADE")
                await cur.execute("DROP TABLE IF EXISTS checkpoint_blobs CASCADE")
                await cur.execute("DROP TABLE IF EXISTS checkpoint_writes CASCADE")
                await cur.execute("DROP TABLE IF EXISTS blobs CASCADE")
                await cur.execute("DROP TABLE IF EXISTS checkpoint_migrations CASCADE")
                print(f"[Checkpointer] Dropped existing tables")
                
                # 创建 checkpoint_migrations 表（用于跟踪迁移版本）
                await cur.execute("""
                    CREATE TABLE IF NOT EXISTS checkpoint_migrations (
                        v INTEGER PRIMARY KEY
                    )
                """)
                
                # 创建 checkpoints 表 - 使用 JSONB 类型
                await cur.execute("""
                    CREATE TABLE IF NOT EXISTS checkpoints (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        checkpoint_id TEXT NOT NULL,
                        parent_checkpoint_id TEXT,
                        type TEXT,
                        checkpoint JSONB NOT NULL DEFAULT '{}',
                        metadata JSONB NOT NULL DEFAULT '{}',
                        PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id)
                    )
                """)
                
                # 创建 checkpoint_writes 表 - 使用 BYTEA 类型
                await cur.execute("""
                    CREATE TABLE IF NOT EXISTS checkpoint_writes (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        checkpoint_id TEXT NOT NULL,
                        task_id TEXT NOT NULL,
                        task_path TEXT NOT NULL DEFAULT '',
                        idx INTEGER NOT NULL,
                        channel TEXT NOT NULL,
                        type TEXT,
                        blob BYTEA,
                        PRIMARY KEY (thread_id, checkpoint_ns, checkpoint_id, task_id, idx)
                    )
                """)
                
                # 创建 checkpoint_blobs 表 - 使用 BYTEA 类型存储二进制数据
                await cur.execute("""
                    CREATE TABLE IF NOT EXISTS checkpoint_blobs (
                        thread_id TEXT NOT NULL,
                        checkpoint_ns TEXT NOT NULL DEFAULT '',
                        channel TEXT NOT NULL,
                        version TEXT NOT NULL,
                        type TEXT NOT NULL,
                        blob BYTEA,
                        PRIMARY KEY (thread_id, checkpoint_ns, channel, version)
                    )
                """)
                
                # 创建索引以提高查询性能
                await cur.execute("""
                    CREATE INDEX IF NOT EXISTS checkpoints_thread_id_idx ON checkpoints(thread_id)
                """)
                await cur.execute("""
                    CREATE INDEX IF NOT EXISTS checkpoint_writes_thread_id_idx ON checkpoint_writes(thread_id)
                """)
                await cur.execute("""
                    CREATE INDEX IF NOT EXISTS checkpoint_blobs_thread_id_idx ON checkpoint_blobs(thread_id)
                """)
                
                # 插入迁移版本记录
                await cur.execute("""
                    INSERT INTO checkpoint_migrations (v) VALUES (1) ON CONFLICT DO NOTHING
                """)
                
                print(f"[Checkpointer] Checkpointer tables created successfully via direct SQL")


def create_checkpointer(checkpointer_type: str = None):
    """
    创建同步 checkpointer 实例
    
    Args:
        checkpointer_type: checkpointer 类型，可选 "memory", "postgres"
    
    Returns:
        checkpointer 实例
    """
    checkpointer_type = checkpointer_type or CHECKPOINTER_TYPE
    
    if checkpointer_type == "memory":
        return MemorySaver()
    
    elif checkpointer_type == "postgres":
        try:
            from langgraph.checkpoint.postgres import PostgresSaver
            from psycopg import Connection
            if not POSTGRES_CONNECTION_STRING:
                raise ValueError("POSTGRES_CONNECTION_STRING 环境变量未设置")
            conn = Connection.connect(POSTGRES_CONNECTION_STRING)
            return PostgresSaver(conn)
        except ImportError:
            print("[WARNING] psycopg 或 langgraph.checkpoint.postgres 未安装，回退到 memory")
            return MemorySaver()
    
    else:
        print(f"[WARNING] 未知的 checkpointer 类型: {checkpointer_type}，使用 memory")
        return MemorySaver()


def create_async_checkpointer(checkpointer_type: str = None):
    """
    创建异步 checkpointer 实例
    
    Args:
        checkpointer_type: checkpointer 类型
    
    Returns:
        异步 checkpointer 实例
    """
    checkpointer_type = checkpointer_type or CHECKPOINTER_TYPE
    print(f"[Checkpointer] 创建异步 checkpointer, 类型配置: {checkpointer_type}")
    
    if checkpointer_type == "memory":
        print("[Checkpointer] 使用 MemorySaver")
        return MemorySaver()
    
    elif checkpointer_type == "postgres":
        try:
            from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
            from psycopg_pool import AsyncConnectionPool
            if not POSTGRES_CONNECTION_STRING:
                raise ValueError("POSTGRES_CONNECTION_STRING 环境变量未设置")
            pool = AsyncConnectionPool(POSTGRES_CONNECTION_STRING)
            return AsyncPostgresSaver(pool)
        except ImportError:
            print("[WARNING] psycopg_pool 或 langgraph.checkpoint.postgres 未安装，回退到 memory")
            return MemorySaver()
    
    else:
        print(f"[WARNING] 未知的 checkpointer 类型: {checkpointer_type}，使用 memory")
        return MemorySaver()


async def get_or_create_async_checkpointer(checkpointer_type: str = None):
    """
    获取或创建异步 checkpointer 实例（正确处理上下文管理器）
    
    Args:
        checkpointer_type: checkpointer 类型
    
    Returns:
        已初始化的异步 checkpointer 实例
    """
    checkpointer = create_async_checkpointer(checkpointer_type)
    
    if hasattr(checkpointer, '__aenter__'):
        checkpointer = await checkpointer.__aenter__()
    
    return checkpointer


class CheckpointerManager:
    """
    Checkpointer 管理器 - 单例模式
    
    用于管理多个工作流的 checkpointer 实例
    """
    
    _instance = None
    _sync_checkpointer = None
    _async_conn_pool = None
    _async_checkpointer = None
    _initialized = False
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def get_checkpointer(self):
        """获取同步 checkpointer 实例"""
        if self._sync_checkpointer is None:
            self._sync_checkpointer = create_checkpointer()
        return self._sync_checkpointer
    
    async def get_async_checkpointer(self):
        """
        获取异步 checkpointer 实例
        """
        if self._async_checkpointer is None:
            if CHECKPOINTER_TYPE == "postgres":
                try:
                    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
                    from psycopg_pool import AsyncConnectionPool
                    if not POSTGRES_CONNECTION_STRING:
                        raise ValueError("POSTGRES_CONNECTION_STRING not set")
                    
                    self._async_conn_pool = AsyncConnectionPool(
                        POSTGRES_CONNECTION_STRING,
                        open=False
                    )
                    await self._async_conn_pool.open()
                    self._async_checkpointer = AsyncPostgresSaver(self._async_conn_pool)
                    
                    # 检查表是否存在
                    tables_exist = await verify_checkpointer_tables_exist()
                    
                    if not tables_exist:
                        # 尝试使用 setup() 方法
                        try:
                            await self._async_checkpointer.setup()
                            print(f"[Checkpointer] AsyncPostgresSaver.setup() executed")
                        except Exception as setup_e:
                            print(f"[Checkpointer] AsyncPostgresSaver.setup() failed: {setup_e}")
                        
                        # 再次验证表是否被创建
                        tables_exist_after = await verify_checkpointer_tables_exist()
                        
                        if not tables_exist_after:
                            print(f"[Checkpointer] Tables not created by setup(), using direct SQL...")
                            await create_checkpointer_tables_directly()
                    else:
                        print(f"[Checkpointer] Database tables already exist")
                    
                    print(f"[Checkpointer] AsyncPostgresSaver created, connection: {POSTGRES_CONNECTION_STRING.split('@')[1] if '@' in POSTGRES_CONNECTION_STRING else POSTGRES_CONNECTION_STRING}")
                except ImportError as e:
                    print(f"[WARNING] psycopg_pool or langgraph.checkpoint.postgres not installed: {e}, fallback to memory")
                    self._async_checkpointer = MemorySaver()
                except Exception as e:
                    print(f"[Checkpointer] ERROR: Failed to create AsyncPostgresSaver: {e}")
                    raise
            elif CHECKPOINTER_TYPE == "memory":
                self._async_checkpointer = MemorySaver()
                print("[Checkpointer] Using MemorySaver")
            else:
                self._async_checkpointer = MemorySaver()
                print(f"[Checkpointer] Unknown type {CHECKPOINTER_TYPE}, using MemorySaver")
        
        print(f"[Checkpointer] Async checkpointer type: {type(self._async_checkpointer).__name__}")
        return self._async_checkpointer
    
    async def setup(self):
        """
        初始化 checkpointer（创建数据库表等）
        """
        if not self._initialized:
            await init_checkpointer_db()
            try:
                from tools.fault_record_db import init_fault_records_table
                await init_fault_records_table()
            except Exception as e:
                print(f"[Checkpointer] 初始化 fault_records 表失败: {e}")
            self._initialized = True
    
    def reset(self):
        """重置 checkpointer（主要用于测试）"""
        self._sync_checkpointer = None
        self._async_checkpointer = None
        self._async_conn_pool = None
        self._initialized = False


checkpointer_manager = CheckpointerManager()


if __name__ == "__main__":
    print("Checkpointer 配置测试:")
    print(f"  类型: {CHECKPOINTER_TYPE}")
    print(f"  PostgreSQL 连接: {'已配置' if POSTGRES_CONNECTION_STRING else '未配置'}")