ax-backend-v1/app/services/content_db.py

"""content_db.py — SQLite 内容数据库操作类"""

import json
import sqlite3
from pathlib import Path
from typing import Optional


class ContentDB:
    """SQLite 内容数据库操作类
    
    每个文档对应一个独立的 SQLite 数据库文件，存储文档的所有 Blocks。
    """
    
    def __init__(self, db_path: str):
        self.db_path = Path(db_path)
        self.conn: Optional[sqlite3.Connection] = None
    
    def connect(self):
        """连接数据库"""
        self.conn = sqlite3.connect(str(self.db_path))
        self.conn.row_factory = sqlite3.Row
        return self
    
    def close(self):
        """关闭连接"""
        if self.conn:
            self.conn.close()
            self.conn = None
    
    def __enter__(self):
        """上下文管理器入口"""
        return self.connect()
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出"""
        self.close()
    
    def create_tables(self):
        """创建 document_blocks 表"""
        self.conn.execute("""
            CREATE TABLE IF NOT EXISTS document_blocks (
                id TEXT PRIMARY KEY,
                block_order INTEGER NOT NULL,
                type TEXT NOT NULL,
                level INTEGER DEFAULT 0,
                "index" INTEGER DEFAULT 0,
                content TEXT NOT NULL,
                word_style TEXT DEFAULT '',
                style TEXT DEFAULT '{}',
                metadata TEXT DEFAULT '{}'
            )
        """)
        
        # 创建索引
        self.conn.execute(
            'CREATE INDEX IF NOT EXISTS idx_block_order ON document_blocks(block_order)'
        )
        self.conn.execute(
            'CREATE INDEX IF NOT EXISTS idx_type ON document_blocks(type)'
        )
        self.conn.execute(
            'CREATE INDEX IF NOT EXISTS idx_level ON document_blocks(level)'
        )
        self.conn.commit()
    
    def insert_blocks(self, blocks: list[dict]):
        """批量插入 blocks"""
        for block in blocks:
            content = block['content']
            if isinstance(content, (dict, list)):
                content = json.dumps(content, ensure_ascii=False)
            
            self.conn.execute("""
                INSERT INTO document_blocks 
                (id, block_order, type, level, "index", content, word_style, style, metadata)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (
                block['id'],
                block['block_order'],
                block['type'],
                block.get('level', 0),
                block.get('index', 0),
                content,
                block.get('word_style', ''),
                json.dumps(block.get('style', {}), ensure_ascii=False),
                json.dumps(block.get('metadata', {}), ensure_ascii=False)
            ))
        self.conn.commit()
    
    def get_blocks(self, order_by: str = 'block_order') -> list[dict]:
        """查询所有 blocks"""
        cursor = self.conn.execute(f"""
            SELECT * FROM document_blocks 
            ORDER BY {order_by}
        """)
        rows = cursor.fetchall()
        return [self._row_to_dict(row) for row in rows]
    
    def get_block_by_id(self, block_id: str) -> Optional[dict]:
        """按 ID 查询单个 block"""
        cursor = self.conn.execute("""
            SELECT * FROM document_blocks WHERE id = ?
        """, (block_id,))
        row = cursor.fetchone()
        return self._row_to_dict(row) if row else None
    
    def update_block(self, block_id: str, updates: dict):
        """更新单个 block"""
        set_clauses = []
        params = []
        
        if 'content' in updates:
            content = updates['content']
            if isinstance(content, (dict, list)):
                content = json.dumps(content, ensure_ascii=False)
            set_clauses.append('content = ?')
            params.append(content)
        
        if 'style' in updates:
            set_clauses.append('style = ?')
            params.append(json.dumps(updates['style'], ensure_ascii=False))
        
        if 'word_style' in updates:
            set_clauses.append('word_style = ?')
            params.append(updates['word_style'])
        
        if 'metadata' in updates:
            set_clauses.append('metadata = ?')
            params.append(json.dumps(updates['metadata'], ensure_ascii=False))
        
        if not set_clauses:
            return
        
        params.append(block_id)
        sql = f"UPDATE document_blocks SET {', '.join(set_clauses)} WHERE id = ?"
        self.conn.execute(sql, params)
        self.conn.commit()
    
    def delete_block(self, block_id: str):
        """删除单个 block"""
        self.conn.execute('DELETE FROM document_blocks WHERE id = ?', (block_id,))
        self.conn.commit()
    
    def search_blocks(self, query: str, block_type: Optional[str] = None) -> list[dict]:
        """搜索 blocks"""
        sql = "SELECT * FROM document_blocks WHERE content LIKE ?"
        params = [f'%{query}%']
        
        if block_type:
            sql += " AND type = ?"
            params.append(block_type)
        
        sql += " ORDER BY block_order"
        
        cursor = self.conn.execute(sql, params)
        rows = cursor.fetchall()
        return [self._row_to_dict(row) for row in rows]
    
    def get_headings(self) -> list[dict]:
        """获取所有标题块"""
        cursor = self.conn.execute("""
            SELECT * FROM document_blocks 
            WHERE type = 'heading' 
            ORDER BY block_order
        """)
        rows = cursor.fetchall()
        return [self._row_to_dict(row) for row in rows]
    
    def get_stats(self) -> dict:
        """获取统计信息"""
        cursor = self.conn.execute("""
            SELECT 
                type,
                COUNT(*) as count
            FROM document_blocks
            GROUP BY type
        """)
        stats = {row['type']: row['count'] for row in cursor.fetchall()}
        
        cursor = self.conn.execute("SELECT COUNT(*) as total FROM document_blocks")
        total = cursor.fetchone()['total']
        
        return {
            'total': total,
            'by_type': stats
        }
    
    def calculate_index_for_insert(
        self,
        after_block_id: str,
        new_type: str,
        new_level: int = 0
    ) -> int:
        """计算插入 block 的 index（稀疏排序）
        
        Args:
            after_block_id: 在哪个 block 后插入（如果为 None 则插入到末尾）
            new_type: 新 block 的类型
            new_level: 新 block 的级别（标题有效）
            
        Returns:
            计算得到的 index
        """
        if after_block_id:
            # 1. 获取 after_block 的 block_order
            after_block = self.get_block_by_id(after_block_id)
            if not after_block:
                raise ValueError(f"Block not found: {after_block_id}")
            after_order = after_block['block_order']
            
            # 2. 查找下一个同类型同级别的 block
            if new_type == 'heading':
                # 标题：按 level 查询
                cursor = self.conn.execute("""
                    SELECT "index" FROM document_blocks
                    WHERE type = ? AND level = ? AND block_order > ?
                    ORDER BY block_order LIMIT 1
                """, (new_type, new_level, after_order))
            else:
                # 其他类型：只按 type 查询
                cursor = self.conn.execute("""
                    SELECT "index" FROM document_blocks
                    WHERE type = ? AND block_order > ?
                    ORDER BY block_order LIMIT 1
                """, (new_type, after_order))
            
            next_row = cursor.fetchone()
            
            if next_row:
                next_index = next_row['index']
                
                # 3. 查找前一个同类型同级别的 block
                if new_type == 'heading':
                    cursor = self.conn.execute("""
                        SELECT "index" FROM document_blocks
                        WHERE type = ? AND level = ? AND block_order <= ?
                        ORDER BY block_order DESC LIMIT 1
                    """, (new_type, new_level, after_order))
                else:
                    cursor = self.conn.execute("""
                        SELECT "index" FROM document_blocks
                        WHERE type = ? AND block_order <= ?
                        ORDER BY block_order DESC LIMIT 1
                    """, (new_type, after_order))
                
                prev_row = cursor.fetchone()
                prev_index = prev_row['index'] if prev_row else -100
                
                # 4. 计算中间值
                gap = next_index - prev_index
                if gap <= 1:
                    # 间隙不足，触发局部重排
                    self._rebalance_indexes_between(
                        new_type, new_level, prev_index, next_index
                    )
                    # 重新查询
                    if new_type == 'heading':
                        cursor = self.conn.execute("""
                            SELECT "index" FROM document_blocks
                            WHERE type = ? AND level = ? AND block_order > ?
                            ORDER BY block_order LIMIT 1
                        """, (new_type, new_level, after_order))
                    else:
                        cursor = self.conn.execute("""
                            SELECT "index" FROM document_blocks
                            WHERE type = ? AND block_order > ?
                            ORDER BY block_order LIMIT 1
                        """, (new_type, after_order))
                    next_row = cursor.fetchone()
                    next_index = next_row['index'] if next_row else prev_index + 200
                
                new_index = (prev_index + next_index) // 2
            else:
                # 没有下一个同类 block，追加到最后
                if new_type == 'heading':
                    cursor = self.conn.execute("""
                        SELECT MAX("index") as max_index FROM document_blocks
                        WHERE type = ? AND level = ?
                    """, (new_type, new_level))
                else:
                    cursor = self.conn.execute("""
                        SELECT MAX("index") as max_index FROM document_blocks
                        WHERE type = ?
                    """, (new_type,))
                
                row = cursor.fetchone()
                max_index = row['max_index'] if row['max_index'] is not None else -100
                new_index = max_index + 100
        else:
            # 插入到文档末尾
            if new_type == 'heading':
                cursor = self.conn.execute("""
                    SELECT MAX("index") as max_index FROM document_blocks
                    WHERE type = ? AND level = ?
                """, (new_type, new_level))
            else:
                cursor = self.conn.execute("""
                    SELECT MAX("index") as max_index FROM document_blocks
                    WHERE type = ?
                """, (new_type,))
            
            row = cursor.fetchone()
            max_index = row['max_index'] if row['max_index'] is not None else -100
            new_index = max_index + 100
        
        return new_index
    
    def calculate_block_order_for_insert(self, after_block_id: str = None) -> int:
        """计算插入 block 的 block_order（稀疏排序）
        
        Args:
            after_block_id: 在哪个 block 后插入（如果为 None 则插入到末尾）
            
        Returns:
            计算得到的 block_order
        """
        if after_block_id:
            # 1. 获取 after_block 的 block_order
            after_block = self.get_block_by_id(after_block_id)
            if not after_block:
                raise ValueError(f"Block not found: {after_block_id}")
            after_order = after_block['block_order']
            
            # 2. 查询下一个 block 的 block_order
            cursor = self.conn.execute("""
                SELECT block_order FROM document_blocks
                WHERE block_order > ?
                ORDER BY block_order LIMIT 1
            """, (after_order,))
            
            next_row = cursor.fetchone()
            
            if next_row:
                next_order = next_row['block_order']
                gap = next_order - after_order
                
                # 3. 检查间隙是否足够
                if gap <= 1:
                    # 触发局部重排
                    self._rebalance_block_orders_between(after_order, next_order)
                    # 重新查询
                    cursor = self.conn.execute("""
                        SELECT block_order FROM document_blocks
                        WHERE block_order > ?
                        ORDER BY block_order LIMIT 1
                    """, (after_order,))
                    next_row = cursor.fetchone()
                    next_order = next_row['block_order'] if next_row else after_order + 200
                
                # 4. 计算中间值
                return (after_order + next_order) // 2
            else:
                # 没有下一个 block，插入到最后
                return after_order + 100
        else:
            # 插入到文档末尾
            cursor = self.conn.execute("""
                SELECT MAX(block_order) as max_order FROM document_blocks
            """)
            row = cursor.fetchone()
            max_order = row['max_order'] if row['max_order'] is not None else 0
            return max_order + 100
    
    def _rebalance_indexes_between(
        self,
        block_type: str,
        level: int,
        start_index: int,
        end_index: int
    ):
        """局部重排：重新分配区间内同类型同级别 blocks 的 index
        
        Args:
            block_type: Block 类型
            level: 级别（标题有效）
            start_index: 起始 index
            end_index: 结束 index
        """
        # 查询区间内的所有同类型同级别 blocks
        if block_type == 'heading':
            cursor = self.conn.execute("""
                SELECT id, "index" 
                FROM document_blocks 
                WHERE type = ? AND level = ? AND "index" > ? AND "index" < ?
                ORDER BY "index"
            """, (block_type, level, start_index, end_index))
        else:
            cursor = self.conn.execute("""
                SELECT id, "index" 
                FROM document_blocks 
                WHERE type = ? AND "index" > ? AND "index" < ?
                ORDER BY "index"
            """, (block_type, start_index, end_index))
        
        blocks = cursor.fetchall()
        
        if not blocks:
            return
        
        # 计算新的间隔
        count = len(blocks)
        gap = end_index - start_index
        step = gap // (count + 1)
        
        # 重新分配 index
        new_index = start_index
        for block in blocks:
            new_index += step
            self.conn.execute(
                'UPDATE document_blocks SET "index" = ? WHERE id = ?',
                (new_index, block['id'])
            )
        
        self.conn.commit()
    
    def _rebalance_block_orders_between(self, start_order: int, end_order: int):
        """局部重排：重新分配区间内所有 blocks 的 block_order
        
        Args:
            start_order: 起始 block_order
            end_order: 结束 block_order
        """
        # 查询区间内的所有 blocks
        cursor = self.conn.execute("""
            SELECT id, block_order 
            FROM document_blocks 
            WHERE block_order > ? AND block_order < ?
            ORDER BY block_order
        """, (start_order, end_order))
        
        blocks = cursor.fetchall()
        
        if not blocks:
            return
        
        # 计算新的间隔
        count = len(blocks)
        gap = end_order - start_order
        step = gap // (count + 1)
        
        # 重新分配 block_order
        new_order = start_order
        for block in blocks:
            new_order += step
            self.conn.execute(
                'UPDATE document_blocks SET block_order = ? WHERE id = ?',
                (new_order, block['id'])
            )
        
        self.conn.commit()
    
    def _row_to_dict(self, row) -> dict:
        """将 sqlite3.Row 转换为字典"""
        d = dict(row)
        
        # 解析 JSON 字段
        if d.get('style'):
            try:
                d['style'] = json.loads(d['style'])
            except (json.JSONDecodeError, TypeError):
                d['style'] = {}
        else:
            d['style'] = {}
        
        if d.get('metadata'):
            try:
                d['metadata'] = json.loads(d['metadata'])
            except (json.JSONDecodeError, TypeError):
                d['metadata'] = {}
        else:
            d['metadata'] = {}
        
        # content 可能是 JSON（富文本或表格）
        if d.get('content'):
            try:
                d['content'] = json.loads(d['content'])
            except (json.JSONDecodeError, TypeError):
                pass  # 保持原字符串
        
        return d
    
    @staticmethod
    def generate_block_id(block_type: str, level: int, index: int) -> str:
        """生成 Block ID
        
        Args:
            block_type: Block 类型
            level: 级别（标题有效）
            index: 序号
            
        Returns:
            生成的 Block ID
        """
        if block_type == 'heading':
            return f'block-h{level}-{index}'
        elif block_type == 'paragraph':
            return f'block-p-{index}'
        elif block_type == 'table':
            return f'block-table-{index}'
        elif block_type == 'image':
            return f'block-img-{index}'
        elif block_type == 'toc':
            return f'block-toc-{index}'
        else:
            return f'block-{block_type}-{index}'