import aiohttp
import asyncio
import pandas as pd
import time
import os

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

async def get_biosample_info(accession_number, session, df):
    """获取单个run_accession的详细信息"""
    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",
    }
    
    # 构建API请求参数 - 移除了无效字段
    params = {
        "accession": accession_number,
        "result": "read_run",
        "format": "json",
        "fields": "study_accession,secondary_study_accession,sample_accession,secondary_sample_accession,experiment_accession,run_accession,submission_accession,tax_id,scientific_name,instrument_platform,instrument_model,library_name,nominal_length,library_layout,library_strategy,library_source,library_selection,read_count,base_count,center_name,first_public,last_updated,experiment_title,study_title,study_alias,experiment_alias,run_alias,fastq_bytes,fastq_md5,fastq_ftp,fastq_aspera,fastq_galaxy,submitted_bytes,submitted_md5,submitted_ftp,submitted_aspera,submitted_galaxy,submitted_format,sra_bytes,sra_md5,sra_ftp,sra_aspera,sra_galaxy,sample_alias,broker_name,sample_title,nominal_sdev,first_created"
    }
    
    url = 'https://www.ebi.ac.uk/ena/portal/api/filereport'
    
    try:
        async with session.get(url, headers=headers, params=params, timeout=aiohttp.ClientTimeout(total=30)) as resp:
            if resp.status == 200:
                content = await resp.text()
                if content.strip():  # 检查响应内容是否为空
                    biosample_info_dic = await resp.json()
                    
                    # 检查返回的数据是否为空
                    if biosample_info_dic and len(biosample_info_dic) > 0:
                        biosample_info_dic = biosample_info_dic[0]
                        print(f'成功获取 {accession_number} 的信息')
                        
                        # 将数据添加到DataFrame
                        for k, v in biosample_info_dic.items():
                            df.loc[accession_number, k] = v
                        return True
                    else:
                        print(f'警告: {accession_number} 返回空数据')
                        return False
                else:
                    print(f'警告: {accession_number} 返回空响应')
                    return False
            else:
                print(f'错误: 获取 {accession_number} 时服务器返回状态码 {resp.status}')
                return False
    except asyncio.TimeoutError:
        print(f'错误: 获取 {accession_number} 时超时')
        return False
    except Exception as e:
        print(f'错误: 获取 {accession_number} 时发生异常: {str(e)}')
        return False

async def process_batch(accession_batch, session, batch_num, input_file):
    """处理一批accession"""
    df = pd.DataFrame()
    tasks = []
    
    for accession in accession_batch:
        task = asyncio.create_task(get_biosample_info(accession, session, df))
        tasks.append(task)
    
    # 等待所有任务完成
    results = await asyncio.gather(*tasks, return_exceptions=True)
    
    # 保存本批次结果
    if not df.empty:
        output_file = f"{os.path.splitext(input_file)[0]}_output_batch{batch_num}.xlsx"
        df.to_excel(output_file)
        print(f"批次 {batch_num} 结果已保存到 {output_file}")
        return df
    else:
        print(f"批次 {batch_num} 没有获取到任何数据")
        return pd.DataFrame()

async def main(accession_list, input_file, batch_size=10, max_retries=3):
    """主函数，处理所有accession"""
    all_results = pd.DataFrame()
    
    if not accession_list:
        print("没有需要处理的accession")
        return all_results
    
    # 分批处理
    batches = [accession_list[i:i+batch_size] for i in range(0, len(accession_list), batch_size)]
    print(f"将 {len(accession_list)} 个accession分成 {len(batches)} 批处理")
    
    # 设置连接参数
    connector = aiohttp.TCPConnector(limit=5, 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)} 个accession")
            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(2)  # 增加延迟时间
    
    # 保存所有结果
    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}")
    else:
        print("没有获取到任何数据")
    
    return all_results

# 测试函数，用于验证API请求
def test_ena_api(accession):
    """测试单个accession的API请求"""
    import requests
    
    # 构建API请求参数 - 移除了无效字段
    params = {
        "accession": accession,
        "result": "read_run",
        "format": "json",
        "fields": "study_accession,secondary_study_accession,sample_accession,secondary_sample_accession,experiment_accession,run_accession,submission_accession,tax_id,scientific_name,instrument_platform,instrument_model,library_name,nominal_length,library_layout,library_strategy,library_source,library_selection,read_count,base_count,center_name,first_public,last_updated,experiment_title,study_title,study_alias,experiment_alias,run_alias,fastq_bytes,fastq_md5,fastq_ftp,fastq_aspera,fastq_galaxy,submitted_bytes,submitted_md5,submitted_ftp,submitted_aspera,submitted_galaxy,submitted_format,sra_bytes,sra_md5,sra_ftp,sra_aspera,sra_galaxy,sample_alias,broker_name,sample_title,nominal_sdev,first_created"
    }
    
    url = 'https://www.ebi.ac.uk/ena/portal/api/filereport'
    
    try:
        response = requests.get(url, params=params, timeout=30)
        if response.status_code == 200:
            data = response.json()
            if data:
                print(f"成功获取 {accession} 的数据:")
                print(f"返回了 {len(data)} 条记录")
                print(f"第一条记录的键: {list(data[0].keys())}")
                return True
            else:
                print(f"警告: {accession} 返回空数据")
                return False
        else:
            print(f"错误: 获取 {accession} 时服务器返回状态码 {response.status_code}")
            return False
    except Exception as e:
        print(f"错误: 获取 {accession} 时发生异常: {str(e)}")
        return False

if __name__ == '__main__':
    # 首先测试一个accession，确保API正常工作
    print("测试API连接...")
    test_result = test_ena_api("SRR12667785")
    
    if not test_result:
        print("API测试失败，请检查网络连接或参数设置")
        exit(1)
    
    print('开始下载')
    
    # 读取accession列表
    input_file = 'results_read_run_tsv.txt'
    accession_list = read_file(input_file)
    
    if not accession_list:
        print("无法读取accession列表，程序退出")
        exit(1)
    
    # 运行主函数
    try:
        # 使用asyncio.run()而不是手动管理事件循环
        result_df = asyncio.run(main(accession_list, input_file))
        print("下载结束")
    except Exception as e:
        print(f"程序执行出错: {str(e)}")
        import traceback
        traceback.print_exc()