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

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):
    """批量获取基因组数据集报告 - 使用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/genome/biosample/{accessions_str}/dataset_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"])} 个基因组数据集报告')
                    return data
                else:
                    print(f'警告: 未找到基因组数据集报告')
                    if "messages" in data:
                        for msg in data["messages"]:
                            if "error" in msg:
                                print(f"错误信息: {msg['error'].get('message', '未知错误')}")
                    # 返回一个空字典，表示没有数据但API调用成功
                    return {"reports": []}
            else:
                print(f'错误: 获取基因组数据集报告时服务器返回状态码 {resp.status}')
                return None
    except asyncio.TimeoutError:
        print(f'错误: 获取基因组数据集报告时超时')
        return None
    except Exception as e:
        print(f'错误: 获取基因组数据集报告时发生异常: {str(e)}')
        return None

def sanitize_column_name(name):
    """清理列名，移除特殊字符"""
    # 移除特殊字符，只保留字母、数字和下划线
    return re.sub(r'[^a-zA-Z0-9_]', '_', name)

def extract_all_info(report):
    """从基因组数据集报告中提取所有信息"""
    all_info = {}
    
    # 基本信息
    all_info['accession'] = report.get('accession', '')
    all_info['current_accession'] = report.get('current_accession', '')
    all_info['paired_accession'] = report.get('paired_accession', '')
    all_info['source_database'] = report.get('source_database', '')
    
    # 生物体信息
    organism = report.get('organism', {})
    all_info['organism_tax_id'] = organism.get('tax_id', '')
    all_info['organism_sci_name'] = organism.get('sci_name', '')
    all_info['organism_organism_name'] = organism.get('organism_name', '')
    all_info['organism_common_name'] = organism.get('common_name', '')
    all_info['organism_strain'] = organism.get('strain', '')
    all_info['organism_pangolin_classification'] = organism.get('pangolin_classification', '')
    
    # 生物体谱系
    lineage = organism.get('lineage', [])
    for i, taxon in enumerate(lineage):
        all_info[f'organism_lineage_{i+1}_tax_id'] = taxon.get('tax_id', '')
        all_info[f'organism_lineage_{i+1}_name'] = taxon.get('name', '')
    
    # 种下名称
    infraspecific_names = organism.get('infraspecific_names', {})
    all_info['organism_infraspecific_breed'] = infraspecific_names.get('breed', '')
    all_info['organism_infraspecific_cultivar'] = infraspecific_names.get('cultivar', '')
    all_info['organism_infraspecific_ecotype'] = infraspecific_names.get('ecotype', '')
    all_info['organism_infraspecific_isolate'] = infraspecific_names.get('isolate', '')
    all_info['organism_infraspecific_sex'] = infraspecific_names.get('sex', '')
    all_info['organism_infraspecific_strain'] = infraspecific_names.get('strain', '')
    
    # 组装信息
    assembly_info = report.get('assembly_info', {})
    all_info['assembly_level'] = assembly_info.get('assembly_level', '')
    all_info['assembly_status'] = assembly_info.get('assembly_status', '')
    all_info['assembly_name'] = assembly_info.get('assembly_name', '')
    all_info['assembly_long_name'] = assembly_info.get('assembly_long_name', '')
    all_info['assembly_type'] = assembly_info.get('assembly_type', '')
    all_info['bioproject_accession'] = assembly_info.get('bioproject_accession', '')
    all_info['submission_date'] = assembly_info.get('submission_date', '')
    all_info['release_date'] = assembly_info.get('release_date', '')
    all_info['description'] = assembly_info.get('description', '')
    all_info['submitter'] = assembly_info.get('submitter', '')
    all_info['refseq_category'] = assembly_info.get('refseq_category', '')
    all_info['synonym'] = assembly_info.get('synonym', '')
    all_info['linked_assembly'] = assembly_info.get('linked_assembly', '')
    all_info['sequencing_tech'] = assembly_info.get('sequencing_tech', '')
    all_info['assembly_method'] = assembly_info.get('assembly_method', '')
    all_info['grouping_method'] = assembly_info.get('grouping_method', '')
    all_info['blast_url'] = assembly_info.get('blast_url', '')
    all_info['comments'] = assembly_info.get('comments', '')
    all_info['suppression_reason'] = assembly_info.get('suppression_reason', '')
    all_info['diploid_role'] = assembly_info.get('diploid_role', '')
    
    # 配对组装信息
    paired_assembly = assembly_info.get('paired_assembly', {})
    all_info['paired_assembly_accession'] = paired_assembly.get('accession', '')
    all_info['paired_assembly_status'] = paired_assembly.get('status', '')
    all_info['paired_assembly_annotation_name'] = paired_assembly.get('annotation_name', '')
    all_info['paired_assembly_only_genbank'] = paired_assembly.get('only_genbank', '')
    all_info['paired_assembly_only_refseq'] = paired_assembly.get('only_refseq', '')
    all_info['paired_assembly_changed'] = paired_assembly.get('changed', '')
    all_info['paired_assembly_manual_diff'] = paired_assembly.get('manual_diff', '')
    all_info['paired_assembly_refseq_genbank_are_different'] = paired_assembly.get('refseq_genbank_are_different', '')
    all_info['paired_assembly_differences'] = paired_assembly.get('differences', '')
    
    # 生物项目谱系
    bioproject_lineage = assembly_info.get('bioproject_lineage', [])
    for i, lineage_item in enumerate(bioproject_lineage):
        bioprojects = lineage_item.get('bioprojects', [])
        for j, bioproject in enumerate(bioprojects):
            all_info[f'bioproject_lineage_{i+1}_{j+1}_accession'] = bioproject.get('accession', '')
            all_info[f'bioproject_lineage_{i+1}_{j+1}_title'] = bioproject.get('title', '')
            all_info[f'bioproject_lineage_{i+1}_{j+1}_parent_accession'] = bioproject.get('parent_accession', '')
            
            parent_accessions = bioproject.get('parent_accessions', [])
            for k, parent_accession in enumerate(parent_accessions):
                all_info[f'bioproject_lineage_{i+1}_{j+1}_parent_accession_{k+1}'] = parent_accession
    
    # 非典型组装信息
    atypical = assembly_info.get('atypical', {})
    all_info['atypical_is_atypical'] = atypical.get('is_atypical', '')
    
    atypical_warnings = atypical.get('warnings', [])
    for i, warning in enumerate(atypical_warnings):
        all_info[f'atypical_warning_{i+1}'] = warning
    
    # 基因组注释
    genome_notes = assembly_info.get('genome_notes', [])
    for i, note in enumerate(genome_notes):
        all_info[f'genome_note_{i+1}'] = note
    
    # 链接组装信息
    linked_assemblies = assembly_info.get('linked_assemblies', [])
    for i, linked_assembly in enumerate(linked_assemblies):
        all_info[f'linked_assembly_{i+1}_linked_assembly'] = linked_assembly.get('linked_assembly', '')
        all_info[f'linked_assembly_{i+1}_assembly_type'] = linked_assembly.get('assembly_type', '')
    
    # 组装统计信息
    assembly_stats = report.get('assembly_stats', {})
    all_info['total_number_of_chromosomes'] = assembly_stats.get('total_number_of_chromosomes', '')
    all_info['total_sequence_length'] = assembly_stats.get('total_sequence_length', '')
    all_info['total_ungapped_length'] = assembly_stats.get('total_ungapped_length', '')
    all_info['number_of_contigs'] = assembly_stats.get('number_of_contigs', '')
    all_info['contig_n50'] = assembly_stats.get('contig_n50', '')
    all_info['contig_l50'] = assembly_stats.get('contig_l50', '')
    all_info['number_of_scaffolds'] = assembly_stats.get('number_of_scaffolds', '')
    all_info['scaffold_n50'] = assembly_stats.get('scaffold_n50', '')
    all_info['scaffold_l50'] = assembly_stats.get('scaffold_l50', '')
    all_info['gaps_between_scaffolds_count'] = assembly_stats.get('gaps_between_scaffolds_count', '')
    all_info['number_of_component_sequences'] = assembly_stats.get('number_of_component_sequences', '')
    all_info['atgc_count'] = assembly_stats.get('atgc_count', '')
    all_info['gc_count'] = assembly_stats.get('gc_count', '')
    all_info['gc_percent'] = assembly_stats.get('gc_percent', '')
    all_info['genome_coverage'] = assembly_stats.get('genome_coverage', '')
    all_info['number_of_organelles'] = assembly_stats.get('number_of_organelles', '')
    
    # 细胞器信息
    organelle_info = report.get('organelle_info', [])
    for i, organelle in enumerate(organelle_info):
        all_info[f'organelle_{i+1}_assembly_name'] = organelle.get('assembly_name', '')
        all_info[f'organelle_{i+1}_infraspecific_name'] = organelle.get('infraspecific_name', '')
        all_info[f'organelle_{i+1}_description'] = organelle.get('description', '')
        all_info[f'organelle_{i+1}_total_seq_length'] = organelle.get('total_seq_length', '')
        all_info[f'organelle_{i+1}_submitter'] = organelle.get('submitter', '')
        
        bioprojects = organelle.get('bioproject', [])
        for j, bioproject in enumerate(bioprojects):
            all_info[f'organelle_{i+1}_bioproject_{j+1}'] = bioproject
    
    # 附加提交者信息
    additional_submitters = report.get('additional_submitters', [])
    for i, submitter in enumerate(additional_submitters):
        all_info[f'additional_submitter_{i+1}_genbank_accession'] = submitter.get('genbank_accession', '')
        all_info[f'additional_submitter_{i+1}_refseq_accession'] = submitter.get('refseq_accession', '')
        all_info[f'additional_submitter_{i+1}_chr_name'] = submitter.get('chr_name', '')
        all_info[f'additional_submitter_{i+1}_molecule_type'] = submitter.get('molecule_type', '')
        all_info[f'additional_submitter_{i+1}_submitter'] = submitter.get('submitter', '')
        all_info[f'additional_submitter_{i+1}_bioproject_accession'] = submitter.get('bioproject_accession', '')
    
    # 注释信息
    annotation_info = report.get('annotation_info', {})
    all_info['annotation_name'] = annotation_info.get('name', '')
    all_info['annotation_provider'] = annotation_info.get('provider', '')
    all_info['annotation_release_date'] = annotation_info.get('release_date', '')
    all_info['annotation_report_url'] = annotation_info.get('report_url', '')
    all_info['annotation_method'] = annotation_info.get('method', '')
    all_info['annotation_pipeline'] = annotation_info.get('pipeline', '')
    all_info['annotation_software_version'] = annotation_info.get('software_version', '')
    all_info['annotation_status'] = annotation_info.get('status', '')
    all_info['annotation_release_version'] = annotation_info.get('release_version', '')
    
    # 注释统计信息
    annotation_stats = annotation_info.get('stats', {})
    gene_counts = annotation_stats.get('gene_counts', {})
    all_info['gene_counts_total'] = gene_counts.get('total', '')
    all_info['gene_counts_protein_coding'] = gene_counts.get('protein_coding', '')
    all_info['gene_counts_non_coding'] = gene_counts.get('non_coding', '')
    all_info['gene_counts_pseudogene'] = gene_counts.get('pseudogene', '')
    all_info['gene_counts_other'] = gene_counts.get('other', '')
    
    # BUSCO信息
    busco = annotation_info.get('busco', {})
    all_info['busco_busco_lineage'] = busco.get('busco_lineage', '')
    all_info['busco_busco_ver'] = busco.get('busco_ver', '')
    all_info['busco_complete'] = busco.get('complete', '')
    all_info['busco_single_copy'] = busco.get('single_copy', '')
    all_info['busco_duplicated'] = busco.get('duplicated', '')
    all_info['busco_fragmented'] = busco.get('fragmented', '')
    all_info['busco_missing'] = busco.get('missing', '')
    all_info['busco_total_count'] = busco.get('total_count', '')
    
    # WGS信息
    wgs_info = report.get('wgs_info', {})
    all_info['wgs_project_accession'] = wgs_info.get('wgs_project_accession', '')
    all_info['wgs_master_wgs_url'] = wgs_info.get('master_wgs_url', '')
    all_info['wgs_wgs_contigs_url'] = wgs_info.get('wgs_contigs_url', '')
    
    # 类型材料信息
    type_material = report.get('type_material', {})
    all_info['type_material_label'] = type_material.get('type_label', '')
    all_info['type_material_display_text'] = type_material.get('type_display_text', '')
    
    # CheckM信息
    checkm_info = report.get('checkm_info', {})
    all_info['checkm_marker_set'] = checkm_info.get('checkm_marker_set', '')
    all_info['checkm_species_tax_id'] = checkm_info.get('checkm_species_tax_id', '')
    all_info['checkm_marker_set_rank'] = checkm_info.get('checkm_marker_set_rank', '')
    all_info['checkm_version'] = checkm_info.get('checkm_version', '')
    all_info['checkm_completeness'] = checkm_info.get('completeness', '')
    all_info['checkm_contamination'] = checkm_info.get('contamination', '')
    all_info['checkm_completeness_percentile'] = checkm_info.get('completeness_percentile', '')
    
    # 平均核苷酸同一性信息
    average_nucleotide_identity = report.get('average_nucleotide_identity', {})
    all_info['ani_taxonomy_check_status'] = average_nucleotide_identity.get('taxonomy_check_status', '')
    all_info['ani_match_status'] = average_nucleotide_identity.get('match_status', '')
    all_info['ani_submitted_organism'] = average_nucleotide_identity.get('submitted_organism', '')
    all_info['ani_submitted_species'] = average_nucleotide_identity.get('submitted_species', '')
    all_info['ani_category'] = average_nucleotide_identity.get('category', '')
    all_info['ani_comment'] = average_nucleotide_identity.get('comment', '')
    
    # 提交的ANI匹配
    submitted_ani_match = average_nucleotide_identity.get('submitted_ani_match', {})
    all_info['submitted_ani_match_assembly'] = submitted_ani_match.get('assembly', '')
    all_info['submitted_ani_match_organism_name'] = submitted_ani_match.get('organism_name', '')
    all_info['submitted_ani_match_category'] = submitted_ani_match.get('category', '')
    all_info['submitted_ani_match_ani'] = submitted_ani_match.get('ani', '')
    all_info['submitted_ani_match_assembly_coverage'] = submitted_ani_match.get('assembly_coverage', '')
    all_info['submitted_ani_match_type_assembly_coverage'] = submitted_ani_match.get('type_assembly_coverage', '')
    
    # 最佳ANI匹配
    best_ani_match = average_nucleotide_identity.get('best_ani_match', {})
    all_info['best_ani_match_assembly'] = best_ani_match.get('assembly', '')
    all_info['best_ani_match_organism_name'] = best_ani_match.get('organism_name', '')
    all_info['best_ani_match_category'] = best_ani_match.get('category', '')
    all_info['best_ani_match_ani'] = best_ani_match.get('ani', '')
    all_info['best_ani_match_assembly_coverage'] = best_ani_match.get('assembly_coverage', '')
    all_info['best_ani_match_type_assembly_coverage'] = best_ani_match.get('type_assembly_coverage', '')
    
    # BioSample信息（如果存在）
    biosample = assembly_info.get('biosample', {})
    if biosample:
        # 提取BioSample基本信息
        all_info['biosample_accession'] = biosample.get('accession', '')
        all_info['biosample_last_updated'] = biosample.get('last_updated', '')
        all_info['biosample_publication_date'] = biosample.get('publication_date', '')
        all_info['biosample_submission_date'] = biosample.get('submission_date', '')
        all_info['biosample_package'] = biosample.get('package', '')
        all_info['biosample_status'] = biosample.get('status', {}).get('status', '')
        all_info['biosample_status_when'] = biosample.get('status', {}).get('when', '')
        
        # 样本ID信息
        sample_ids = biosample.get('sample_ids', [])
        for i, sample_id in enumerate(sample_ids):
            all_info[f'biosample_sample_id_{i+1}_db'] = sample_id.get('db', '')
            all_info[f'biosample_sample_id_{i+1}_label'] = sample_id.get('label', '')
            all_info[f'biosample_sample_id_{i+1}_value'] = sample_id.get('value', '')
        
        # 描述信息
        description = biosample.get('description', {})
        all_info['biosample_description_title'] = description.get('title', '')
        all_info['biosample_description_comment'] = description.get('comment', '')
        
        # 生物体信息
        biosample_organism = description.get('organism', {})
        all_info['biosample_organism_tax_id'] = biosample_organism.get('tax_id', '')
        all_info['biosample_organism_sci_name'] = biosample_organism.get('sci_name', '')
        all_info['biosample_organism_organism_name'] = biosample_organism.get('organism_name', '')
        all_info['biosample_organism_common_name'] = biosample_organism.get('common_name', '')
        all_info['biosample_organism_strain'] = biosample_organism.get('strain', '')
        all_info['biosample_organism_pangolin_classification'] = biosample_organism.get('pangolin_classification', '')
        
        # 生物体谱系
        biosample_lineage = biosample_organism.get('lineage', [])
        for i, taxon in enumerate(biosample_lineage):
            all_info[f'biosample_organism_lineage_{i+1}_tax_id'] = taxon.get('tax_id', '')
            all_info[f'biosample_organism_lineage_{i+1}_name'] = taxon.get('name', '')
        
        # 种下名称
        biosample_infraspecific_names = biosample_organism.get('infraspecific_names', {})
        all_info['biosample_organism_infraspecific_breed'] = biosample_infraspecific_names.get('breed', '')
        all_info['biosample_organism_infraspecific_cultivar'] = biosample_infraspecific_names.get('cultivar', '')
        all_info['biosample_organism_infraspecific_ecotype'] = biosample_infraspecific_names.get('ecotype', '')
        all_info['biosample_organism_infraspecific_isolate'] = biosample_infraspecific_names.get('isolate', '')
        all_info['biosample_organism_infraspecific_sex'] = biosample_infraspecific_names.get('sex', '')
        all_info['biosample_organism_infraspecific_strain'] = biosample_infraspecific_names.get('strain', '')
        
        # 所有者信息
        owner = biosample.get('owner', {})
        all_info['biosample_owner_name'] = owner.get('name', '')
        
        # 所有者联系人
        contacts = owner.get('contacts', [])
        for i, contact in enumerate(contacts):
            all_info[f'biosample_owner_contact_{i+1}_lab'] = contact.get('lab', '')
        
        # 模型信息
        models = biosample.get('models', [])
        for i, model in enumerate(models):
            all_info[f'biosample_model_{i+1}'] = model
        
        # 生物项目信息
        bioprojects = biosample.get('bioprojects', [])
        for i, bioproject in enumerate(bioprojects):
            all_info[f'biosample_bioproject_{i+1}_accession'] = bioproject.get('accession', '')
            all_info[f'biosample_bioproject_{i+1}_title'] = bioproject.get('title', '')
            all_info[f'biosample_bioproject_{i+1}_parent_accession'] = bioproject.get('parent_accession', '')
            
            parent_accessions = bioproject.get('parent_accessions', [])
            for j, parent_accession in enumerate(parent_accessions):
                all_info[f'biosample_bioproject_{i+1}_parent_accession_{j+1}'] = parent_accession
        
        # 属性信息
        attributes = biosample.get('attributes', [])
        for attr in attributes:
            name = attr.get('name', '')
            value = attr.get('value', '')
            if name and value:
                # 将属性名称转换为有效的列名
                col_name = f"biosample_attr_{sanitize_column_name(name)}"
                all_info[col_name] = value
        
        # 其他字段
        all_info['biosample_age'] = biosample.get('age', '')
        all_info['biosample_biomaterial_provider'] = biosample.get('biomaterial_provider', '')
        all_info['biosample_breed'] = biosample.get('breed', '')
        all_info['biosample_collected_by'] = biosample.get('collected_by', '')
        all_info['biosample_collection_date'] = biosample.get('collection_date', '')
        all_info['biosample_cultivar'] = biosample.get('cultivar', '')
        all_info['biosample_dev_stage'] = biosample.get('dev_stage', '')
        all_info['biosample_ecotype'] = biosample.get('ecotype', '')
        all_info['biosample_geo_loc_name'] = biosample.get('geo_loc_name', '')
        all_info['biosample_host'] = biosample.get('host', '')
        all_info['biosample_host_disease'] = biosample.get('host_disease', '')
        all_info['biosample_identified_by'] = biosample.get('identified_by', '')
        all_info['biosample_ifsac_category'] = biosample.get('ifsac_category', '')
        all_info['biosample_isolate'] = biosample.get('isolate', '')
        all_info['biosample_isolate_name_alias'] = biosample.get('isolate_name_alias', '')
        all_info['biosample_isolation_source'] = biosample.get('isolation_source', '')
        all_info['biosample_lat_lon'] = biosample.get('lat_lon', '')
        all_info['biosample_project_name'] = biosample.get('project_name', '')
        all_info['biosample_sample_name'] = biosample.get('sample_name', '')
        all_info['biosample_serovar'] = biosample.get('serovar', '')
        all_info['biosample_sex'] = biosample.get('sex', '')
        all_info['biosample_source_type'] = biosample.get('source_type', '')
        all_info['biosample_strain'] = biosample.get('strain', '')
        all_info['biosample_sub_species'] = biosample.get('sub_species', '')
        all_info['biosample_tissue'] = biosample.get('tissue', '')
        all_info['biosample_serotype'] = biosample.get('serotype', '')
    
    # 添加原始JSON数据（用于调试和进一步分析）
    all_info['raw_json'] = json.dumps(report, ensure_ascii=False)
    
    return all_info

async def process_batch(biosample_batch, session, batch_num, input_file, all_columns=None):
    """处理一批biosample ID"""
    # 获取API数据
    data = await get_sample_info_batch(biosample_batch, session)
    
    # 提取所有信息
    all_info_list = []
    
    if data and "reports" in data:
        for report in data["reports"]:
            info = extract_all_info(report)
            if info:
                all_info_list.append(info)
    
    # 创建DataFrame
    if all_info_list:
        df_samples = pd.DataFrame(all_info_list)
        
        # 设置索引
        if 'biosample_accession' in df_samples.columns:
            df_samples.set_index('biosample_accession', inplace=True)
        elif 'accession' in df_samples.columns:
            df_samples.set_index('accession', inplace=True)
        
        # 保存本批次结果
        output_file = f"{os.path.splitext(input_file)[0]}_output_batch{batch_num}.xlsx"
        df_samples.to_excel(output_file)
        print(f"批次 {batch_num} 结果已保存到 {output_file}")
        
        # 更新所有列名
        if all_columns is not None:
            all_columns.update(df_samples.columns.tolist())
        
        return df_samples
    else:
        print(f"批次 {batch_num} 没有获取到任何数据")
        # 创建一个空的DataFrame，但保留biosample ID
        empty_df = pd.DataFrame(index=biosample_batch)
        empty_df.index.name = 'biosample_accession'
        
        # 保存空结果
        output_file = f"{os.path.splitext(input_file)[0]}_output_batch{batch_num}.xlsx"
        empty_df.to_excel(output_file)
        print(f"批次 {batch_num} 空结果已保存到 {output_file}")
        
        return empty_df

async def main(biosample_list, input_file, batch_size=5):
    """主函数，处理所有biosample ID"""
    all_results = pd.DataFrame()
    all_columns = set()  # 用于收集所有列名
    
    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, all_columns)
            
            if not batch_result.empty:
                if all_results.empty:
                    all_results = batch_result
                else:
                    # 确保列名一致
                    for col in all_columns:
                        if col not in all_results.columns:
                            all_results[col] = None
                        if col not in batch_result.columns:
                            batch_result[col] = None
                    
                    all_results = pd.concat([all_results, batch_result])
            
            # 批次间延迟，避免对服务器造成过大压力
            await asyncio.sleep(2)
    
    # 确保所有biosample ID都在结果中
    missing_biosamples = set(biosample_list) - set(all_results.index)
    if missing_biosamples:
        print(f"添加 {len(missing_biosamples)} 个缺失的biosample ID到结果中")
        missing_df = pd.DataFrame(index=list(missing_biosamples))
        missing_df.index.name = 'biosample_accession'
        
        # 添加所有列
        for col in all_columns:
            missing_df[col] = None
        
        all_results = pd.concat([all_results, missing_df])
    
    # 保存所有结果
    if not all_results.empty:
        final_output = f"{os.path.splitext(input_file)[0]}_final_output.xlsx"
        all_results.to_excel(final_output)
        print(f"所有结果已保存到 {final_output}")
        
        # 同时保存为CSV格式（更适合大数据集）
        csv_output = f"{os.path.splitext(input_file)[0]}_final_output.csv"
        all_results.to_csv(csv_output)
        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/genome/biosample/SAMN15960293/dataset_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]
                info = extract_all_info(report)
                
                if info:
                    print(f"成功获取数据:")
                    print(f"访问号: {info.get('accession', '')}")
                    print(f"生物体: {info.get('organism_sci_name', '')}")
                    print(f"分类ID: {info.get('organism_tax_id', '')}")
                    print(f"组装级别: {info.get('assembly_level', '')}")
                    
                    # 检查是否有BioSample信息
                    if 'biosample_accession' in info:
                        print(f"BioSample访问号: {info.get('biosample_accession', '')}")
                        print(f"采集日期: {info.get('biosample_collection_date', '')}")
                        print(f"宿主: {info.get('biosample_host', '')}")
                        print(f"分离源: {info.get('biosample_isolation_source', '')}")
                else:
                    print("警告: 未提取到信息")
            else:
                print("警告: 基因组数据集报告返回空数据")
                if "messages" in data:
                    for msg in data["messages"]:
                        if "error" in msg:
                            print(f"错误信息: {msg['error'].get('message', '未知错误')}")
        else:
            print(f"错误: 获取基因组数据集报告时服务器返回状态码 {response.status_code}")
    except Exception as e:
        print(f"错误: 获取基因组数据集报告时发生异常: {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()