"""
异步 PostgreSQL 存储实现 - 用于生产环境。

使用 asyncpg 实现真正的异步 PostgreSQL 文档存储，支持高并发访问。
"""

import asyncio
import json
import logging
from typing import List, Dict, Any, Optional, Iterator, Tuple, Sequence

from langchain_core.documents import Document
from langchain_core.stores import BaseStore

import asyncpg

logger = logging.getLogger(__name__)

class PostgresDocStore(BaseStore[str, Any]):
    """
    异步 PostgreSQL 文档存储实现。
    
    使用 asyncpg 作为异步 PostgreSQL 客户端，支持：
    - 真正的异步操作
    - 连接池管理
    - 自动表创建
    - 批量操作（amget/amset/amdelete）
    - JSONB 数据存储
    - 并发控制
    
    适用于生产环境，提供高性能的异步数据持久化。
    
    Attributes:
        dsn: PostgreSQL 连接字符串
        table_name: 存储表名，默认为 "parent_documents"
        _pool: asyncpg 连接池实例
        _semaphore: 控制并发数的信号量（可选）
    """
    
    def __init__(
        self,
        connection_string: str,
        table_name: str = "parent_documents",
        pool_config: Optional[Dict[str, Any]] = None,
        max_concurrency: Optional[int] = None
    ):
        """
        初始化异步 PostgreSQL 文档存储。
        
        Args:
            connection_string: PostgreSQL 连接 URL，格式：
                "postgresql://user:password@host:port/database?sslmode=disable"
            table_name: 存储表名，默认为 "parent_documents"
            pool_config: 连接池配置字典，包含：
                - min_size: 最小连接数（默认 2）
                - max_size: 最大连接数（默认 10）
            max_concurrency: 最大并发操作数，如果为 None 则不限制
        
        Raises:
            ImportError: 未安装 asyncpg 时抛出
        
        Example:
            >>> store = PostgresDocStore(
            ...     "postgresql://user:pass@localhost:5432/mydb",
            ...     table_name="parent_docs",
            ...     pool_config={"min_size": 5, "max_size": 20},
            ...     max_concurrency=10
            ... )
        """

        
        self.dsn = connection_string
        self.table_name = table_name
        self._pool: Optional["asyncpg.Pool"] = None
        self._pool_config = pool_config or {}
        
        # 并发控制信号量
        self._semaphore = None
        if max_concurrency is not None and max_concurrency > 0:
            self._semaphore = asyncio.Semaphore(max_concurrency)
        
        # 注意：连接池的异步初始化延迟到第一次使用时
        # 表结构创建也延迟到第一次操作时
    
    async def _get_pool(self):
        """获取或创建 asyncpg 连接池。"""
        if self._pool is None:
            import asyncpg
            min_size = self._pool_config.get("min_size", 2)
            max_size = self._pool_config.get("max_size", 10)
            
            try:
                self._pool = await asyncpg.create_pool(
                    dsn=self.dsn,
                    min_size=min_size,
                    max_size=max_size
                )
                logger.info(f"PostgreSQL 异步连接池已创建: {self.table_name}")
                
                # 初始化表结构
                await self._create_table()
            except Exception as e:
                raise RuntimeError(f"PostgreSQL 异步连接池创建失败: {e}")
        
        return self._pool
    
    async def _create_table(self):
        """创建存储表（如果不存在）。"""
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            async with conn.transaction():
                await conn.execute(f"""
                    CREATE TABLE IF NOT EXISTS {self.table_name} (
                        key TEXT PRIMARY KEY,
                        value JSONB NOT NULL,
                        created_at TIMESTAMPTZ DEFAULT NOW()
                    )
                """)
                logger.info(f"表 {self.table_name} 已就绪")
    
    async def _with_concurrency_control(self, coro):
        """使用信号量控制并发执行。"""
        if self._semaphore is None:
            return await coro
        async with self._semaphore:
            return await coro
    
    # --- 同步方法（保持兼容性，但功能有限）---
    
    def mget(self, keys: Sequence[str]) -> List[Optional[Any]]:
        """不支持同步操作，请使用异步 amget 方法。"""
        raise NotImplementedError("不支持同步操作，请使用异步 amget 方法。")
    
    def mset(self, key_value_pairs: Sequence[Tuple[str, Any]]) -> None:
        """不支持同步操作，请使用异步 amset 方法。"""
        raise NotImplementedError("不支持同步操作，请使用异步 amset 方法。")
    
    def mdelete(self, keys: Sequence[str]) -> None:
        """不支持同步操作，请使用异步 amdelete 方法。"""
        raise NotImplementedError("不支持同步操作，请使用异步 amdelete 方法。")
    
    def yield_keys(self, *, prefix: str | None = None) -> Iterator[str]:
        """不支持同步操作，请使用异步 ayield_keys 方法。"""
        raise NotImplementedError("不支持同步操作，请使用异步 ayield_keys 方法。")
    
    # --- 异步方法（真正的实现）---
    
    async def amget(self, keys: Sequence[str]) -> List[Optional[Any]]:
        """异步批量获取文档。"""
        if not keys:
            return []
        
        async def _amget():
            pool = await self._get_pool()
            async with pool.acquire() as conn:
                rows = await conn.fetch(
                    f"SELECT key, value FROM {self.table_name} WHERE key = ANY($1)",
                    keys
                )
                result_map = {}
                for row in rows:
                    val = row['value']
                    if isinstance(val, str):
                        val = json.loads(val)
                    if isinstance(val, dict) and 'page_content' in val:
                        result_map[row['key']] = Document(**val)
                    else:
                        result_map[row['key']] = val
                return [result_map.get(key) for key in keys]
        
        return await self._with_concurrency_control(_amget())
    
    async def amset(self, key_value_pairs: Sequence[Tuple[str, Any]]) -> None:
        """异步批量设置文档。"""
        if not key_value_pairs:
            return
        
        async def _amset():
            pool = await self._get_pool()
            async with pool.acquire() as conn:
                async with conn.transaction():
                    await conn.executemany(
                        f"""
                        INSERT INTO {self.table_name} (key, value)
                        VALUES ($1, $2)
                        ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value
                        """,
                        [
                            (k, json.dumps(v.dict() if isinstance(v, Document) else v, ensure_ascii=False))
                            for k, v in key_value_pairs
                        ]
                    )
                    logger.debug(f"已异步批量设置 {len(key_value_pairs)} 个文档")
        
        await self._with_concurrency_control(_amset())
    
    async def amdelete(self, keys: Sequence[str]) -> None:
        """异步批量删除文档。"""
        if not keys:
            return
        
        async def _amdelete():
            pool = await self._get_pool()
            async with pool.acquire() as conn:
                async with conn.transaction():
                    await conn.execute(
                        f"DELETE FROM {self.table_name} WHERE key = ANY($1)",
                        keys
                    )
                    logger.debug(f"已异步批量删除 {len(keys)} 个文档")
        
        await self._with_concurrency_control(_amdelete())
    
    async def ayield_keys(self, *, prefix: str | None = None) -> Iterator[str]:
        """异步迭代所有键。
        
        注意：这是一个异步生成器，需要使用 async for 迭代。
        """
        pool = await self._get_pool()
        async with pool.acquire() as conn:
            if prefix:
                rows = await conn.fetch(
                    f"SELECT key FROM {self.table_name} WHERE key LIKE $1 ORDER BY key",
                    f"{prefix}%"
                )
            else:
                rows = await conn.fetch(
                    f"SELECT key FROM {self.table_name} ORDER BY key"
                )
            
            for row in rows:
                yield row['key']
    
    async def aclose(self) -> None:
        """异步关闭连接池，释放资源。"""
        if self._pool:
            await self._pool.close()
            self._pool = None
            logger.info("PostgreSQL 异步连接池已关闭")
    
    def close(self) -> None:
        """同步关闭连接池（功能有限）。
        
        注意：在异步环境中，请使用 aclose 方法。
        """
        pass