import re
import os
from typing import TypeVar
from functools import cache
import logging

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)
from peft import PeftModel
from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
from langchain_community.chat_models import ChatLiteLLM
from langchain.chains import LLMChain
from langchain.output_parsers import PydanticOutputParser
from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    PromptTemplate,
)
from langchain.schema import BaseOutputParser, OutputParserException
import spaces

from message_classes import ActionType, AgentAction
from utils import format_docstring
from langchain_callback_handler import LoggingCallbackHandler

HF_TOKEN_KEY_FILE="./hf_token.key"
if os.path.exists(HF_TOKEN_KEY_FILE):
    with open(HF_TOKEN_KEY_FILE, "r") as f:
        os.environ["HF_TOKEN"] = f.read().strip()

OutputType = TypeVar("OutputType", bound=object)
log = logging.getLogger("generate")
logging_handler = LoggingCallbackHandler("langchain")

def generate_action(
    model_name: str,
    history: str,
    turn_number: int,
    action_types: list[ActionType],
    agent: str,
    temperature: float = 0.7,
) -> AgentAction:
    """
    Using langchain to generate an example episode
    """
    # try:
        # Normal case, model as agent
    template = """
        Imagine you are {agent}, your task is to act/speak as {agent} would, keeping in mind {agent}'s social goal.
        You can find {agent}'s goal (or background) in the 'Here is the context of the interaction' field.
        Note that {agent}'s goal is only visible to you.
        You should try your best to achieve {agent}'s goal in a way that align with their character traits.
        Additionally, maintaining the conversation's naturalness and realism is essential (e.g., do not repeat what other people has already said before).\n
        {history}.
        You are at Turn #{turn_number}. Your available action types are
        {action_list}.
        Note: You can "leave" this conversation if 1. you have achieved your social goals, 2. this conversation makes you uncomfortable, 3. you find it uninteresting/you lose your patience, 4. or for other reasons you want to leave.

        Please only generate a JSON string including the action type and the argument.
        Your action should follow the given format:
        {format_instructions}
    """
    return generate(
        model_name=model_name,
        template=template,
        input_values=dict(
            agent=agent,
            turn_number=str(turn_number),
            history=history,
            action_list=" ".join(action_types),
        ),
        output_parser=PydanticOutputParser(pydantic_object=AgentAction),
        temperature=temperature,
    )
    # except Exception as e:
    #     print(e)
    #     return AgentAction(action_type="none", argument="")

@cache
def prepare_model(model_name):
    compute_type = torch.float16
    
    if model_name == 'cmu-lti/sotopia-pi-mistral-7b-BC_SR':
        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", model_max_length=4096)
        model = AutoModelForCausalLM.from_pretrained(
        "mistralai/Mistral-7B-Instruct-v0.1",
        cache_dir="./.cache",
        device_map='cuda'
        )
        model = PeftModel.from_pretrained(model, model_name).to("cuda")
        
    elif model_name == 'cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit':
        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", model_max_length=4096)
        model = AutoModelForCausalLM.from_pretrained(
        "mistralai/Mistral-7B-Instruct-v0.1",
        cache_dir="./.cache",
        device_map='cuda',
        quantization_config=BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=compute_type,
            )
        )
        model = PeftModel.from_pretrained(model, model_name[0:-5]).to("cuda")
    
    elif model_name == 'mistralai/Mistral-7B-Instruct-v0.1':
        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", model_max_length=4096)
        tokenizer.model_max_length = 4096
        model = AutoModelForCausalLM.from_pretrained(
        "mistralai/Mistral-7B-Instruct-v0.1",
        cache_dir="./.cache",
        # device_map='cuda'
        )
        
    else:
         raise RuntimeError(f"Model {model_name} not supported")
     
    return model, tokenizer

def obtain_chain_hf(
    model_name: str,
    template: str,
    input_variables: list[str],
    temperature: float = 0.7,
    max_retries: int = 6,
    max_tokens: int = 2700
) -> LLMChain:
    human_message_prompt = HumanMessagePromptTemplate(
        prompt=PromptTemplate(template="[INST] " + template + " [/INST]", input_variables=input_variables)
    )
    chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])
    model, tokenizer = prepare_model(model_name)
    pipe = pipeline("text-generation", 
                    model=model, 
                    tokenizer=tokenizer, 
                    max_new_tokens=100, 
                    temperature=temperature, 
                    return_full_text=False, 
                    do_sample=True,
                    num_beams=3,
                    )
    hf = HuggingFacePipeline(pipeline=pipe)
    chain = LLMChain(llm=hf, prompt=chat_prompt_template)
    return chain

