#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：Auto-BioinfoGPT 
@File    ：app.py
@Author  ：Juexiao Zhou
@Contact : juexiao.zhou@gmail.com
@Date    ：2023/5/2 10:51 
'''

from src.agent import Agent
import yaml
import argparse
import time

def main(init_data_list,
         output_dir,
         init_goal_description,
         model_engine,
         openai_api,
         execute,
         blacklist,
         gui_mode,
         cpu,
         rag):
    AIAgent = Agent(initial_data_list=init_data_list,
                    output_dir=output_dir,
                    initial_goal_description=init_goal_description,
                    model_engine=model_engine,
                    openai_api=openai_api,
                    execute=execute,
                    blacklist=blacklist,
                    gui_mode=gui_mode,
                    cpu=cpu,
                    rag=rag)
    AIAgent.run()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description="ABC", add_help=False)
    parser.add_argument('--config',
                        help='path/to/config.yaml',
                        default='./examples/case1.1/config.yaml')
    parser.add_argument('--openai',
                        help='openai api',
                        default='SET_YOUR_OPENAI_API')
    parser.add_argument('--model',
                        help='name of model engine, e.g. gpt-4, please refer to model zoo',
                        default='gpt-4')
    parser.add_argument('--execute',
                        help='execute code or only write codes',
                        default=False,
                        type=bool)
    parser.add_argument('--blacklist',
                        help='list of softwares in blacklist, default: java,perl,annovar',
                        default='java,perl,annovar,Cutadapt,STAR',
                        type=str)
    parser.add_argument('--gui_mode',
                        default=False,
                        type=bool)
    parser.add_argument('--cpu',
                        default=False,
                        type=bool)
    parser.add_argument('--rag',
                        help='Use RAG or not',
                        default=False,
                        type=bool)
    args = parser.parse_args()

    print("""

  /$$$$$$              /$$               /$$$$$$$   /$$$$$$ 
 /$$__  $$            | $$              | $$__  $$ /$$__  $$
| $$  \ $$ /$$   /$$ /$$$$$$    /$$$$$$ | $$  \ $$| $$  \ $$
| $$$$$$$$| $$  | $$|_  $$_/   /$$__  $$| $$$$$$$ | $$$$$$$$
| $$__  $$| $$  | $$  | $$    | $$  \ $$| $$__  $$| $$__  $$
| $$  | $$| $$  | $$  | $$ /$$| $$  | $$| $$  \ $$| $$  | $$
| $$  | $$|  $$$$$$/  |  $$$$/|  $$$$$$/| $$$$$$$/| $$  | $$
|__/  |__/ \______/    \___/   \______/ |_______/ |__/  |__/
                                                            
           Automated Bioinformatics Analysis
                  www.joshuachou.ink
    """)

    with open(args.config, 'r') as file:
        configs = yaml.safe_load(file)
    print(configs)
    init_data_list = configs['data_list']
    output_dir = configs['output_dir']
    init_goal_description = configs['goal_description']
    start_time = time.time()
    main(init_data_list,
         output_dir,
         init_goal_description,
         args.model,
         args.openai,
         args.execute,
         args.blacklist,
         args.gui_mode,
         args.cpu,
         args.rag)
    end_time = time.time()
    print(f'\033[31m[Total time cost: {end_time-start_time}]\033[0m')