"""LocalRAGService - 本地文档 RAG 服务

实现 KnowledgeBase 协议，支持文档摄取、语义检索、删除和来源追溯。
提供两种实现：
- LocalRAGService: 基于 pgvector + PostgreSQL（生产环境）
- InMemoryLocalRAGService: 基于内存（测试和开发环境）
"""

from __future__ import annotations

import json
import logging
import re
import uuid
from datetime import datetime, timezone
from typing import Any

_SAFE_TABLE_NAME_PATTERN = re.compile(r'^[a-zA-Z_][a-zA-Z0-9_]*$')

from agentkit.memory.chunking import Chunk, StructuralChunker, TextChunker
from agentkit.memory.document_loader import Document as LoaderDocument
from agentkit.memory.embedder import Embedder
from agentkit.memory.knowledge_base import (
    Document,
    KnowledgeBase,
    QueryResult,
    SourceInfo,
)
from agentkit.utils.vector_math import compute_cosine_similarity

logger = logging.getLogger(__name__)


def _loader_doc_to_kb_doc(loader_doc: LoaderDocument) -> Document:
    """将 document_loader.Document 转换为 knowledge_base.Document"""
    return Document(
        doc_id=loader_doc.doc_id,
        content=loader_doc.content,
        title=loader_doc.title,
        source_id=loader_doc.metadata.get("source", ""),
        metadata=loader_doc.metadata,
    )


