# 参考：https://www.ncbi.nlm.nih.gov/datasets/docs/v2/api/rest-api/#get-/biosample/accession/-accessions-/biosample_report
import aiohttp
import asyncio
import pandas as pd
import time
import os
import json
from datetime import datetime

def read_file(file_name):
    """读取TSV文件并提取biosample_id列表"""
    try:
        # 使用pandas读取TSV文件
        df = pd.read_csv(file_name, sep='\t')
        # 假设biosample ID列名为'biosample_id'，如果不是，请修改为实际的列名
        biosample_list = df['biosample_id'].tolist()
        print(f"成功读取 {len(biosample_list)} 个biosample_id")
        return biosample_list
    except Exception as e:
        print(f"读取文件时出错: {str(e)}")
        # 如果pandas读取失败，尝试原始方法
        try:
            with open(file_name, 'r', encoding='utf-8') as f:
                lines = f.readlines()
            biosample_list = []
            for i in lines[1:]:  # 跳过标题行
                parts = i.strip().split('\t')
                if parts:  # 确保行不为空
                    # 假设biosample ID在第一列，如果不是，请修改索引
                    biosample_id = parts[0].strip('"')
                    biosample_list.append(biosample_id)
            print(f"使用备用方法读取 {len(biosample_list)} 个biosample_id")
            return biosample_list
        except Exception as e2:
            print(f"备用方法也失败: {str(e2)}")
            return []

async def get_sample_info_batch(accessions_batch, session, df):
    """批量获取BioSample信息 - 使用NCBI Datasets API"""
    headers = {
        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36",
        "Accept": "application/json"
    }
    
    # 构建API URL
    accessions_str = ",".join(accessions_batch)
    url = f"https://api.ncbi.nlm.nih.gov/datasets/v2/biosample/accession/{accessions_str}/biosample_report"
    
    try:
        async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=30)) as resp:
            if resp.status == 200:
                data = await resp.json()
                
                if "reports" in data and data["reports"]:
                    print(f'成功获取 {len(data["reports"])} 个样本的信息')
                    
                    for report in data["reports"]:
                        accession = report.get("accession", "")
                        if not accession:
                            continue
                            
                        # 提取所有可用信息
                        extracted_info = extract_all_info(report)
                        
                        # 将数据添加到DataFrame
                        for k, v in extracted_info.items():
                            df.loc[accession, k] = v
                    
                    return True
                else:
                    print(f'警告: 未找到样本信息')
                    if "messages" in data:
                        for msg in data["messages"]:
                            if "error" in msg:
                                print(f"错误信息: {msg['error'].get('message', '未知错误')}")
                    return False
            else:
                print(f'错误: 获取样本信息时服务器返回状态码 {resp.status}')
                return False
    except asyncio.TimeoutError:
        print(f'错误: 获取样本信息时超时')
        return False
    except Exception as e:
        print(f'错误: 获取样本信息时发生异常: {str(e)}')
        return False

