#!/usr/bin/env python3
"""
FTS5 记忆索引系统
用于全文搜索 memory/*.md 文件
"""

import sqlite3
import os
import re
from pathlib import Path
from datetime import datetime

# 配置
MEMORY_DIR = Path.home() / ".openclaw" / "workspace" / "memory"
DB_PATH = MEMORY_DIR / ".fts5" / "memory.db"


def init_db():
    """初始化 FTS5 数据库"""
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 创建文件表
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS files (
            id INTEGER PRIMARY KEY,
            path TEXT UNIQUE NOT NULL,
            modified_at TIMESTAMP,
            indexed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)
    
    # 创建 FTS5 虚拟表用于全文搜索
    cursor.execute("""
        CREATE VIRTUAL TABLE IF NOT EXISTS memory_index USING fts5(
            content,
            file_id,
            chunk_index,
            tokenize='porter unicode61'
        )
    """)
    
    # 创建 chunks 表存储分块信息
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS chunks (
            id INTEGER PRIMARY KEY,
            file_id INTEGER,
            chunk_index INTEGER,
            content TEXT,
            start_line INTEGER,
            end_line INTEGER,
            FOREIGN KEY (file_id) REFERENCES files(id)
        )
    """)
    
    conn.commit()
    conn.close()
    print(f"✓ 数据库初始化完成: {DB_PATH}")


def chunk_content(content, chunk_size=500, overlap=100):
    """
    将内容分块，带重叠以保持上下文
    
    Args:
        content: 文件内容
        chunk_size: 每块字符数
        overlap: 重叠字符数
    
    Returns:
        list of (chunk_text, start_line, end_line)
    """
    lines = content.split('\n')
    chunks = []
    
    current_chunk = []
    current_size = 0
    start_line = 0
    
    for i, line in enumerate(lines):
        line_with_newline = line + '\n'
        
        if current_size + len(line_with_newline) > chunk_size and current_chunk:
            # 保存当前块
            chunk_text = ''.join(current_chunk)
            chunks.append((chunk_text, start_line + 1, i))
            
            # 重叠部分
            overlap_text = chunk_text[-overlap:] if len(chunk_text) > overlap else chunk_text
            current_chunk = [overlap_text]
            current_size = len(overlap_text)
            start_line = i - len(overlap_text.split('\n')) + 1
        
        current_chunk.append(line_with_newline)
        current_size += len(line_with_newline)
    
    # 最后一块
    if current_chunk:
        chunk_text = ''.join(current_chunk)
        chunks.append((chunk_text, start_line + 1, len(lines)))
    
    return chunks


def index_file(file_path, conn=None):
    """索引单个文件"""
    close_conn = False
    if conn is None:
        conn = sqlite3.connect(DB_PATH)
        close_conn = True
    
    cursor = conn.cursor()
    
    # 读取文件
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except Exception as e:
        print(f"✗ 读取失败 {file_path}: {e}")
        return
    
    # 获取文件信息
    stat = os.stat(file_path)
    modified_at = datetime.fromtimestamp(stat.st_mtime)
    rel_path = str(Path(file_path).relative_to(MEMORY_DIR.parent))
    
    # 插入或更新文件记录
    cursor.execute("""
        INSERT OR REPLACE INTO files (path, modified_at, indexed_at)
        VALUES (?, ?, CURRENT_TIMESTAMP)
    """, (rel_path, modified_at))
    
    file_id = cursor.lastrowid
    
    # 删除旧索引
    cursor.execute("DELETE FROM memory_index WHERE file_id = ?", (file_id,))
    cursor.execute("DELETE FROM chunks WHERE file_id = ?", (file_id,))
    
    # 分块并索引
    chunks = chunk_content(content)
    
    for chunk_idx, (chunk_text, start_line, end_line) in enumerate(chunks):
        # 插入 chunks 表
        cursor.execute("""
            INSERT INTO chunks (file_id, chunk_index, content, start_line, end_line)
            VALUES (?, ?, ?, ?, ?)
        """, (file_id, chunk_idx, chunk_text, start_line, end_line))
        
        # 插入 FTS5 索引
        cursor.execute("""
            INSERT INTO memory_index (content, file_id, chunk_index)
            VALUES (?, ?, ?)
        """, (chunk_text, file_id, chunk_idx))
    
    conn.commit()
    
    if close_conn:
        conn.close()
    
    print(f"✓ 已索引: {rel_path} ({len(chunks)} 块)")


def index_all():
    """索引所有 memory 文件"""
    init_db()
    
    conn = sqlite3.connect(DB_PATH)
    
    # 查找所有 .md 文件
    md_files = list(MEMORY_DIR.glob("*.md"))
    
    print(f"\n开始索引 {len(md_files)} 个文件...")
    
    for file_path in md_files:
        if file_path.name.startswith('.'):
            continue
        index_file(file_path, conn)
    
    conn.close()
    print(f"\n✓ 索引完成")


def search(query, limit=10):
    """
    全文搜索
    
    Args:
        query: 搜索关键词
        limit: 返回结果数量
    
    Returns:
        list of dict with path, snippet, rank
    """
    conn = sqlite3.connect(DB_PATH)
    cursor = conn.cursor()
    
    # 使用 FTS5 的 snippet 功能生成摘要
    cursor.execute("""
        SELECT 
            f.path,
            snippet(memory_index, 0, '[', ']', '...', 32) as snippet,
            c.start_line,
            c.end_line,
            rank
        FROM memory_index
        JOIN files f ON memory_index.file_id = f.id
        JOIN chunks c ON memory_index.file_id = c.file_id 
                     AND memory_index.chunk_index = c.chunk_index
        WHERE memory_index MATCH ?
        ORDER BY rank
        LIMIT ?
    """, (query, limit))
    
    results = []
    for row in cursor.fetchall():
        results.append({
            'path': row[0],
            'snippet': row[1],
            'start_line': row[2],
            'end_line': row[3],
            'rank': row[4]
        })
    
    conn.close()
    return results


def print_search_results(query, results):
    """打印搜索结果"""
    print(f"\n🔍 搜索: '{query}'")
    print(f"找到 {len(results)} 条结果:\n")
    
    for i, r in enumerate(results, 1):
        print(f"{i}. {r['path']} (行 {r['start_line']}-{r['end_line']})")
        print(f"   {r['snippet']}")
        print()


if __name__ == "__main__":
    import sys
    
    if len(sys.argv) < 2:
        print("用法:")
        print(f"  {sys.argv[0]} index       # 索引所有文件")
        print(f"  {sys.argv[0]} search <query>  # 搜索")
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "index":
        index_all()
    elif command == "search":
        if len(sys.argv) < 3:
            print("错误: 需要提供搜索关键词")
            sys.exit(1)
        query = sys.argv[2]
        results = search(query)
        print_search_results(query, results)
    else:
        print(f"未知命令: {command}")
        sys.exit(1)