"""
知识库服务 — 写入向量 + 语义检索
使用 pgvector 原生 SQL 向量检索（<=> 余弦距离算子），不在 Python 侧计算
"""

from typing import Optional

from sqlalchemy import text
from sqlmodel import Session, select
from pgvector.sqlalchemy import Vector

from app.models.knowledge import KnowledgeItem, KnowledgeEmbedding
from app.services.embedding_service import EmbeddingService
from app.db.session import engine


class KnowledgeService:
    """知识库 CRUD + 语义检索"""

    def __init__(self):
        self.embedder = EmbeddingService()

    def store_md_file(
        self,
        title: str,
        content_md: str,
        source_file_name: Optional[str] = None,
        source_type: str = "manual",
        author: Optional[str] = None,
    ) -> KnowledgeItem:
        """
        读取一篇 md 内容，调用 embo-01 拿到向量，写入 knowledge_items + knowledge_embeddings
        """
        # 调用 embedding（type="db" 表示存入知识库）
        embedding_list = self.embedder.embed_single(content_md, embed_type="db")

        with Session(engine) as session:
            # 写入 knowledge_items
            item = KnowledgeItem(
                title=title,
                content_md=content_md,
                source_type=source_type,
                source_file_name=source_file_name,
                author=author,
            )
            session.add(item)
            session.flush()  # 拿到 id

            # 写入 knowledge_embeddings（单 chunk，chunk_index=0）
            # 直接传 list，pgvector.sqlalchemy.Vector 会自动处理转换
            emb = KnowledgeEmbedding(
                knowledge_id=item.id,
                chunk_index=0,
                chunk_text=content_md,
                embedding=embedding_list,
            )
            session.add(emb)
            session.commit()
            session.refresh(item)
            return item

    def search_similar(self, query_text: str, top_k: int = 5) -> list[dict]:
        """
        语义检索：查询句转为向量，用 SQL 余弦距离（<=>）在数据库层检索
        返回 top_k 条相似笔记，含相似度分数
        """
        # 查询向量（type="query"）
        query_vector = self.embedder.embed_single(query_text, embed_type="query")

        # 将向量列表转为 pgvector SQL 字符串格式
        vec_str = "[" + ",".join(str(v) for v in query_vector) + "]"

        with Session(engine) as session:
            # pgvector 原生 SQL：<=> 是余弦距离，1 - 距离 = 相似度
            # 用字符串插注向量，避免 psycopg2 参数化问题
            sql = f"""
                SELECT
                    ki.id,
                    ki.title,
                    ki.source_type,
                    1 - (ke.embedding <=> '{vec_str}'::vector) AS similarity
                FROM knowledge_embeddings ke
                JOIN knowledge_items ki ON ke.knowledge_id = ki.id
                WHERE ke.chunk_index = 0
                ORDER BY ke.embedding <=> '{vec_str}'::vector
                LIMIT {top_k}
            """
            stmt = text(sql)
            rows = session.execute(stmt).all()

            results = []
            for row in rows:
                results.append({
                    "id": row.id,
                    "title": row.title,
                    "source_type": row.source_type,
                    "similarity": round(row.similarity, 4),
                })
            return results

    def get_item_count(self) -> int:
        """返回 knowledge_items 表行数"""
        with Session(engine) as session:
            count = session.exec(select(KnowledgeItem)).all()
            return len(count)

    def get_embedding_count(self) -> int:
        """返回 knowledge_embeddings 表行数"""
        with Session(engine) as session:
            count = session.exec(select(KnowledgeEmbedding)).all()
            return len(count)