"""SQLite + FTS5 + sqlite-vec DAL 层

所有数据库操作集中在此文件。sqlite-vec 调用仅在 _load_vec_extension
和 vector_search 中出现，方便未来替换。
"""

from __future__ import annotations

import re
import sqlite3
import struct
from dataclasses import dataclass
from pathlib import Path

import sqlite_vec

# CJK 字符检测（中日韩）
_CJK_RE = re.compile(r"[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]")

# SQLite 版本检测：contentless_delete 需要 3.43.0+
_SQLITE_VERSION = tuple(int(x) for x in sqlite3.sqlite_version.split("."))
_SUPPORTS_CONTENTLESS_DELETE = _SQLITE_VERSION >= (3, 43, 0)


def tokenize_for_fts5(text: str) -> str:
    """对含 CJK 字符的文本做分词，空格连接，供 FTS5 索引和查询使用。

    - 纯英文文本原样返回（零开销）
    - jieba 未安装时降级原样返回（不影响英文用户）
    """
    if not text or not _CJK_RE.search(text):
        return text
    try:
        import jieba
        return " ".join(jieba.cut(text))
    except ImportError:
        return text


@dataclass
class ChunkRow:
    """从数据库读出的 chunk 数据"""

    id: int
    doc_id: int
    content: str
    section: str | None
    parent_section: str | None
    chunk_index: int
    tags: str | None
    token_count: int | None
    summary: str | None
    source: str | None = None  # join documents.path