def extract_all_info(report):
    """从API响应中提取所有可用信息"""
    extracted_info = {}
    
    # 基本信息
    extracted_info['accession'] = report.get('accession', '')
    extracted_info['last_updated'] = report.get('last_updated', '')
    extracted_info['publication_date'] = report.get('publication_date', '')
    extracted_info['submission_date'] = report.get('submission_date', '')
    extracted_info['package'] = report.get('package', '')
    extracted_info['status'] = report.get('status', {}).get('status', '')
    extracted_info['status_when'] = report.get('status', {}).get('when', '')
    
    # 样本ID信息
    sample_ids = report.get('sample_ids', [])
    for i, sample_id in enumerate(sample_ids):
        extracted_info[f'sample_id_{i+1}_db'] = sample_id.get('db', '')
        extracted_info[f'sample_id_{i+1}_label'] = sample_id.get('label', '')
        extracted_info[f'sample_id_{i+1}_value'] = sample_id.get('value', '')
    
    # 描述信息
    description = report.get('description', {})
    extracted_info['description_title'] = description.get('title', '')
    extracted_info['description_comment'] = description.get('comment', '')
    
    # 生物体信息
    organism = description.get('organism', {})
    extracted_info['organism_tax_id'] = organism.get('tax_id', '')
    extracted_info['organism_sci_name'] = organism.get('sci_name', '')
    extracted_info['organism_organism_name'] = organism.get('organism_name', '')
    extracted_info['organism_common_name'] = organism.get('common_name', '')
    extracted_info['organism_strain'] = organism.get('strain', '')
    extracted_info['organism_pangolin_classification'] = organism.get('pangolin_classification', '')
    
    # 生物体谱系
    lineage = organism.get('lineage', [])
    for i, taxon in enumerate(lineage):
        extracted_info[f'organism_lineage_{i+1}_tax_id'] = taxon.get('tax_id', '')
        extracted_info[f'organism_lineage_{i+1}_name'] = taxon.get('name', '')
    
    # 种下名称
    infraspecific_names = organism.get('infraspecific_names', {})
    extracted_info['organism_infraspecific_breed'] = infraspecific_names.get('breed', '')
    extracted_info['organism_infraspecific_cultivar'] = infraspecific_names.get('cultivar', '')
    extracted_info['organism_infraspecific_ecotype'] = infraspecific_names.get('ecotype', '')
    extracted_info['organism_infraspecific_isolate'] = infraspecific_names.get('isolate', '')
    extracted_info['organism_infraspecific_sex'] = infraspecific_names.get('sex', '')
    extracted_info['organism_infraspecific_strain'] = infraspecific_names.get('strain', '')
    
    # 所有者信息
    owner = report.get('owner', {})
    extracted_info['owner_name'] = owner.get('name', '')
    
    # 所有者联系人
    contacts = owner.get('contacts', [])
    for i, contact in enumerate(contacts):
        extracted_info[f'owner_contact_{i+1}_lab'] = contact.get('lab', '')
    
    # 模型信息
    models = report.get('models', [])
    for i, model in enumerate(models):
        extracted_info[f'model_{i+1}'] = model
    
    # 生物项目信息
    bioprojects = report.get('bioprojects', [])
    for i, bioproject in enumerate(bioprojects):
        extracted_info[f'bioproject_{i+1}_accession'] = bioproject.get('accession', '')
        extracted_info[f'bioproject_{i+1}_title'] = bioproject.get('title', '')
        extracted_info[f'bioproject_{i+1}_parent_accession'] = bioproject.get('parent_accession', '')
        
        # 父级访问号列表
        parent_accessions = bioproject.get('parent_accessions', [])
        for j, parent_accession in enumerate(parent_accessions):
            extracted_info[f'bioproject_{i+1}_parent_accession_{j+1}'] = parent_accession
    
    # 属性信息
    attributes = report.get('attributes', [])
    for attr in attributes:
        name = attr.get('name', '')
        value = attr.get('value', '')
        if name and value:
            # 将属性名称转换为有效的列名
            col_name = f"attr_{name.replace(' ', '_').replace('/', '_').lower()}"
            extracted_info[col_name] = value
    
    # 其他字段
    extracted_info['age'] = report.get('age', '')
    extracted_info['biomaterial_provider'] = report.get('biomaterial_provider', '')
    extracted_info['breed'] = report.get('breed', '')
    extracted_info['collected_by'] = report.get('collected_by', '')
    extracted_info['collection_date'] = report.get('collection_date', '')
    extracted_info['cultivar'] = report.get('cultivar', '')
    extracted_info['dev_stage'] = report.get('dev_stage', '')
    extracted_info['ecotype'] = report.get('ecotype', '')
    extracted_info['geo_loc_name'] = report.get('geo_loc_name', '')
    extracted_info['host'] = report.get('host', '')
    extracted_info['host_disease'] = report.get('host_disease', '')
    extracted_info['identified_by'] = report.get('identified_by', '')
    extracted_info['ifsac_category'] = report.get('ifsac_category', '')
    extracted_info['isolate'] = report.get('isolate', '')
    extracted_info['isolate_name_alias'] = report.get('isolate_name_alias', '')
    extracted_info['isolation_source'] = report.get('isolation_source', '')
    extracted_info['lat_lon'] = report.get('lat_lon', '')
    extracted_info['project_name'] = report.get('project_name', '')
    extracted_info['sample_name'] = report.get('sample_name', '')
    extracted_info['serovar'] = report.get('serovar', '')
    extracted_info['sex'] = report.get('sex', '')
    extracted_info['source_type'] = report.get('source_type', '')
    extracted_info['strain'] = report.get('strain', '')
    extracted_info['sub_species'] = report.get('sub_species', '')
    extracted_info['tissue'] = report.get('tissue', '')
    extracted_info['serotype'] = report.get('serotype', '')
    
    # 添加原始JSON数据（用于调试和进一步分析）
    extracted_info['raw_json'] = json.dumps(report, ensure_ascii=False)
    
    return extracted_info

