#!/usr/bin/env python3
"""
方法学提取 API v2 - 支持9个阶段（完整版）
"""

import os
import sys
import json
import uuid
from pathlib import Path
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from typing import Dict, List

sys.path.insert(0, str(Path(__file__).parent.parent))

from src.rag.extractors.methodology.simple_extractor import SimpleExtractor

app = FastAPI(title="Methodology Extraction API v2 - Full Data")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

results = {}


@app.post("/extract")
async def extract_methodology(pdf: UploadFile = File(...)):
    """从 PDF 提取方法学参数 - 返回完整数据"""
    
    job_id = str(uuid.uuid4())[:8]
    
    upload_dir = Path(__file__).parent.parent / 'uploads'
    upload_dir.mkdir(exist_ok=True)
    pdf_path = upload_dir / f"{job_id}_{pdf.filename}"
    
    with open(pdf_path, "wb") as f:
        content = await pdf.read()
        f.write(content)
    
    try:
        extractor = SimpleExtractor()
        result_text = extractor.extract(pdf_path)
        
        structured = parse_extraction_result(result_text)
        
        # TODO: 修复 parse_extraction_result 真正解析 LLM 结果
        # 临时方案：返回原始文本 + 尝试解析
        structured = parse_extraction_result(result_text)
        
        results[job_id] = {
            'status': 'completed',
            'pdf_path': str(pdf_path),
            'raw_text': result_text,
            'data': structured
        }
        
        return {
            'job_id': job_id,
            'status': 'completed',
            'raw_text': result_text,  # 添加原始文本供调试
            'data': structured
        }
        
    except Exception as e:
        results[job_id] = {'status': 'failed', 'error': str(e)}
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/result/{job_id}")
async def get_result(job_id: str):
    """获取提取结果"""
    if job_id not in results:
        raise HTTPException(status_code=404, detail="Job not found")
    return results[job_id]


def parse_extraction_result(text: str) -> Dict:
    """
    解析提取结果（9个阶段）- 从 LLM JSON 输出解析
    """
    import re
    
    # 尝试从文本中提取 JSON
    try:
        # 查找 ```json ... ``` 代码块
        json_match = re.search(r'```json\s*(.*?)\s*```', text, re.DOTALL)
        if json_match:
            json_str = json_match.group(1)
        else:
            # 尝试直接解析整个文本为 JSON
            json_str = text.strip()
        
        parsed = json.loads(json_str)
        
        # 验证必要字段
        if 'species' in parsed and 'stages' in parsed:
            return parsed
    except Exception as e:
        print(f"JSON 解析失败: {e}")
    
    # 如果 JSON 解析失败，返回空结构
    return {
        'species': {'target': '', 'references': []},
        'gene_family': {'name': '', 'pfam_id': ''},
        'stages': []
    }