class Database:
    """SQLite + FTS5 + sqlite-vec 数据访问层"""

    def __init__(self, db_path: Path | str, embedding_dim: int = 768):
        import threading
        self._db_path = str(db_path)
        self._embedding_dim = embedding_dim
        self._conn: sqlite3.Connection | None = None
        self._vec_loaded = False
        self._conn_lock = threading.RLock()

    def connect(self) -> None:
        """打开数据库连接并初始化"""
        parent = Path(self._db_path).parent
        try:
            parent.mkdir(parents=True, exist_ok=True)
        except OSError:
            pass

        self._conn = sqlite3.connect(self._db_path, timeout=10.0)
        self._conn.row_factory = sqlite3.Row
        self._conn.execute("PRAGMA journal_mode=WAL")
        self._conn.execute("PRAGMA foreign_keys=ON")
        self._conn.execute("PRAGMA busy_timeout=5000")

        self._load_vec_extension()
        self._init_schema()

    def close(self) -> None:
        if self._conn:
            self._conn.close()
            self._conn = None

    def _reconnect(self) -> None:
        """关闭旧连接并重建"""
        import logging
        if self._conn:
            try:
                self._conn.close()
            except Exception as e:
                logging.getLogger(__name__).warning("关闭旧连接失败: %s", e)
        self._conn = None
        self.connect()

    @property
    def conn(self) -> sqlite3.Connection:
        with self._conn_lock:
            if self._conn is None:
                self.connect()
            else:
                # 健康检查：连接可能已断开或损坏
                try:
                    self._conn.execute("SELECT 1")
                except sqlite3.OperationalError as e:
                    import logging
                    logging.getLogger(__name__).warning(
                        "SQLite 连接已失效（OperationalError: %s），正在重连...", e,
                    )
                    self._reconnect()
                except sqlite3.ProgrammingError as e:
                    import logging
                    logging.getLogger(__name__).warning(
                        "SQLite 连接已损坏（ProgrammingError: %s），正在重连...", e,
                    )
                    self._reconnect()
            if self._conn is None:
                raise RuntimeError("数据库连接失败，无法恢复")
            return self._conn

    # ── 文档操作 ──

    def upsert_document(
        self,
        path: str,
        hash: str,
        collection: str,
        doc_type: str,
        metadata: str | None = None,
    ) -> int:
        """插入或更新文档记录，返回 doc_id"""
        cur = self.conn.execute(
            """INSERT INTO documents (path, hash, collection, doc_type, metadata, indexed_at)
               VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
               ON CONFLICT(path) DO UPDATE SET
                 hash=excluded.hash,
                 collection=excluded.collection,
                 doc_type=excluded.doc_type,
                 metadata=excluded.metadata,
                 indexed_at=CURRENT_TIMESTAMP
               RETURNING id""",
            (path, hash, collection, doc_type, metadata),
        )
        row = cur.fetchone()
        self.conn.commit()
        return row[0]

    def get_document_hash(self, path: str) -> str | None:
        """获取文档的已索引 hash，用于变更检测"""
        row = self.conn.execute(
            "SELECT hash FROM documents WHERE path = ?", (path,)
        ).fetchone()
        return row[0] if row else None

    def delete_document(self, path: str) -> None:
        """删除文档及其所有 chunks（ON DELETE CASCADE）"""
        self.conn.execute("DELETE FROM documents WHERE path = ?", (path,))
        self.conn.commit()

    # ── Chunk 操作 ──

    def insert_chunks(self, doc_id: int, chunks: list[dict]) -> None:
        """批量插入 chunks，事务内完成

        chunks: list of dict with keys:
            content, section, parent_section, chunk_index, tags,
            embedding (list[float] | None), token_count, summary
        """
        with self.conn:
            for chunk in chunks:
                embedding_blob = None
                if chunk.get("embedding"):
                    embedding_blob = _floats_to_blob(chunk["embedding"])

                tags_str = chunk.get("tags")
                if isinstance(tags_str, list):
                    tags_str = ", ".join(tags_str)

                cur = self.conn.execute(
                    """INSERT INTO chunks
                       (doc_id, content, section, parent_section, chunk_index,
                        tags, embedding, token_count, summary)
                       VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                    (
                        doc_id,
                        chunk["content"],
                        chunk.get("section"),
                        chunk.get("parent_section"),
                        chunk.get("chunk_index", 0),
                        tags_str,
                        embedding_blob,
                        chunk.get("token_count"),
                        chunk.get("summary"),
                    ),
                )
                chunk_id = cur.lastrowid

                # FTS5 同步（对 content 做 CJK 分词以支持中文 BM25）
                self.conn.execute(
                    """INSERT INTO chunks_fts(rowid, content, section, parent_section, tags)
                       VALUES (?, ?, ?, ?, ?)""",
                    (
                        chunk_id,
                        tokenize_for_fts5(chunk["content"]),
                        chunk.get("section"),
                        chunk.get("parent_section"),
                        tags_str,
                    ),
                )

                # sqlite-vec 同步
                if embedding_blob and self._vec_loaded:
                    self.conn.execute(
                        "INSERT INTO chunks_vec(rowid, embedding) VALUES (?, ?)",
                        (chunk_id, embedding_blob),
                    )

    def delete_chunks_by_doc(self, doc_id: int) -> None:
        """删除文档的所有旧 chunks（重索引前调用）"""
        # 获取要删除的 chunk 数据（用于清理 FTS5 和 vec）
        rows = self.conn.execute(
            "SELECT id, content, section, parent_section, tags FROM chunks WHERE doc_id = ?",
            (doc_id,),
        ).fetchall()

        if not rows:
            return

        _fts5_delete_failed = False

        with self.conn:
            for row in rows:
                cid, content, section, parent_section, tags = row
                if _SUPPORTS_CONTENTLESS_DELETE:
                    # SQLite >= 3.43: contentless_delete=1 支持直接 DELETE
                    self.conn.execute(
                        "DELETE FROM chunks_fts WHERE rowid = ?", (cid,)
                    )
                else:
                    # SQLite < 3.43: contentless 表用 'delete' 命令删除
                    try:
                        self.conn.execute(
                            "INSERT INTO chunks_fts(chunks_fts, rowid, content, section, parent_section, tags) "
                            "VALUES('delete', ?, ?, ?, ?, ?)",
                            (cid, content or "", section or "", parent_section or "", tags or ""),
                        )
                    except sqlite3.OperationalError:
                        _fts5_delete_failed = True
                # sqlite-vec 删除
                if self._vec_loaded:
                    self.conn.execute(
                        "DELETE FROM chunks_vec WHERE rowid = ?", (cid,)
                    )

            self.conn.execute("DELETE FROM chunks WHERE doc_id = ?", (doc_id,))

            # 旧版 SQLite FTS5 逐条删除失败时，rebuild 整个 FTS5 索引
            if _fts5_delete_failed:
                try:
                    self.conn.execute(
                        "INSERT INTO chunks_fts(chunks_fts) VALUES('rebuild')"
                    )
                except sqlite3.OperationalError:
                    pass  # rebuild 也失败则放弃，BM25 可能有少量幽灵结果

    def get_chunk(self, chunk_id: int) -> ChunkRow | None:
        """按 ID 获取完整 chunk（含 source 路径）"""
        row = self.conn.execute(
            """SELECT c.id, c.doc_id, c.content, c.section, c.parent_section,
                      c.chunk_index, c.tags, c.token_count, c.summary, d.path as source
               FROM chunks c
               JOIN documents d ON c.doc_id = d.id
               WHERE c.id = ?""",
            (chunk_id,),
        ).fetchone()
        if not row:
            return None
        return ChunkRow(
            id=row[0], doc_id=row[1], content=row[2], section=row[3],
            parent_section=row[4], chunk_index=row[5], tags=row[6],
            token_count=row[7], summary=row[8], source=row[9],
        )

    def get_chunk_context(self, chunk_id: int) -> dict:
        """获取 chunk 的上下文（前后 chunk 的 section）"""
        row = self.conn.execute(
            "SELECT doc_id, chunk_index FROM chunks WHERE id = ?", (chunk_id,)
        ).fetchone()
        if not row:
            return {}

        doc_id, chunk_index = row[0], row[1]
        prev_row = self.conn.execute(
            "SELECT section FROM chunks WHERE doc_id = ? AND chunk_index = ?",
            (doc_id, chunk_index - 1),
        ).fetchone()
        next_row = self.conn.execute(
            "SELECT section FROM chunks WHERE doc_id = ? AND chunk_index = ?",
            (doc_id, chunk_index + 1),
        ).fetchone()

        return {
            "prev_section": prev_row[0] if prev_row else None,
            "next_section": next_row[0] if next_row else None,
        }

    # ── 搜索操作 ──

    def fts5_search(
        self, query: str, collection: str | None, limit: int
    ) -> list[tuple[int, float]]:
        """BM25 全文搜索，返回 [(chunk_id, bm25_score)]

        注意：调用方应先对查询做 tokenize_for_fts5() 分词，
        与 insert_chunks() 中的索引分词保持一致。
        """
        if collection:
            rows = self.conn.execute(
                """SELECT cf.rowid, bm25(chunks_fts) as score
                   FROM chunks_fts cf
                   JOIN chunks c ON cf.rowid = c.id
                   JOIN documents d ON c.doc_id = d.id
                   WHERE chunks_fts MATCH ? AND d.collection = ?
                   ORDER BY score
                   LIMIT ?""",
                (query, collection, limit),
            ).fetchall()
        else:
            rows = self.conn.execute(
                """SELECT rowid, bm25(chunks_fts) as score
                   FROM chunks_fts
                   WHERE chunks_fts MATCH ?
                   ORDER BY score
                   LIMIT ?""",
                (query, limit),
            ).fetchall()

        return [(r[0], abs(r[1])) for r in rows]

    def vector_search(
        self, embedding: list[float], collection: str | None, limit: int
    ) -> list[tuple[int, float]]:
        """sqlite-vec KNN 搜索，返回 [(chunk_id, distance)]"""
        if not self._vec_loaded:
            return []
        if limit <= 0:
            return []
        limit = min(limit, 1000)

        query_blob = _floats_to_blob(embedding)

        if collection:
            rows = self.conn.execute(
                """SELECT v.rowid, v.distance
                   FROM chunks_vec v
                   JOIN chunks c ON v.rowid = c.id
                   JOIN documents d ON c.doc_id = d.id
                   WHERE v.embedding MATCH ? AND k = ? AND d.collection = ?
                   ORDER BY v.distance""",
                (query_blob, limit * 2, collection),
            ).fetchall()
            return [(r[0], r[1]) for r in rows[:limit]]
        else:
            rows = self.conn.execute(
                """SELECT rowid, distance
                   FROM chunks_vec
                   WHERE embedding MATCH ? AND k = ?
                   ORDER BY distance""",
                (query_blob, limit),
            ).fetchall()
            return [(r[0], r[1]) for r in rows]

    # ── 统计 ──

    def get_stats(self) -> dict:
        """返回索引统计信息"""
        doc_count = self.conn.execute("SELECT COUNT(*) FROM documents").fetchone()[0]
        chunk_count = self.conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
        collections = [
            r[0]
            for r in self.conn.execute(
                "SELECT DISTINCT collection FROM documents"
            ).fetchall()
        ]
        last_indexed = self.conn.execute(
            "SELECT MAX(indexed_at) FROM documents"
        ).fetchone()[0]
        doc_types = [
            r[0]
            for r in self.conn.execute(
                "SELECT DISTINCT doc_type FROM documents"
            ).fetchall()
        ]

        return {
            "doc_count": doc_count,
            "chunk_count": chunk_count,
            "collections": collections,
            "doc_types": doc_types,
            "last_indexed": last_indexed,
            "db_path": self._db_path,
            "vec_loaded": self._vec_loaded,
        }

    def get_doc_chunks(self, doc_id: int) -> list[ChunkRow]:
        """获取文档的所有 chunks（走 idx_chunks_doc 索引）"""
        rows = self.conn.execute(
            """SELECT c.id, c.doc_id, c.content, c.section, c.parent_section,
                      c.chunk_index, c.tags, c.token_count, c.summary, d.path
               FROM chunks c
               JOIN documents d ON c.doc_id = d.id
               WHERE c.doc_id = ?
               ORDER BY c.chunk_index""",
            (doc_id,),
        ).fetchall()
        return [
            ChunkRow(
                id=r[0], doc_id=r[1], content=r[2], section=r[3],
                parent_section=r[4], chunk_index=r[5], tags=r[6],
                token_count=r[7], summary=r[8], source=r[9],
            )
            for r in rows
        ]

    def browse_collection(
        self, collection: str | None, doc_type: str | None = None,
        limit: int = 50, offset: int = 0,
    ) -> list[dict]:
        """浏览集合中的文档列表

        使用子查询先定位文档（走 idx_docs_collection_path 索引），
        再用相关子查询获取 chunk 统计，避免 JOIN 展开。
        """
        where_clauses = []
        params: list = []
        if collection:
            where_clauses.append("d.collection = ?")
            params.append(collection)
        if doc_type:
            where_clauses.append("d.doc_type = ?")
            params.append(doc_type)

        where_sql = ("WHERE " + " AND ".join(where_clauses)) if where_clauses else ""
        order = "d.path" if collection else "d.collection, d.path"

        rows = self.conn.execute(
            f"""SELECT d.id, d.path, d.collection, d.doc_type,
                       d.indexed_at,
                       (SELECT COUNT(*) FROM chunks WHERE doc_id = d.id) as chunk_count,
                       (SELECT COALESCE(SUM(token_count), 0) FROM chunks WHERE doc_id = d.id) as total_tokens
                FROM documents d
                {where_sql}
                ORDER BY {order}
                LIMIT ? OFFSET ?""",
            params + [limit, offset],
        ).fetchall()
        return [
            {
                "id": r[0],
                "path": r[1],
                "collection": r[2],
                "doc_type": r[3],
                "indexed_at": r[4],
                "chunk_count": r[5],
                "total_tokens": r[6],
            }
            for r in rows
        ]

    def get_collection_stats(self) -> list[dict]:
        """返回每个集合的详细统计

        使用相关子查询避免 LEFT JOIN 展开后的大结果集。
        documents 表按 collection 分组走 idx_docs_collection 索引。
        """
        rows = self.conn.execute("""
            SELECT
                d.collection,
                COUNT(*) AS doc_count,
                (SELECT COUNT(*) FROM chunks c
                 JOIN documents d2 ON c.doc_id = d2.id
                 WHERE d2.collection = d.collection) AS chunk_count,
                MAX(d.indexed_at) AS last_indexed
            FROM documents d
            GROUP BY d.collection
            ORDER BY d.collection
        """).fetchall()
        return [
            {
                "collection": r[0],
                "doc_count": r[1],
                "chunk_count": r[2],
                "last_indexed": r[3],
            }
            for r in rows
        ]

    # ── 内部方法 ──

    def _load_vec_extension(self) -> None:
        """加载 sqlite-vec 扩展"""
        try:
            self._conn.enable_load_extension(True)
            sqlite_vec.load(self._conn)
            self._conn.enable_load_extension(False)
            self._vec_loaded = True
        except Exception:
            self._vec_loaded = False

    def _init_schema(self) -> None:
        """建表 + FTS5（幂等）"""
        with self.conn:
            self.conn.executescript("""
                CREATE TABLE IF NOT EXISTS documents (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    path TEXT UNIQUE NOT NULL,
                    hash TEXT NOT NULL,
                    collection TEXT NOT NULL DEFAULT 'default',
                    doc_type TEXT NOT NULL,
                    indexed_at DATETIME DEFAULT CURRENT_TIMESTAMP,
                    metadata JSON
                );

                CREATE TABLE IF NOT EXISTS chunks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    doc_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
                    content TEXT NOT NULL,
                    section TEXT,
                    parent_section TEXT,
                    chunk_index INTEGER NOT NULL,
                    tags TEXT,
                    embedding BLOB,
                    token_count INTEGER,
                    summary TEXT
                );

                CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);
                CREATE INDEX IF NOT EXISTS idx_docs_collection ON documents(collection);
                CREATE INDEX IF NOT EXISTS idx_docs_collection_path ON documents(collection, path);
            """)

            # FTS5 虚拟表（不能用 IF NOT EXISTS，检查是否已存在）
            exists = self.conn.execute(
                "SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_fts'"
            ).fetchone()
            if not exists:
                if _SUPPORTS_CONTENTLESS_DELETE:
                    self.conn.execute("""
                        CREATE VIRTUAL TABLE chunks_fts USING fts5(
                            content, section, parent_section, tags,
                            content='',
                            contentless_delete=1
                        )
                    """)
                else:
                    self.conn.execute("""
                        CREATE VIRTUAL TABLE chunks_fts USING fts5(
                            content, section, parent_section, tags,
                            content=''
                        )
                    """)

            # sqlite-vec 虚拟表
            if self._vec_loaded:
                vec_exists = self.conn.execute(
                    "SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_vec'"
                ).fetchone()
                if not vec_exists:
                    self.conn.execute(
                        f"CREATE VIRTUAL TABLE chunks_vec USING vec0(embedding float[{self._embedding_dim}])"
                    )


def _floats_to_blob(floats: list[float]) -> bytes:
    """float 列表转 sqlite-vec 需要的 bytes"""
    return struct.pack(f"{len(floats)}f", *floats)


def _blob_to_floats(blob: bytes, dim: int) -> list[float]:
    """bytes 转 float 列表"""
    return list(struct.unpack(f"{dim}f", blob))