async def process_batch(biosample_batch, session, batch_num, input_file):
    """处理一批biosample ID"""
    df_samples = pd.DataFrame()  # 存储样本信息
    
    # 使用批量API获取数据
    success = await get_sample_info_batch(biosample_batch, session, df_samples)
    
    # 保存本批次结果
    if not df_samples.empty:
        output_file = f"{os.path.splitext(input_file)[0]}_output_batch{batch_num}.xlsx"
        df_samples.to_excel(output_file, index_label='accession')
        print(f"批次 {batch_num} 结果已保存到 {output_file}")
        return df_samples
    else:
        print(f"批次 {batch_num} 没有获取到任何数据")
        return pd.DataFrame()

async def main(biosample_list, input_file, batch_size=10):
    """主函数，处理所有biosample ID"""
    all_results = pd.DataFrame()
    
    if not biosample_list:
        print("没有需要处理的biosample ID")
        return all_results
    
    # 分批处理
    batches = [biosample_list[i:i+batch_size] for i in range(0, len(biosample_list), batch_size)]
    print(f"将 {len(biosample_list)} 个biosample ID分成 {len(batches)} 批处理")
    
    # 设置连接参数
    connector = aiohttp.TCPConnector(limit=3, ssl=False)
    timeout = aiohttp.ClientTimeout(total=300)
    
    async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
        for batch_num, batch in enumerate(batches, 1):
            print(f"正在处理第 {batch_num} 批，共 {len(batch)} 个biosample ID")
            batch_result = await process_batch(batch, session, batch_num, input_file)
            
            if not batch_result.empty:
                all_results = pd.concat([all_results, batch_result])
            
            # 批次间延迟，避免对服务器造成过大压力
            await asyncio.sleep(1)
    
    # 保存所有结果
    if not all_results.empty:
        final_output = f"{os.path.splitext(input_file)[0]}_final_output.xlsx"
        all_results.to_excel(final_output, index_label='accession')
        print(f"所有结果已保存到 {final_output}")
        
        # 同时保存为CSV格式（更适合大数据集）
        csv_output = f"{os.path.splitext(input_file)[0]}_final_output.csv"
        all_results.to_csv(csv_output, index_label='accession')
        print(f"所有结果已保存到 {csv_output}")
    else:
        print("没有获取到任何数据")
    
    return all_results

# 测试函数，用于验证API请求
def test_ncbi_api():
    """测试NCBI Datasets API请求"""
    import requests
    
    # 测试sample_accession
    print("测试NCBI Datasets API...")
    url = "https://api.ncbi.nlm.nih.gov/datasets/v2/biosample/accession/SAMN39983294/biosample_report"
    headers = {
        "Accept": "application/json"
    }
    
    try:
        response = requests.get(url, headers=headers, timeout=30)
        if response.status_code == 200:
            data = response.json()
            if "reports" in data and data["reports"]:
                report = data["reports"][0]
                print(f"成功获取sample_accession数据:")
                print(f"标题: {report.get('description', {}).get('title', '')}")
                print(f"生物体: {report.get('description', {}).get('organism', {}).get('sci_name', '')}")
                print(f"分类ID: {report.get('description', {}).get('organism', {}).get('tax_id', '')}")
                print(f"采集日期: {report.get('collection_date', '')}")
                print(f"宿主: {report.get('host', '')}")
                print(f"分离源: {report.get('isolation_source', '')}")
                
                # 提取属性
                attributes = report.get('attributes', [])
                for attr in attributes:
                    name = attr.get('name', '')
                    value = attr.get('value', '')
                    if name and value:
                        print(f"{name}: {value}")
            else:
                print("警告: sample_accession返回空数据")
                if "messages" in data:
                    for msg in data["messages"]:
                        if "error" in msg:
                            print(f"错误信息: {msg['error'].get('message', '未知错误')}")
        else:
            print(f"错误: 获取sample_accession时服务器返回状态码 {response.status_code}")
    except Exception as e:
        print(f"错误: 获取sample_accession时发生异常: {str(e)}")

if __name__ == '__main__':
    # 首先测试API连接
    print("测试NCBI Datasets API连接...")
    test_ncbi_api()
    
    print('\n开始下载')
    
    # 读取biosample ID列表
    input_file = 'biosample_list.txt'  # 请替换为您的文件名
    biosample_list = read_file(input_file)
    
    if not biosample_list:
        print("无法读取biosample ID列表，程序退出")
        exit(1)
    
    # 运行主函数
    try:
        result_df = asyncio.run(main(biosample_list, input_file))
        print("下载结束")
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
        import traceback
        traceback.print_exc()