import pickle
from pathlib import Path

from agentdriver.llm_core.timeout import timeout
from agentdriver.reasoning.chain_of_thoughts import (
    generate_chain_of_thoughts,
)
from agentdriver.reasoning.prompt_reasoning import (
    generate_reasoning_results,
)
import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

class ReasoningAgent:
    def __init__(self, model_name="gpt-3.5-turbo", verbose=True) -> None:
        self.verbose = verbose
        self.model_name = model_name

        self.model = {}
        if "llama" in model_name:
            model_id = "/project/Meta-Llama-3-8B-Instruct"
            local_tokenizer = AutoTokenizer.from_pretrained(model_id)
            local_model = AutoModelForCausalLM.from_pretrained(
                # "meta-llama/Meta-Llama-3-8B-Instruct",
                model_id,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                load_in_8bit=True)

            # pipeline = transformers.pipeline("text-generation", model="/project/Meta-Llama-3-8B-Instruct", model_kwargs={"torch_dtype": torch.bfloat16, "load_in_8bit":True}, device_map="auto")
            
            self.model["model_name"] = "meta-llama/Meta-Llama-3-8B-Instruct"
            self.model["model"] = local_model
            self.model["tokenizer"] = local_tokenizer
            # self.model["pipeline"] = pipeline
    
    def generate_chain_of_thoughts_target(self, data_dict, working_memory):
        """Generating reasoning targets by rules, can be used as fine-tuning"""
        reasoning = generate_chain_of_thoughts(data_dict, working_memory)
        if self.verbose:
            print(reasoning)
        return reasoning

    @timeout(15)
    def generate_chain_of_thoughts_reasoning(self, env_info_prompts, system_message, model=None):
        """Generating chain_of_thoughts reasoning by GPT in-context learning"""
        reasoning = generate_reasoning_results(env_info_prompts, self.model_name, system_message, model)
        if self.verbose:
            print(reasoning)
        return reasoning
    
    @timeout(15)
    def run(self, data_dict, env_info_prompts, system_message, working_memory, use_cot_rules=False):
        """Generate planning target and chain_of_thoughts reasoning"""
        if use_cot_rules:
            reasoning = self.generate_chain_of_thoughts_target(data_dict, working_memory)
        else:
            reasoning = self.generate_chain_of_thoughts_reasoning(env_info_prompts, system_message, self.model)
        return reasoning