"""pgvector 索引管理 — LlamaIndex PGVectorStore 封装。

Schema 隔离：使用显式表名 `rag_platform_kb_chunks`，不可使用默认 `data_<name>`。
确认 `create_if_not_exists` 不会触碰 `episodic_memory` 或任何 `memory/` 所属表。
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from llama_index.vector_stores.postgres import PGVectorStore

# 显式表名 — 与 memory/episodic.py 的 episodic_memories 表完全隔离
KB_CHUNKS_TABLE = "rag_platform_kb_chunks"

# 默认 embedding 维度（OpenAI text-embedding-3-small）
DEFAULT_EMBED_DIM = 1536


def _async_to_sync_url(database_url: str) -> str:
    """将 asyncpg URL 转换为 psycopg2 URL（LlamaIndex PGVectorStore 使用同步驱动）。"""
    return database_url.replace("postgresql+asyncpg://", "postgresql://")


def create_vector_store(
    database_url: str,
    embed_dim: int = DEFAULT_EMBED_DIM,
    hybrid_search: bool = True,
) -> "PGVectorStore":
    """创建 PGVectorStore，使用显式表名实现 schema 隔离。

    Args:
        database_url: SQLAlchemy 数据库 URL（async 或 sync 均可）
        embed_dim: embedding 向量维度
        hybrid_search: 是否启用混合搜索（向量 + 全文）

    Returns:
        LlamaIndex PGVectorStore 实例

    Raises:
        ImportError: 如果 llama-index-vector-stores-postgres 未安装
    """
    from llama_index.vector_stores.postgres import PGVectorStore

    sync_url = _async_to_sync_url(database_url)

    logger.info(
        "Creating PGVectorStore: table=%s, embed_dim=%d, hybrid=%s",
        KB_CHUNKS_TABLE,
        embed_dim,
        hybrid_search,
    )

    return PGVectorStore.from_params(
        database=sync_url,
        table_name=KB_CHUNKS_TABLE,
        embed_dim=embed_dim,
        hybrid_search=hybrid_search,
        text_search_config="english",  # U4 将用 jieba 替换为中文分词
        # ponytail: 不设 schema_name，默认 public — 避免创建独立 schema 的运维复杂度
        # 表名前缀 rag_platform_ 已足够隔离
    )


async def ensure_vector_store_schema(database_url: str, embed_dim: int = DEFAULT_EMBED_DIM) -> None:
    """确保 vector store 表存在（幂等）。

    在应用启动时调用，创建表结构（如果不存在）。
    不会触碰 episodic_memory 或任何 memory/ 所属表。
    """
    vs = create_vector_store(database_url, embed_dim=embed_dim)
    # PGVectorStore.__init__ 内部会调用 create_if_not_exists
    # 显式调用 _initialize 来确保表创建
    vs._initialize()
    logger.info("Vector store schema ensured: table=%s", KB_CHUNKS_TABLE)