#!/usr/bin/env python3
"""
按章节分块存储 - 参考 PlantConnectome 格式
"""

import fitz
import json


def extract_sections(pdf_path: str) -> dict:
    """提取PDF各章节"""
    doc = fitz.open(pdf_path)
    
    # 读取全文
    full_text = ""
    for page in doc:
        full_text += page.get_text() + "\n"
    doc.close()
    
    # 简单章节识别
    sections = {
        'full_text': full_text,
        'abstract': '',
        'background': '',
        'methods': '',
        'results_phylogeny': '',
        'results_expression': '',
        'results_synteny': '',
        'discussion': ''
    }
    
    lines = full_text.split('\n')
    current_section = None
    current_text = []
    
    for line in lines:
        line_lower = line.lower().strip()
        
        # 检测章节
        if 'abstract' in line_lower and len(line_lower) < 20:
            if current_section:
                sections[current_section] = '\n'.join(current_text).strip()
            current_section = 'abstract'
            current_text = []
        elif any(x in line_lower for x in ['background', 'introduction']) and len(line_lower) < 30:
            if current_section:
                sections[current_section] = '\n'.join(current_text).strip()
            current_section = 'background'
            current_text = []
        elif 'methods' in line_lower and len(line_lower) < 20:
            if current_section:
                sections[current_section] = '\n'.join(current_text).strip()
            current_section = 'methods'
            current_text = []
        elif 'results' in line_lower and len(line_lower) < 20:
            if current_section:
                sections[current_section] = '\n'.join(current_text).strip()
            current_section = 'results_phylogeny'
            current_text = []
        elif 'discussion' in line_lower and len(line_lower) < 20:
            if current_section:
                sections[current_section] = '\n'.join(current_text).strip()
            current_section = 'discussion'
            current_text = []
        elif current_section:
            # Results 内子章节检测
            if current_section.startswith('results_'):
                if any(x in line_lower for x in ['phylogenetic', 'gene structure', 'conserved motif']):
                    sections[current_section] = '\n'.join(current_text).strip()
                    current_text = []
                    current_section = 'results_phylogeny'
                elif any(x in line_lower for x in ['expression pattern', 'tissue', 'qrt-pcr']):
                    sections[current_section] = '\n'.join(current_text).strip()
                    current_text = []
                    current_section = 'results_expression'
                elif 'synteny' in line_lower:
                    sections[current_section] = '\n'.join(current_text).strip()
                    current_text = []
                    current_section = 'results_synteny'
            
            current_text.append(line)
    
    # 保存最后一个章节
    if current_section and current_text:
        sections[current_section] = '\n'.join(current_text).strip()
    
    return sections


def main():
    pdf_path = "paper/test_v2_reproduce/input/苦荞SPL2019BMCplantbios12870-019-1916-6.pdf"
    pmid = "31286919"
    
    print("=" * 70)
    print("Extracting Sections (PlantConnectome Format)")
    print("=" * 70)
    
    sections = extract_sections(pdf_path)
    
    # 创建 PlantConnectome 格式的 chunks
    chunks = {}
    for section_name, text in sections.items():
        if text and len(text) > 100:
            chunk_id = f"{pmid}_{section_name}"
            chunks[chunk_id] = {
                "pmid": pmid,
                "section": section_name,
                "text": text[:2000],  # 限制长度
                "char_count": len(text)
            }
            print(f"\n✓ {chunk_id}: {len(text)} chars")
    
    # 保存
    output_path = "evaluation/chunks/pmid31286919_chunks.json"
    import os
    os.makedirs("evaluation/chunks", exist_ok=True)
    
    with open(output_path, 'w') as f:
        json.dump(chunks, f, indent=2)
    
    print(f"\n✓ Saved to: {output_path}")
    print(f"Total chunks: {len(chunks)}")


if __name__ == "__main__":
    main()