PlantAI-rag/api/species_routes.py

"""
物种查询API路由
提供基因家族-物种索引查询功能
"""

from fastapi import APIRouter, HTTPException, Query
from pydantic import BaseModel
from typing import List, Dict, Optional
from pathlib import Path
import json

from src.rag.project_config import PROJECT_ROOT

router = APIRouter()

INDEX_FILE = PROJECT_ROOT / "data" / "gene_family_species_index.json"


class SpeciesInfo(BaseModel):
    """物种信息"""
    paper_id: str
    scientific_name: str
    common_name_en: str
    common_name_cn: str
    comparison_species: List[str]
    relations_count: int
    entities_count: int
    pmid: Optional[str]


class GeneFamilyInfo(BaseModel):
    """基因家族信息"""
    gene_family: str
    description: str
    description_cn: str
    total_species: int
    species_list: List[SpeciesInfo]


class SpeciesListResponse(BaseModel):
    """物种列表响应"""
    gene_family: str
    total_species: int
    species: List[SpeciesInfo]


class SpeciesQueryResponse(BaseModel):
    """物种查询响应"""
    query: str
    matched_species: List[Dict]
    total_matches: int


def load_species_index() -> dict:
    """加载物种索引"""
    if not INDEX_FILE.exists():
        raise HTTPException(status_code=500, detail="物种索引文件不存在，请先运行索引构建脚本")
    
    with open(INDEX_FILE, 'r', encoding='utf-8') as f:
        return json.load(f)


@router.get("/species/list", response_model=SpeciesListResponse)
async def list_species_by_gene_family(
    gene_family: str = Query("SPL", description="基因家族名称，如SPL、NAC、WRKY等")
):
    """
    获取指定基因家族已分析的所有物种列表
    
    示例:
    - /api/v1/species/list?gene_family=SPL
    """
    try:
        index = load_species_index()
        
        gene_family_upper = gene_family.upper()
        
        if gene_family_upper not in index.get("gene_families", {}):
            return SpeciesListResponse(
                gene_family=gene_family,
                total_species=0,
                species=[]
            )
        
        gf_info = index["gene_families"][gene_family_upper]
        
        # 构建响应
        species_list = []
        for species in gf_info.get("species_list", []):
            species_list.append(SpeciesInfo(
                paper_id=species["paper_id"],
                scientific_name=species["scientific_name"],
                common_name_en=species["common_name_en"],
                common_name_cn=species["common_name_cn"],
                comparison_species=species["comparison_species"],
                relations_count=species["relations_count"],
                entities_count=species["entities_count"],
                pmid=species.get("pmid") or None
            ))
        
        return SpeciesListResponse(
            gene_family=gene_family_upper,
            total_species=gf_info["total_species"],
            species=species_list
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@router.get("/species/search", response_model=SpeciesQueryResponse)
async def search_species(
    query: str = Query(..., description="物种查询关键词，支持中文、英文、拉丁学名")
):
    """
    搜索特定物种是否做过基因家族分析
    
    示例:
    - /api/v1/species/search?query=苦荞
    - /api/v1/species/search?query=tartary
    - /api/v1/species/search?query=Fagopyrum
    """
    try:
        index = load_species_index()
        
        query_lower = query.lower()
        matches = []
        
        # 遍历所有基因家族
        for gf_name, gf_info in index.get("gene_families", {}).items():
            for species in gf_info.get("species_list", []):
                # 匹配多个字段
                match_fields = [
                    species["scientific_name"].lower(),
                    species["common_name_en"].lower(),
                    species["common_name_cn"]
                ]
                
                if any(query_lower in field for field in match_fields):
                    matches.append({
                        "gene_family": gf_name,
                        "paper_id": species["paper_id"],
                        "scientific_name": species["scientific_name"],
                        "common_name_en": species["common_name_en"],
                        "common_name_cn": species["common_name_cn"],
                        "pmid": species.get("pmid"),
                        "relations_count": species["relations_count"],
                        "entities_count": species["entities_count"]
                    })
        
        return SpeciesQueryResponse(
            query=query,
            matched_species=matches,
            total_matches=len(matches)
        )
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@router.get("/species/gene-families", response_model=List[str])
async def list_gene_families():
    """
    获取所有已索引的基因家族列表
    
    示例:
    - /api/v1/species/gene-families
    """
    try:
        index = load_species_index()
        return list(index.get("gene_families", {}).keys())
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")


@router.get("/species/stats")
async def get_statistics():
    """
    获取基因家族-物种索引统计信息
    
    示例:
    - /api/v1/species/stats
    """
    try:
        index = load_species_index()
        
        stats = {
            "total_gene_families": len(index.get("gene_families", {})),
            "total_papers": len(index.get("papers", {})),
            "total_species_indexed": len(index.get("species_index", {})),
            "metadata": index.get("metadata", {}),
            "gene_families": {}
        }
        
        for gf_name, gf_info in index.get("gene_families", {}).items():
            stats["gene_families"][gf_name] = {
                "description_cn": gf_info.get("description_cn", ""),
                "total_species": gf_info.get("total_species", 0),
                "species_names": [s["common_name_cn"] for s in gf_info.get("species_list", [])]
            }
        
        return stats
        
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"查询失败: {str(e)}")