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

"""SQLite 内容数据库操作类"""
import json
import sqlite3
from pathlib import Path
from typing import Optional


class ContentDB:
    """SQLite 内容数据库操作类，每个文档对应一个独立的 SQLite 数据库文件"""
    
    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 _build_type_filter(self, block_type: str, level: Optional[int] = None) -> tuple[str, list]:
        """构建类型过滤条件（用于 index 查询）"""
        if block_type == 'heading' and level is not None:
            return "type = ? AND level = ?", [block_type, level]
        return "type = ?", [block_type]
    
    def _query_next_value(self, field: str, after_order: int, type_filter: str, params: list) -> Optional[int]:
        """查询下一个值（index 或 block_order）"""
        sql = f'SELECT "{field}" FROM document_blocks WHERE {type_filter} AND block_order > ? ORDER BY block_order LIMIT 1'
        cursor = self.conn.execute(sql, params + [after_order])
        row = cursor.fetchone()
        return row[field] if row else None
    
    def _query_prev_value(self, field: str, after_order: int, type_filter: str, params: list) -> Optional[int]:
        """查询前一个值（index 或 block_order）"""
        sql = f'SELECT "{field}" FROM document_blocks WHERE {type_filter} AND block_order <= ? ORDER BY block_order DESC LIMIT 1'
        cursor = self.conn.execute(sql, params + [after_order])
        row = cursor.fetchone()
        return row[field] if row else None
    
    def _query_max_value(self, field: str, type_filter: str, params: list) -> Optional[int]:
        """查询最大值（index 或 block_order）"""
        sql = f'SELECT MAX("{field}") as max_val FROM document_blocks'
        if type_filter:
            sql += f' WHERE {type_filter}'
        cursor = self.conn.execute(sql, params)
        row = cursor.fetchone()
        return row['max_val']
    
    def _calculate_sparse_value(
        self, 
        field: str,
        after_block_id: Optional[str],
        type_filter: str,
        params: list,
        default_min: int,
        rebalance_func: callable
    ) -> int:
        """通用稀疏值计算逻辑"""
        if after_block_id:
            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']
            
            next_val = self._query_next_value(field, after_order, type_filter, params)
            
            if next_val is not None:
                prev_val = self._query_prev_value(field, after_order, type_filter, params)
                if prev_val is None:
                    prev_val = default_min
                
                gap = next_val - prev_val
                if gap <= 1:
                    rebalance_func(prev_val, next_val)
                    next_val = self._query_next_value(field, after_order, type_filter, params)
                    if next_val is None:
                        next_val = prev_val + 200
                
                return (prev_val + next_val) // 2
            else:
                max_val = self._query_max_value(field, type_filter, params)
                return (max_val if max_val is not None else default_min) + 100
        else:
            max_val = self._query_max_value(field, type_filter, params)
            return (max_val if max_val is not None else default_min) + 100
    
    def calculate_index_for_insert(self, after_block_id: str, new_type: str, new_level: int = 0) -> int:
        """计算插入 block 的 index（稀疏排序）"""
        type_filter, params = self._build_type_filter(new_type, new_level)
        return self._calculate_sparse_value(
            field='index',
            after_block_id=after_block_id,
            type_filter=type_filter,
            params=params,
            default_min=-100,
            rebalance_func=lambda prev, next: self._rebalance_indexes_between(new_type, new_level, prev, next)
        )
    
    def calculate_block_order_for_insert(self, after_block_id: str = None) -> int:
        """计算插入 block 的 block_order（稀疏排序）"""
        return self._calculate_sparse_value(
            field='block_order',
            after_block_id=after_block_id,
            type_filter='1=1',
            params=[],
            default_min=0,
            rebalance_func=self._rebalance_block_orders_between
        )
    
    def _rebalance_indexes_between(self, block_type: str, level: int, start_index: int, end_index: int):
        """局部重排：重新分配区间内同类型同级别 blocks 的 index"""
        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)
        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"""
        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)
        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)
        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'] = {}
        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"""
        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}'