class LocalRAGService:
    """基于 pgvector 的本地 RAG 服务

    实现 KnowledgeBase 协议，使用 pgvector 存储 + 检索。
    复用 EpisodicMemory 的 pgvector 基础设施模式。

    摄取 Pipeline：上传 → 解析 → 分块 → 嵌入 → 写入 pgvector
    """

    def __init__(
        self,
        session_factory: Any,
        embedder: Embedder,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
        table_name: str = "knowledge_chunks",
        pgvector_enabled: bool = True,
    ):
        """
        Args:
            session_factory: 返回 async context manager 的工厂
            embedder: 嵌入器，用于生成向量
            chunk_size: 分块大小
            chunk_overlap: 分块重叠
            table_name: pgvector 查询使用的表名
            pgvector_enabled: 是否使用 pgvector 原生检索
        """
        self._session_factory = session_factory
        self._embedder = embedder
        self._chunk_size = chunk_size
        self._chunk_overlap = chunk_overlap
        self._table_name = table_name
        if not _SAFE_TABLE_NAME_PATTERN.match(self._table_name):
            raise ValueError(f"Invalid table_name: {self._table_name}. Must match [a-zA-Z_][a-zA-Z0-9_]*")
        self._pgvector_enabled = pgvector_enabled
        self._text_chunker = TextChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        self._structural_chunker = StructuralChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    async def ingest(self, documents: list[Document]) -> list[str]:
        """摄取文档列表

        Args:
            documents: knowledge_base.Document 对象列表

        Returns:
            成功摄取的文档 ID 列表
        """
        ingested_ids = []

        for doc in documents:
            try:
                chunks = self._chunk_document(doc)
                await self._store_chunks(doc, chunks)
                ingested_ids.append(doc.doc_id)
                logger.info(f"Ingested document '{doc.title}' with {len(chunks)} chunks")
            except Exception as e:
                logger.error(f"Failed to ingest document '{doc.title}': {e}")

        return ingested_ids

    async def query(self, text: str, top_k: int = 5) -> list[QueryResult]:
        """语义检索

        Args:
            text: 查询文本
            top_k: 返回结果数量

        Returns:
            检索结果列表
        """
        query_embedding = await self._embedder.embed(text)

        async with self._session_factory() as db:
            try:
                if self._pgvector_enabled:
                    return await self._query_pgvector(db, query_embedding, top_k)
                return await self._query_client_side(db, query_embedding, top_k)
            except Exception as e:
                logger.error(f"Failed to query knowledge base: {e}")
                return []

    async def delete_by_id(self, id: str) -> bool:
        """按文档 ID 删除

        Args:
            id: 文档 ID

        Returns:
            是否删除成功
        """
        async with self._session_factory() as db:
            try:
                from sqlalchemy import text as sql_text

                sql = sql_text(
                    f"DELETE FROM {self._table_name} WHERE source_doc_id = :doc_id"
                )
                await db.execute(sql, {"doc_id": id})
                await db.commit()
                return True
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to delete document {id}: {e}")
                return False

    async def list_sources(self) -> list[SourceInfo]:
        """列出已摄取的文档

        Returns:
            文档元信息列表
        """
        async with self._session_factory() as db:
            try:
                from sqlalchemy import text as sql_text

                sql = sql_text(
                    f"SELECT source_doc_id, source_title, doc_format, "
                    f"COUNT(*) as chunk_count, "
                    f"MIN(created_at) as created_at, "
                    f"MIN(doc_metadata) as doc_metadata "
                    f"FROM {self._table_name} "
                    f"GROUP BY source_doc_id, source_title, doc_format "
                    f"ORDER BY MIN(created_at) DESC"
                )
                result = await db.execute(sql)
                rows = result.mappings().all()

                sources = []
                for row in rows:
                    meta = {}
                    if row.get("doc_metadata"):
                        try:
                            meta = json.loads(row["doc_metadata"])
                        except (json.JSONDecodeError, TypeError):
                            pass

                    sources.append(SourceInfo(
                        source_id=row["source_doc_id"],
                        source_name=row.get("source_title", ""),
                        source_type=row.get("doc_format", "local"),
                        document_count=row.get("chunk_count", 0),
                        last_updated=row["created_at"] if row.get("created_at") else None,
                    ))
                return sources
            except Exception as e:
                logger.error(f"Failed to list sources: {e}")
                return []

    async def health_check(self) -> bool:
        """检查服务健康状态"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import text as sql_text

                await db.execute(sql_text(f"SELECT 1 FROM {self._table_name} LIMIT 1"))
                return True
            except Exception as e:
                logger.error(f"Health check failed: {e}")
                return False

    def _chunk_document(self, doc: Document) -> list[Chunk]:
        """将文档分块"""
        doc_format = doc.metadata.get("format", "text")

        # Markdown 和 HTML 使用结构化分块
        if doc_format in ("markdown", "html"):
            chunks = self._structural_chunker.chunk(
                doc.content,
                source_doc_id=doc.doc_id,
                metadata=doc.metadata,
            )
        else:
            chunks = self._text_chunker.chunk(
                doc.content,
                source_doc_id=doc.doc_id,
                metadata=doc.metadata,
            )

        return chunks

    async def _store_chunks(self, doc: Document, chunks: list[Chunk]) -> None:
        """存储文档块到 pgvector"""
        async with self._session_factory() as db:
            try:
                from sqlalchemy import text as sql_text

                # Batch embedding generation
                embeddings: list[list[float]] = []
                if hasattr(self._embedder, "embed_batch"):
                    embeddings = await self._embedder.embed_batch([c.content for c in chunks])
                else:
                    for chunk in chunks:
                        embedding = await self._embedder.embed(chunk.content)
                        embeddings.append(embedding)

                # Batch INSERT using executemany
                sql = sql_text(
                    f"INSERT INTO {self._table_name} "
                    f"(chunk_id, source_doc_id, source_title, doc_format, "
                    f"content, embedding, chunk_metadata, doc_metadata, created_at) "
                    f"VALUES (:chunk_id, :doc_id, :title, :format, "
                    f":content, :embedding, :chunk_meta, :doc_meta, :created_at)"
                )

                now = datetime.now(timezone.utc)
                params_list = [
                    {
                        "chunk_id": chunk.chunk_id,
                        "doc_id": doc.doc_id,
                        "title": doc.title,
                        "format": doc.metadata.get("format", "unknown"),
                        "content": chunk.content,
                        "embedding": str(embedding),
                        "chunk_meta": json.dumps(chunk.metadata, ensure_ascii=False),
                        "doc_meta": json.dumps(doc.metadata, ensure_ascii=False),
                        "created_at": now,
                    }
                    for chunk, embedding in zip(chunks, embeddings)
                ]

                await db.execute(sql, params_list)
                await db.commit()
            except Exception as e:
                await db.rollback()
                logger.error(f"Failed to store chunks for document '{doc.title}': {e}")
                raise

    async def _query_pgvector(
        self,
        db: Any,
        query_embedding: list[float],
        top_k: int,
    ) -> list[QueryResult]:
        """使用 pgvector <=> 算符检索"""
        from sqlalchemy import text as sql_text

        sql = sql_text(
            f"SELECT chunk_id, source_doc_id, source_title, content, "
            f"chunk_metadata, embedding <=> :query_vec AS distance "
            f"FROM {self._table_name} "
            f"ORDER BY embedding <=> :query_vec "
            f"LIMIT :lim"
        )

        result = await db.execute(sql, {
            "query_vec": str(query_embedding),
            "lim": top_k,
        })
        rows = result.mappings().all()

        results = []
        for row in rows:
            # 从 distance 计算 cosine similarity
            distance = row.get("distance", 0.0)
            # pgvector <=> 返回 cosine distance = 1 - cosine_similarity
            cosine = max(0.0, 1.0 - float(distance))

            chunk_meta = {}
            if row.get("chunk_metadata"):
                try:
                    chunk_meta = json.loads(row["chunk_metadata"])
                except (json.JSONDecodeError, TypeError):
                    pass

            results.append(QueryResult(
                content=row["content"],
                source_id=row["source_doc_id"],
                source_name=row.get("source_title", ""),
                score=cosine,
                metadata=chunk_meta,
                doc_id=row["source_doc_id"],
                title=row.get("source_title", ""),
            ))

        return results

    async def _query_client_side(
        self,
        db: Any,
        query_embedding: list[float],
        top_k: int,
    ) -> list[QueryResult]:
        """客户端 O(N) cosine similarity 检索（回退路径）"""
        from sqlalchemy import text as sql_text

        sql = sql_text(
            f"SELECT chunk_id, source_doc_id, source_title, content, "
            f"chunk_metadata, embedding "
            f"FROM {self._table_name} "
            f"LIMIT 500"
        )

        result = await db.execute(sql)
        rows = result.mappings().all()

        candidates = []
        for row in rows:
            row_embedding = row.get("embedding")
            if row_embedding is None:
                continue

            # 解析存储的 embedding
            try:
                if isinstance(row_embedding, str):
                    stored_embedding = json.loads(row_embedding)
                else:
                    stored_embedding = list(row_embedding)
            except (json.JSONDecodeError, TypeError):
                continue

            cosine = compute_cosine_similarity(query_embedding, stored_embedding)
            if cosine < 0.1:
                continue

            chunk_meta = {}
            if row.get("chunk_metadata"):
                try:
                    chunk_meta = json.loads(row["chunk_metadata"])
                except (json.JSONDecodeError, TypeError):
                    pass

            candidates.append(QueryResult(
                content=row["content"],
                source_id=row["source_doc_id"],
                source_name=row.get("source_title", ""),
                score=cosine,
                metadata=chunk_meta,
                doc_id=row["source_doc_id"],
                title=row.get("source_title", ""),
            ))

        candidates.sort(key=lambda x: x.score, reverse=True)
        return candidates[:top_k]


class InMemoryLocalRAGService:
    """基于内存的本地 RAG 服务

    用于测试和开发环境，无需 pgvector 依赖。
    实现 KnowledgeBase 协议。
    """

    def __init__(
        self,
        embedder: Embedder,
        chunk_size: int = 1000,
        chunk_overlap: int = 200,
    ):
        """
        Args:
            embedder: 嵌入器
            chunk_size: 分块大小
            chunk_overlap: 分块重叠
        """
        self._embedder = embedder
        self._text_chunker = TextChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
        self._structural_chunker = StructuralChunker(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

        # 内存存储
        self._chunks: dict[str, dict[str, Any]] = {}  # chunk_id → {content, embedding, metadata}
        self._documents: dict[str, dict[str, Any]] = {}  # doc_id → {title, format, chunk_ids, metadata, created_at}

    async def ingest(self, documents: list[Document]) -> list[str]:
        """摄取文档列表

        也支持传入 document_loader.Document，会自动转换为 knowledge_base.Document。
        """
        ingested_ids = []

        for doc in documents:
            # 支持 document_loader.Document 自动转换
            if isinstance(doc, LoaderDocument):
                doc = _loader_doc_to_kb_doc(doc)

            try:
                chunks = self._chunk_document(doc)
                chunk_ids = []

                for chunk in chunks:
                    embedding = await self._embedder.embed(chunk.content)
                    self._chunks[chunk.chunk_id] = {
                        "content": chunk.content,
                        "embedding": embedding,
                        "metadata": chunk.metadata,
                        "source_doc_id": doc.doc_id,
                    }
                    chunk_ids.append(chunk.chunk_id)

                self._documents[doc.doc_id] = {
                    "title": doc.title,
                    "source_id": doc.source_id,
                    "format": doc.metadata.get("format", "unknown"),
                    "chunk_ids": chunk_ids,
                    "metadata": doc.metadata,
                    "created_at": datetime.now(timezone.utc),
                }
                ingested_ids.append(doc.doc_id)
                logger.info(f"Ingested document '{doc.title}' with {len(chunks)} chunks")
            except Exception as e:
                logger.error(f"Failed to ingest document '{doc.title}': {e}")

        return ingested_ids

    async def query(self, text: str, top_k: int = 5) -> list[QueryResult]:
        """语义检索"""
        query_embedding = await self._embedder.embed(text)

        candidates = []
        for chunk_id, chunk_data in self._chunks.items():
            stored_embedding = chunk_data["embedding"]
            cosine = compute_cosine_similarity(query_embedding, stored_embedding)
            if cosine < 0.1:
                continue

            source_doc_id = chunk_data["source_doc_id"]
            doc_info = self._documents.get(source_doc_id, {})

            candidates.append(QueryResult(
                content=chunk_data["content"],
                source_id=source_doc_id,
                source_name=doc_info.get("title", ""),
                score=cosine,
                metadata=chunk_data.get("metadata", {}),
                doc_id=source_doc_id,
                title=doc_info.get("title", ""),
            ))

        candidates.sort(key=lambda x: x.score, reverse=True)
        return candidates[:top_k]

    async def delete_by_id(self, id: str) -> bool:
        """按文档 ID 删除"""
        if id not in self._documents:
            return False

        doc_info = self._documents[id]
        for chunk_id in doc_info.get("chunk_ids", []):
            self._chunks.pop(chunk_id, None)

        del self._documents[id]
        return True

    async def list_sources(self) -> list[SourceInfo]:
        """列出已摄取的文档"""
        sources = []
        for doc_id, doc_info in self._documents.items():
            sources.append(SourceInfo(
                source_id=doc_id,
                source_name=doc_info["title"],
                source_type=doc_info.get("format", "local"),
                document_count=len(doc_info.get("chunk_ids", [])),
                last_updated=doc_info.get("created_at"),
            ))
        return sources

    async def health_check(self) -> bool:
        """检查服务健康状态"""
        return True

    def _chunk_document(self, doc: Document) -> list[Chunk]:
        """将文档分块"""
        doc_format = doc.metadata.get("format", "text")

        if doc_format in ("markdown", "html"):
            return self._structural_chunker.chunk(
                doc.content,
                source_doc_id=doc.doc_id,
                metadata=doc.metadata,
            )
        return self._text_chunker.chunk(
            doc.content,
            source_doc_id=doc.doc_id,
            metadata=doc.metadata,
        )
