#!/usr/bin/env python

# -*- coding: UTF-8 -*-

'''

@Project ：Auto-BioinfoGPT 

@File    ：agent.py

@Author  ：Juexiao Zhou

@Contact : juexiao.zhou@gmail.com

@Date    ：2023/5/3 13:24 

'''

import os.path

import os

import torch.cuda

from src.prompt import PromptGenerator

from src.spinner import Spinner

from src.executor import CodeExecutor

from src.build_RAG_private import preload_retriever

from src.local_llm import api_preload, api_generator, api_preload_deepseek, api_generator_deepseek

import openai

from openai import OpenAI

import time

import json

class Agent:

    def __init__(self,

                 initial_data_list,

                 output_dir,

                 initial_goal_description,

                 model_engine,

                 openai_api,

                 execute = True,

                 blacklist='',

                 gui_mode=False,

                 cpu=False,

                 rag=False):

        self.initial_data_list = initial_data_list

        self.initial_goal_description = initial_goal_description

        self.tasks = []

        self.update_data_lists = [_ for _ in initial_data_list]

        self.output_dir = output_dir

        self.update_data_lists.append(f'{output_dir}: all outputs should be stored under this dir')

        self.model_engine = model_engine

        self.rag = rag

        self.local_model_engines = ['codellama-7bi',

                                    'codellama-13bi',

                                    'codellama-34bi',

                                    'llama2-7bc',

                                    'llama2-13bc',

                                    'llama2-70bc',

                                    'deepseek-6.7bi',

                                    'deepseek-7bi',

                                    'deepseek-33bi',

                                    'deepseek-67bc']

        self.gpt_model_engines = ['gpt-3.5-turbo',

                                  'gpt-4-turbo',

                                  'gpt-4',

                                  'gpt-4o',

                                  'gpt-4o-mini',

                                  'gpt-3.5-turbo-1106',

                                  'gpt-4-0613',

                                  'gpt-4-32k-0613',

                                  'gpt-4-1106-preview']

        self.ollama_engines = ['ollama_llama3.1']

        self.valid_model_engines = self.local_model_engines + self.gpt_model_engines

        self.openai_api = openai_api

        if self.rag:

            if self.model_engine in self.gpt_model_engines:

                self.retriever = preload_retriever(False,

                                                   self.openai_api,

                                                   PERSIST_DIR = "./softwares_database_RAG_openai",

                                                   SOURCE_DIR = "./softwares_database"

                                                   )

            else:

                self.retriever = preload_retriever(True,

                                                   None,

                                                   PERSIST_DIR="./softwares_database_RAG_local",

                                                   SOURCE_DIR="./softwares_database"

                                                   )

        else:

            self.retriever = None

        self.generator = PromptGenerator(blacklist=blacklist,

                                         engine = self.model_engine,

                                         rag = self.rag,

                                         retriever = self.retriever)

        self.global_round = 0

        self.execute = execute

        self.execute_success = True

        self.execute_info = ''

        self.code_executor = CodeExecutor()

        self.gui_mode = gui_mode

        self.cpu = cpu

        if self.model_engine.startswith('ollama_'):

            print('[INFO] using ollama engine!')

        elif self.model_engine not in self.valid_model_engines:

            print('[ERROR] model invalid, please check the model engine selected!')

            exit()

        # use gpt model

        if self.model_engine in self.gpt_model_engines:

            self.openai_client = OpenAI(

                # This is the default and can be omitted

                api_key=self.openai_api,

            )

        # load local model with ollama

        if self.model_engine.startswith('ollama_'):

            from langchain_community.llms import Ollama

            self.local_llm_generator = Ollama(model=self.model_engine.split('ollama_')[-1])

        # preload local model

        if 'llama' in self.model_engine and 'ollama_' not in self.model_engine:

            import torch.distributed as dist

            os.environ['MASTER_ADDR'] = 'localhost'

            os.environ['MASTER_PORT'] = '5678'

            if torch.cuda.is_available():

                dist.init_process_group(backend='nccl', init_method='env://', rank=0, world_size=1)

            else:

                dist.init_process_group(backend='gloo', init_method='env://', rank=0, world_size=1)

            if self.model_engine == 'codellama-7bi':

                self.local_llm_generator = api_preload(ckpt_dir='src/codellama-main/CodeLlama-7b-Instruct/',

                                        tokenizer_path='src/codellama-main/CodeLlama-7b-Instruct/tokenizer.model',

                                        max_seq_len=4096)

            elif self.model_engine == 'codellama-13bi':

                self.local_llm_generator = api_preload(ckpt_dir='src/codellama-main/CodeLlama-13b-Instruct/one-gpu/',

                                        tokenizer_path='src/codellama-main/CodeLlama-13b-Instruct/tokenizer.model',

                                        max_seq_len=4096)

            elif self.model_engine == 'codellama-34bi':

                self.local_llm_generator = api_preload(ckpt_dir='src/codellama-main/CodeLlama-34b-Instruct/one-gpu/',

                                        tokenizer_path='src/codellama-main/CodeLlama-34b-Instruct/tokenizer.model',

                                        max_seq_len=4096)

            elif self.model_engine == 'llama2-7bc':

                self.local_llm_generator = api_preload(ckpt_dir='src/llama-main/llama-2-7b-chat/',

                                        tokenizer_path='src/llama-main/tokenizer.model',

                                        max_seq_len=4096)

            elif self.model_engine == 'llama2-13bc':

                self.local_llm_generator = api_preload(ckpt_dir='src/llama-main/llama-2-13b-chat/one-gpu/',

                                        tokenizer_path='src/llama-main/tokenizer.model',

                                        max_seq_len=4096)

            elif self.model_engine == 'llama2-70bc':

                self.local_llm_generator = api_preload(ckpt_dir='src/llama-main/llama-2-70b-chat/',

                                        tokenizer_path='src/llama-main/tokenizer.model',

                                        max_seq_len=4096)

            elif self.model_engine == 'deepseek-6.7bi':

                self.tokenizer, self.local_llm_generator = api_preload_deepseek(

                    ckpt_dir='src/deepseek/deepseek-coder-6.7b-instruct/',

                    tokenizer_path='src/deepseek/deepseek-coder-6.7b-instruct/',

                    cpu = self.cpu

                )

            elif self.model_engine == 'deepseek-7bi':

                self.tokenizer, self.local_llm_generator = api_preload_deepseek(

                    ckpt_dir='src/deepseek/deepseek-coder-7b-instruct-v1.5/',

                    tokenizer_path='src/deepseek/deepseek-coder-7b-instruct-v1.5/',

                    cpu = self.cpu

                )

            elif self.model_engine == 'deepseek-33bi':

                self.tokenizer, self.local_llm_generator = api_preload_deepseek(

                    ckpt_dir='src/deepseek/deepseek-coder-33b-instruct/',

                    tokenizer_path='src/deepseek/deepseek-coder-33b-instruct/',

                    cpu = self.cpu

                )

            elif self.model_engine == 'deepseek-67bc':

                self.tokenizer, self.local_llm_generator = api_preload_deepseek(

                    ckpt_dir='src/deepseek/deepseek-llm-67b-chat/',

                    tokenizer_path='src/deepseek/deepseek-llm-67b-chat/',

                    cpu = self.cpu

                )

    def get_single_response(self, prompt):

        # use openai

        if self.model_engine in self.gpt_model_engines:

            if self.model_engine in ['gpt-3.5-turbo-1106', 'gpt-4-1106-preview']:

                response = self.openai_client.chat.completions.create(

                    model=self.model_engine,

                    response_format={"type": "json_object"},

                    messages=[

                        {"role": "user", "content": str(prompt)}],

                    max_tokens=1024,

                    temperature=0,

                )

            else:

                response = self.openai_client.chat.completions.create(

                    model=self.model_engine,

                      messages=[

                        {"role": "user", "content": str(prompt)}],

                    max_tokens=1024,

                    temperature=0,

                )

            """

            {

              "choices": [

                {

                  "finish_reason": "stop",

                  "index": 0,

                  "message": {

                    "content": "Hello! As an AI language model, I don't have emotions, but I'm functioning well. I'm here to assist you with any questions or tasks you may have. How can I help you today?",

                    "role": "assistant"

                  }

                }

              ],

              "created": 1683014436,

              "id": "chatcmpl-7BfE4AdTo5YlSIWyMDS6nL6CYv5is",

              "model": "gpt-3.5-turbo-0301",

              "object": "chat.completion",

              "usage": {

                "completion_tokens": 42,

                "prompt_tokens": 20,

                "total_tokens": 62

              }

            }

            """

            response_message = response.choices[0].message.content

        elif self.model_engine in self.local_model_engines:

            instructions = [

                [

                    {

                        "role": "user",

                        "content": str(prompt),

                    }

                ],

            ]

            if 'deepseek' in self.model_engine:

                results = api_generator_deepseek(instructions = instructions,

                                       tokenizer = self.tokenizer,

                                       generator = self.local_llm_generator,

                                       max_new_tokens=4096,

                                       top_k=50,

                                       top_p=0.95)

            elif 'llama' in self.model_engine:

                results = api_generator(instructions=instructions,

                                        generator=self.local_llm_generator,

                                        temperature=0.6)

            response_message = results[0]['generation']['content']

        elif self.model_engine.startswith('ollama_'):

            response_message = self.local_llm_generator.invoke(str(prompt))

        return response_message

    def valid_json_response(self, response_message):

        if not os.path.isdir(f'{self.output_dir}'):

            os.makedirs(f'{self.output_dir}')

        try:

            with open(f'{self.output_dir}/{self.global_round}_response.json', 'w') as w:

                json.dump(json.loads(response_message), w)

            json.load(open(f'{self.output_dir}/{self.global_round}_response.json'))

        except:

            print('[INVALID RESSPONSE]\n', response_message)

            return False

        return True

    def valid_json_response_executor(self, response_message):

        if not os.path.isdir(f'{self.output_dir}'):

            os.makedirs(f'{self.output_dir}')

        try:

            with open(f'{self.output_dir}/executor_response.json', 'w') as w:

                json.dump(json.loads(response_message), w)

            tmp_data = json.load(open(f'{self.output_dir}/executor_response.json'))

            if str(tmp_data['stat']) not in ['0', '1']:

                return False

        except:

            print('[INVALID RESSPONSE]\n', response_message)

            return False

        return True

    def process_tasks(self, response_message):

        self.tasks = response_message['plan']

    def find_json(self, response_message):

        if "```json\n" in response_message:

            start_index = response_message.find("{")

            end_index = response_message.rfind("}") + 1

            # 提取 JSON 部分

            return response_message[start_index:end_index]

        elif "```bash\n" in response_message:

            start_index = response_message.find("```bash\n")

            end_index = response_message.find("```\n")

            return str({'tool':'','code':response_message[start_index:end_index].lstrip('```bash\n')})

        else:

            start_index = response_message.find("{")

            end_index = response_message.rfind("}") + 1

            # 提取 JSON 部分

            return response_message[start_index:end_index]

    def execute_code(self, response_message):

        if not os.path.isdir(f'{self.output_dir}'):

            os.makedirs(f'{self.output_dir}')

        try:

            with open(f'{self.output_dir}/{self.global_round}.sh', 'w') as w:

                w.write(response_message['code'])

            if self.execute:

                self.last_execute_code = response_message['code']

                executor_info = self.code_executor.execute(bash_code_path=f'{self.output_dir}/{self.global_round}.sh')

                if len(executor_info) == 0:

                    execute_statu, execute_info = True, 'No error message'

                else:

                    executor_response_message = self.get_single_response(self.generator.get_executor_prompt(executor_info=executor_info))

                    print('[CHECKING EXECUTION RESULTS]\n')

                    if 'llama' in self.model_engine or 'deepseek' in self.model_engine:

                        executor_response_message = self.find_json(executor_response_message)

                    max_tries = 10

                    n_tries = 0

                    while not self.valid_json_response_executor(executor_response_message):

                        print(f'[EXECUTOR RESPONSE CHECKING TEST #{n_tries}/{max_tries}]')

                        if 'gpt' in self.model_engine:

                            time.sleep(20)

                        executor_response_message = self.get_single_response(

                            self.generator.get_executor_prompt(executor_info=executor_info))

                        print('[CHECKING EXECUTION RESULTS]\n')

                        if 'llama' in self.model_engine or 'deepseek' in self.model_engine:

                            executor_response_message = self.find_json(executor_response_message)

                        if n_tries > max_tries:

                            executor_response_message = {'stat':0, 'info':'None'}

                            break

                        n_tries += 1

                    executor_response_message = json.load(open(f'{self.output_dir}/executor_response.json'))

                    execute_statu, execute_info = executor_response_message['stat'], executor_response_message['info']

                    #execute_statu, execute_info = executor_response_message['stat'], executor_info

                #os.system(f'bash {self.output_dir}/{self.global_round}.sh')

                return bool(int(execute_statu)), execute_info

            return True, 'Success without executing'

        except Exception as e:

            return False, e

    def run_plan_phase(self):

        # initial prompt

        init_prompt = self.generator.get_prompt(

            data_list=self.initial_data_list,

            goal_description=self.initial_goal_description,

            global_round=self.global_round,

            execute_success=self.execute_success,

            execute_info=self.execute_info

        )

        INFO_STR_USER = self.generator.format_user_prompt(init_prompt, self.global_round, self.gui_mode)

        if self.gui_mode:

            print('[AI Thinking...]')

            response_message = self.get_single_response(init_prompt)

            if 'llama' in self.model_engine or 'deepseek' in self.model_engine:

                response_message = self.find_json(response_message)

            while not self.valid_json_response(response_message):

                print(f'[Invalid Response, Waiting for 20s and Retrying...]')

                print(f'invalid response: {response_message}')

                if 'gpt' in self.model_engine:

                    time.sleep(20)

                response_message = self.get_single_response(init_prompt)

            response_message = json.load(open(f'{self.output_dir}/{self.global_round}_response.json'))

        else:

            with Spinner(f'\033[32m[AI Thinking...]\033[0m'):

                response_message = self.get_single_response(init_prompt)

                if 'llama' in self.model_engine or 'deepseek' in self.model_engine:

                    response_message = self.find_json(response_message)

                while not self.valid_json_response(response_message):

                    print(f'\033[32m[Invalid Response, Waiting for 20s and Retrying...]\033[0m')

                    print(f'invalid response: {response_message}')

                    if 'gpt' in self.model_engine:

                        time.sleep(20)

                    response_message = self.get_single_response(init_prompt)

                response_message = json.load(open(f'{self.output_dir}/{self.global_round}_response.json'))

        INFO_STR_AI = self.generator.format_ai_response(response_message, self.gui_mode)

        # process tasks

        self.process_tasks(response_message)

        self.generator.set_tasks(self.tasks)

        self.generator.add_history(None, self.global_round, self.update_data_lists)

        self.global_round += 1

        if self.execute == False:

            time.sleep(15)

        else:

            pass

    def run_code_generation_phase(self):

        # finish task one-by-one with code

        # print('[DEBUG] ', self.tasks)

        while len(self.tasks) > 0:

            task = self.tasks.pop(0)

            prompt = self.generator.get_prompt(

                data_list=self.update_data_lists,

                goal_description=task,

                global_round=self.global_round,

                execute_success=self.execute_success,

                execute_info=self.execute_info

            )

            self.first_prompt = True

            self.execute_success = False

            while self.execute_success == False:

                if self.first_prompt == False:

                    prompt = self.generator.get_prompt(

                        data_list=self.update_data_lists,

                        goal_description=task,

                        global_round=self.global_round,

                        execute_success=self.execute_success,

                        execute_info=self.execute_info,

                        last_execute_code=self.last_execute_code

                    )

                INFO_STR_USER = self.generator.format_user_prompt(prompt, self.global_round, self.gui_mode)

                if self.gui_mode:

                    print('[AI Thinking...]')

                    response_message = self.get_single_response(prompt)

                    if 'llama' in self.model_engine or 'deepseek' in self.model_engine:

                        response_message = self.find_json(response_message)

                    while not self.valid_json_response(response_message):

                        print(f'[Invalid Response, Waiting for 20s and Retrying...]')

                        print(f'invalid response: {response_message}')

                        if 'gpt' in self.model_engine:

                            time.sleep(20)

                        response_message = self.get_single_response(prompt)

                    response_message = json.load(open(f'{self.output_dir}/{self.global_round}_response.json'))

                else:

                    with Spinner(f'\033[32m[AI Thinking...]\033[0m'):

                        response_message = self.get_single_response(prompt)

                        if 'llama' in self.model_engine or 'deepseek' in self.model_engine:

                            response_message = self.find_json(response_message)

                        while not self.valid_json_response(response_message):

                            print(f'\033[32m[Invalid Response, Waiting for 20s and Retrying...]\033[0m')

                            print(f'invalid response: {response_message}')

                            if 'gpt' in self.model_engine:

                                time.sleep(20)

                            response_message = self.get_single_response(prompt)

                        response_message = json.load(open(f'{self.output_dir}/{self.global_round}_response.json'))

                INFO_STR_AI = self.generator.format_ai_response(response_message, self.gui_mode)

                # execute code

                if self.gui_mode:

                    print('[AI Executing codes...]')

                    print(f'[Execute Code Start]')

                    execute_success, execute_info = self.execute_code(response_message)

                    self.execute_success = execute_success

                    self.execute_info = execute_info

                    print('[Execute Code Finish]', self.execute_success, self.execute_info)

                    if self.execute_success:

                        print(f'[Execute Code Success!]')

                        self.execute_info = ''

                    else:

                        print(f'[Execute Code Failed!]')

                        self.first_prompt = False

                else:

                    with Spinner(f'\033[32m[AI Executing codes...]\033[0m'):

                        print(f'\033[32m[Execute Code Start]\033[0m')

                        execute_success, execute_info = self.execute_code(response_message)

                        self.execute_success = execute_success

                        self.execute_info = execute_info

                        print('\033[32m[Execute Code Finish]\033[0m', self.execute_success, self.execute_info)

                        if self.execute_success:

                            print(f'\033[32m[Execute Code Success!]\033[0m')

                            self.execute_info = ''

                        else:

                            print(f'\033[31m[Execute Code Failed!]\033[0m')

                            self.first_prompt = False

            self.generator.add_history(task, self.global_round, self.update_data_lists, code=response_message['code'])

            self.global_round += 1

            if self.execute == False:

                time.sleep(15)

    def run(self):

        self.run_plan_phase()

        self.run_code_generation_phase()

        print(f'\033[31m[Job Finished! Cheers!]\033[0m')