def _old_parse_extraction_result(text: str) -> Dict:
    """
    [已废弃] 解析提取结果（9个阶段）- 完整版
    从完整JSON模板加载，确保所有字段完整
    """
    # 加载完整JSON模板
    template_path = Path(__file__).parent.parent / 'output' / 'spl_methods_extraction_complete.json'
    if template_path.exists():
        with open(template_path, 'r', encoding='utf-8') as f:
            full_data = json.load(f)
    else:
        full_data = {}
    
    # 构建前端需要的格式
    data = {
        'species': {
            'target': full_data.get('species', {}).get('target', {}).get('scientific_name', ''),
            'references': [s.get('name', '') for s in full_data.get('species', {}).get('references', [])]
        },
        'gene_family': {
            'name': full_data.get('gene_family', {}).get('name', ''),
            'full_name': full_data.get('gene_family', {}).get('full_name', ''),
            'pfam_id': full_data.get('gene_family', {}).get('pfam_id', ''),
            'domain': full_data.get('gene_family', {}).get('domain', ''),
            'gene_count': full_data.get('gene_family', {}).get('gene_count', 0)
        },
        'stages': []
    }
    
    # 从完整数据提取各阶段
    protein_analysis = full_data.get('protein_sequence_analysis', {})
    chromosomal = full_data.get('chromosomal_analysis', {})
    expression = full_data.get('expression_analysis', {})
    statistics = full_data.get('statistical_analysis', {})
    
    # 1. 基因家族鉴定
    identification = protein_analysis.get('identification', {})
    data['stages'].append({
        'id': 'identify',
        'name': '1. 基因家族鉴定',
        'params': [
            {'name': 'tool', 'label': '工具', 'value': ', '.join(identification.get('tools', []))},
            {'name': 'genome_source', 'label': '基因组来源', 'value': full_data.get('species', {}).get('target', {}).get('genome_database', '')},
            {'name': 'genome_url', 'label': '基因组URL', 'value': full_data.get('species', {}).get('target', {}).get('genome_url', '')},
            {'name': 'pfam_id', 'label': 'Pfam ID', 'value': identification.get('domain', '').split('(')[-1].replace(')', '') if '(' in identification.get('domain', '') else ''},
            {'name': 'pfam_url', 'label': 'Pfam URL', 'value': 'http://pfam.xfam.org/'},
            {'name': 'reference_sequences', 'label': '参考序列', 'value': identification.get('reference_sequences', '')},
            {'name': 'reference_db', 'label': '参考数据库', 'value': full_data.get('species', {}).get('references', [{}])[0].get('database', '')},
            {'name': 'ref_seq_length', 'label': '参考序列长度', 'value': '131aa-1035aa'},
            {'name': 'e_value', 'label': 'E-value', 'value': '原文未明确'},
            {'name': 'identity', 'label': 'Identity', 'value': '原文未明确'},
            {'name': 'coverage', 'label': 'Coverage', 'value': '原文未明确'},
            {'name': 'filter_criteria', 'label': '筛选标准', 'value': identification.get('filtering', '')},
            {'name': 'final_count', 'label': '最终基因数', 'value': str(identification.get('final_count', ''))},
        ]
    })
    
    # 2. 蛋白理化性质
    pep = protein_analysis.get('physicochemical_properties', {})
    pep_results = pep.get('results', {})
    data['stages'].append({
        'id': 'pep',
        'name': '2. 蛋白理化性质',
        'params': [
            {'name': 'tool', 'label': '工具', 'value': pep.get('tool', '')},
            {'name': 'url', 'label': 'URL', 'value': pep.get('url', '')},
            {'name': 'properties_analyzed', 'label': '分析内容', 'value': ', '.join(pep.get('properties_analyzed', []))},
            {'name': 'protein_length', 'label': '蛋白长度', 'value': pep_results.get('protein_length', '')},
            {'name': 'molecular_weight', 'label': '分子量', 'value': pep_results.get('molecular_weight', '')},
            {'name': 'isoelectric_point', 'label': '等电点 PI', 'value': pep_results.get('isoelectric_point', '')},
            {'name': 'subcellular_localization', 'label': '亚细胞定位', 'value': pep_results.get('subcellular_localization', '')},
        ]
    })
    
    # 3. 系统发育分析
    phylo = protein_analysis.get('phylogenetic_analysis', {})
    data['stages'].append({
        'id': 'phylogeny',
        'name': '3. 系统发育分析',
        'params': [
            {'name': 'software', 'label': '软件', 'value': f"{phylo.get('software', '')} {phylo.get('version', '')}".strip()},
            {'name': 'alignment_tool', 'label': '比对工具', 'value': f"{protein_analysis.get('sequence_alignment', {}).get('tool', '')} {protein_analysis.get('sequence_alignment', {}).get('version', '')}".strip()},
            {'name': 'method', 'label': '建树方法', 'value': phylo.get('method', '').split('(')[0].strip() if '(' in phylo.get('method', '') else phylo.get('method', '')},
            {'name': 'bootstrap', 'label': 'Bootstrap', 'value': str(phylo.get('bootstrap', ''))},
            {'name': 'model', 'label': '替代模型', 'value': phylo.get('substitution_model', '')},
            {'name': 'comparison', 'label': '比较对象', 'value': phylo.get('comparison', '')},
            {'name': 'grouping', 'label': '分组结果', 'value': phylo.get('grouping', '')},
            {'name': 'reference_species', 'label': '参考物种', 'value': ', '.join([s.get('name', '') for s in full_data.get('species', {}).get('references', [])])},
            {'name': 'database', 'label': '序列数据库', 'value': full_data.get('species', {}).get('references', [{}])[0].get('database', '')},
        ]
    })
    
    # 4. 基因结构分析
    gene_struct = protein_analysis.get('gene_structure', {})
    data['stages'].append({
        'id': 'structure',
        'name': '4. 基因结构分析',
        'params': [
            {'name': 'tool', 'label': '工具', 'value': gene_struct.get('tool', '')},
            {'name': 'url', 'label': 'URL', 'value': gene_struct.get('url', '')},
            {'name': 'method', 'label': '方法', 'value': gene_struct.get('method', '')},
            {'name': 'analysis', 'label': '分析内容', 'value': gene_struct.get('analysis', '')},
        ]
    })
    
    # 5. 保守基序分析
    motifs = protein_analysis.get('conserved_domains', {}).get('conserved_motifs', {})
    data['stages'].append({
        'id': 'motif',
        'name': '5. 保守基序分析',
        'params': [
            {'name': 'tool', 'label': '工具', 'value': motifs.get('tool', '')},
            {'name': 'url', 'label': 'URL', 'value': motifs.get('url', '')},
            {'name': 'mode', 'label': '模式', 'value': motifs.get('parameters', {}).get('mode', '')},
            {'name': 'max_motifs', 'label': 'Max motifs', 'value': str(motifs.get('parameters', {}).get('max_motifs', ''))},
            {'name': 'width', 'label': 'Width', 'value': motifs.get('parameters', {}).get('width', '')},
            {'name': 'detected_motifs', 'label': '检测到的motifs', 'value': motifs.get('detected_motifs', '')},
            {'name': 'conserved_sequences', 'label': '保守序列', 'value': ', '.join(motifs.get('conserved_sequences', []))},
        ]
    })
    
    # 6. 染色体定位
    data['stages'].append({
        'id': 'karyotype',
        'name': '6. 染色体定位',
        'params': [
            {'name': 'mapping_method', 'label': '定位方法', 'value': chromosomal.get('mapping_method', '')},
            {'name': 'location', 'label': '定位结果', 'value': chromosomal.get('location', '')},
        ]
    })
    
    # 7. 共线性分析
    duplication = chromosomal.get('duplication_analysis', {})
    data['stages'].append({
        'id': 'synteny',
        'name': '7. 共线性分析',
        'params': [
            {'name': 'tool', 'label': '工具', 'value': duplication.get('tool', '')},
            {'name': 'full_name', 'label': '完整名称', 'value': duplication.get('full_name', '')},
            {'name': 'synteny_tool', 'label': '共线性绘图工具', 'value': duplication.get('synteny_tool', '')},
            {'name': 'synteny_url', 'label': '共线性工具URL', 'value': duplication.get('synteny_url', '')},
            {'name': 'segmental_duplications', 'label': '片段复制数', 'value': str(duplication.get('results', {}).get('segmental_duplications', ''))},
            {'name': 'tandem_duplications', 'label': '串联复制数', 'value': str(duplication.get('results', {}).get('tandem_duplications', ''))},
        ]
    })
    
    # 8. 表达分析
    primer_design = expression.get('primer_design', {})
    primer_params = primer_design.get('parameters', {})
    data['stages'].append({
        'id': 'expression',
        'name': '8. 表达分析',
        'params': [
            {'name': 'method', 'label': '方法', 'value': expression.get('method', '')},
            {'name': 'full_name', 'label': '方法全称', 'value': expression.get('full_name', '')},
            {'name': 'kit', 'label': '试剂盒', 'value': expression.get('kit', '')},
            {'name': 'reference_gene', 'label': '内参基因', 'value': expression.get('reference_gene', '')},
            {'name': 'genes_analyzed', 'label': '分析基因数', 'value': str(expression.get('genes_analyzed', ''))},
            {'name': 'primer_tool', 'label': '引物设计工具', 'value': primer_design.get('tool', '')},
            {'name': 'primer_url', 'label': '引物工具URL', 'value': primer_design.get('url', '')},
            {'name': 'product_length', 'label': '产物长度', 'value': primer_params.get('product_length', '')},
            {'name': 'gc_content', 'label': 'GC含量', 'value': primer_params.get('gc_content', '')},
            {'name': 'melting_temp', 'label': '退火温度', 'value': primer_params.get('melting_temperature', '')},
            {'name': 'primer_source', 'label': '引物序列来源', 'value': primer_design.get('primer_sequences_file', '')},
            {'name': 'data_analysis', 'label': '数据分析方法', 'value': expression.get('data_analysis', '')},
            {'name': 'tissues', 'label': '分析组织', 'value': ', '.join(expression.get('tissues_analyzed', []))},
            {'name': 'stages', 'label': '发育阶段', 'value': ', '.join(expression.get('developmental_stages', []))},
        ]
    })
    
    # 9. 统计分析
    data['stages'].append({
        'id': 'statistics',
        'name': '9. 统计分析',
        'params': [
            {'name': 'software', 'label': '软件', 'value': statistics.get('software', '')},
            {'name': 'company', 'label': '公司', 'value': statistics.get('company', '')},
            {'name': 'method', 'label': '方法', 'value': ', '.join(statistics.get('methods', []))},
            {'name': 'significance_levels', 'label': '显著性水平', 'value': ', '.join(statistics.get('significance_levels', []))},
        ]
    })
    
    return data


if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5002)