@spaces.GPU
def generate(
    model_name: str,
    template: str,
    input_values: dict[str, str],
    output_parser: BaseOutputParser[OutputType],
    temperature: float = 0.7,
) -> OutputType:
    input_variables = re.findall(r"{(.*?)}", template)
    assert (
        set(input_variables) == set(list(input_values.keys()) + ["format_instructions"])
        or set(input_variables) == set(list(input_values.keys()))
    ), f"The variables in the template must match input_values except for format_instructions. Got {sorted(input_values.keys())}, expect {sorted(input_variables)}"
    # process template
    template = format_docstring(template)
    chain = obtain_chain(model_name, template, input_variables, temperature)
    if "format_instructions" not in input_values:
        input_values["format_instructions"] = output_parser.get_format_instructions()
    result = chain.predict([logging_handler], **input_values)
    prompt = logging_handler.retrive_prompt()
    print(f"Prompt:\n {prompt}")
    print(f"Result:\n {result}")
    try:
        parsed_result = output_parser.parse(result)
    except KeyboardInterrupt:
        raise KeyboardInterrupt
    except Exception as e:
        log.debug(
            f"[red] Failed to parse result: {result}\nEncounter Exception {e}\nstart to reparse",
            extra={"markup": True},
        )
        reformat_parsed_result = format_bad_output(
            result, format_instructions=output_parser.get_format_instructions()
        )
        print(f"Reformatted result:\n {reformat_parsed_result}")
        parsed_result = output_parser.parse(reformat_parsed_result)
    log.info(f"Generated result: {parsed_result}")
    return parsed_result

def format_bad_output(
    ill_formed_output: str,
    format_instructions: str,
    model_name: str = "gpt-3.5-turbo",
) -> str:
    template = """
    Given the string that can not be parsed by json parser, reformat it to a string that can be parsed by json parser.
    Original string: {ill_formed_output}

    Format instructions: {format_instructions}

    Please only generate the JSON:
    """
    chain = obtain_chain(
        model_name=model_name,
        template=template,
        input_variables=re.findall(r"{(.*?)}", template),
    )
    input_values = {
        "ill_formed_output": ill_formed_output,
        "format_instructions": format_instructions,
    }
    reformat = chain.predict([logging_handler], **input_values)
    log.info(f"Reformated output: {reformat}")
    return reformat

def obtain_chain(
    model_name: str,
    template: str,
    input_variables: list[str],
    temperature: float = 0.7,
    max_retries: int = 6,
) -> LLMChain:
    """
    Using langchain to sample profiles for participants
    """
    if model_name in ["cmu-lti/sotopia-pi-mistral-7b-BC_SR", "cmu-lti/sotopia-pi-mistral-7b-BC_SR_4bit", "mistralai/Mistral-7B-Instruct-v0.1"]:
        return obtain_chain_hf(
            model_name=model_name,
            template=template,
            input_variables=input_variables,
            temperature=temperature,
            max_retries=max_retries,
        )
    
    model_name = _return_fixed_model_version(model_name)
    chat = ChatLiteLLM(
        model=model_name,
        temperature=temperature,
        max_tokens=2700,  # tweak as needed
        max_retries=max_retries,
    )
    human_message_prompt = HumanMessagePromptTemplate(
        prompt=PromptTemplate(template=template, input_variables=input_variables)
    )
    chat_prompt_template = ChatPromptTemplate.from_messages([human_message_prompt])
    chain = LLMChain(llm=chat, prompt=chat_prompt_template)
    return chain

def _return_fixed_model_version(model_name: str) -> str:
    model_version_map = {
        "gpt-3.5-turbo": "gpt-3.5-turbo-0613",
        "gpt-4": "gpt-4-0613",
        "gpt-4-turbo": "gpt-4-1106-preview",
    }
    return model_version_map[model_name] if model_name in model_version_map else model_name