#!/usr/bin/env python3
"""
构建 RAG 索引: 文本分块 → JSON 索引文件
简单但实用的 BM25 风格索引（无需向量数据库）
"""
import json
import os
import re
from pathlib import Path

PROCESSED_DIR = Path('/home/67/knowledge/maritime/processed')
INDEX_FILE = Path('/home/67/knowledge/maritime/rag_index.json')

CHUNK_SIZE = 512      # 每块字符数
CHUNK_OVERLAP = 50    # 重叠字符数

def split_chunks(text, source_file, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
    """将文本分割为带重叠的块"""
    chunks = []
    # 按段落优先分割
    paragraphs = re.split(r'\n\s*\n', text)
    
    current_chunk = ''
    chunk_id = 0
    
    for para in paragraphs:
        para = para.strip()
        if not para:
            continue
        
        if len(current_chunk) + len(para) <= chunk_size:
            current_chunk += '\n' + para if current_chunk else para
        else:
            if current_chunk:
                chunks.append({
                    'id': f'{source_file}#chunk{chunk_id}',
                    'source': source_file,
                    'chunk_id': chunk_id,
                    'text': current_chunk,
                    'length': len(current_chunk),
                })
                chunk_id += 1
                # 重叠: 保留上一块末尾
                overlap_text = current_chunk[-overlap:] if len(current_chunk) > overlap else ''
                current_chunk = overlap_text + '\n' + para
            else:
                # 单段落超长，按字符切
                for i in range(0, len(para), chunk_size - overlap):
                    chunk_text = para[i:i+chunk_size]
                    chunks.append({
                        'id': f'{source_file}#chunk{chunk_id}',
                        'source': source_file,
                        'chunk_id': chunk_id,
                        'text': chunk_text,
                        'length': len(chunk_text),
                    })
                    chunk_id += 1
                current_chunk = ''
    
    # 最后一块
    if current_chunk.strip():
        chunks.append({
            'id': f'{source_file}#chunk{chunk_id}',
            'source': source_file,
            'chunk_id': chunk_id,
            'text': current_chunk,
            'length': len(current_chunk),
        })
    
    return chunks

# 构建索引
all_chunks = []
txt_files = sorted(PROCESSED_DIR.glob('*.txt'))
print(f'处理 {len(txt_files)} 个文本文件...')

for i, txt_path in enumerate(txt_files):
    text = txt_path.read_text(encoding='utf-8')
    if '[扫描版PDF' in text:
        continue
    
    chunks = split_chunks(text, txt_path.stem)
    all_chunks.extend(chunks)
    
    if (i+1) % 200 == 0:
        print(f'  进度: {i+1}/{len(txt_files)}, 累计块数: {len(all_chunks)}')

# 构建关键词倒排索引（简单 BM25 近似）
print(f'构建倒排索引... ({len(all_chunks)} 个文本块)')

keyword_index = {}
for chunk in all_chunks:
    # 分词: 中文按字符 bigram + 英文按空格
    words = set()
    text = chunk['text'].lower()
    # 中文 bigram
    chinese_chars = re.findall(r'[\u4e00-\u9fff]+', text)
    for cc in chinese_chars:
        for j in range(len(cc) - 1):
            words.add(cc[j:j+2])
        if len(cc) >= 3:
            for j in range(len(cc) - 2):
                words.add(cc[j:j+3])
    # 英文单词
    eng_words = re.findall(r'[a-z]+', text)
    words.update(eng_words)
    
    for w in words:
        if w not in keyword_index:
            keyword_index[w] = []
        keyword_index[w].append(chunk['id'])

# 保存索引
index_data = {
    'total_chunks': len(all_chunks),
    'total_files': len(txt_files),
    'chunk_size': CHUNK_SIZE,
    'chunks': all_chunks,
    'keyword_index_size': len(keyword_index),
}

with open(INDEX_FILE, 'w', encoding='utf-8') as f:
    json.dump(index_data, f, ensure_ascii=False)

# 单独保存倒排索引
with open(INDEX_FILE.with_suffix('.keywords.json'), 'w', encoding='utf-8') as f:
    json.dump(keyword_index, f, ensure_ascii=False)

print(f'完成!')
print(f'  文本块: {len(all_chunks)}')
print(f'  关键词: {len(keyword_index)}')
print(f'  索引文件: {INDEX_FILE} ({INDEX_FILE.stat().st_size / 1024 / 1024:.1f} MB)')