import ast
import copy
import json
import logging
from typing import List, Tuple, Dict, Callable

import langchain
from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, AIMessagePromptTemplate
from langchain.prompts.chat import BaseMessagePromptTemplate
from langchain.schema import LLMResult
from langchain.schema.language_model import BaseLanguageModel

#from src.generators import LMGenerator
from blangchain.async_openai import JitterWaitChatOpenAI, JitterWaitOpenAI
from blangchain.tracking_utils import TokensTracker

logger = logging.getLogger(__name__)
from langchain import LLMChain, PromptTemplate, FewShotPromptTemplate
import asyncio

completion_model_map = {
    'gpt3': 'text-davinci-003',
    'gpt-3.5-turbo-instruct': 'gpt-3.5-turbo-instruct',
    'turbo-instruct': 'gpt-3.5-turbo-instruct',
}

chat_model_map = {
    'chatgpt': "gpt-3.5-turbo-0613",
    'gpt-3.5-turbo-16k': "gpt-3.5-turbo-16k",
    'chatgpt-16k': "gpt-3.5-turbo-16k",
    'gpt-4': 'gpt-4',
}


class LMGenerator:

    def generate(self, inputs: List[dict], **gen_kwargs) -> List[List[str]]:
        raise NotImplementedError()

class OpenAIGenerator(LMGenerator):
    def __init__(self, prompt=None, model='gpt3'):
        """

        :param prompt:
        :param model: either "gpt3" or "Chatgpt"
        """
        self.tracker = TokensTracker
        self.model_type = model
        self.lm_class: BaseLanguageModel = None
        if model in completion_model_map:
            self.gen_kwargs = {
                "n": 1,
                'temperature': 1,
                'model_name': completion_model_map.get(model),
                # "top_p": 1,
                "max_tokens": 1000,
                "max_retries": 100,
            }
            self.lm_class = JitterWaitOpenAI

        elif model in chat_model_map:
            self.gen_kwargs = {
                "n": 1,
                'model_name': chat_model_map.get(model),
                'temperature': 1,
                # "top_p": 1,
                "request_timeout": 600,
                "max_retries": 100,
            }
            # self.lm_class = CachedChatOpenAI
            self.lm_class = JitterWaitChatOpenAI
        else:
            raise NotImplementedError()
        self.batch_size = 50
        self.prompt = prompt
        self.total_tokens = 0

    def generate(self, inputs: List[dict], parallel=False, **gen_kwargs) -> List[List[str]]:
        _gkwargs = copy.deepcopy(self.gen_kwargs)
        _gkwargs.update(**gen_kwargs)
        if self.model_type == 'gpt3' and _gkwargs.get('n', 1) > 1:
            _gkwargs['best_of'] = _gkwargs['n']

        assert langchain.llm_cache is not None
        lm = self.lm_class(**_gkwargs)
        chain = LLMChain(llm=lm, prompt=self.prompt)
        ret = []
        for i in range(0, len(inputs), self.batch_size):
            in_batch = inputs[i:i + self.batch_size]
            if parallel:
                async def gen():
                    tasks = [chain.agenerate([ib]) for ib in in_batch]
                    ret_list = await asyncio.gather(*tasks)
                    for lm_out_i in ret_list:
                        logger.info(lm_out_i.llm_output)
                        TokensTracker.update(lm_out_i.llm_output, module=type(self).__name__)
                    return LLMResult(generations=[lm_out_i.generations[0] for lm_out_i in ret_list], )
                lm_output = asyncio.run(gen())
            else:
                lm_output = chain.generate(in_batch)
                logger.info(lm_output.llm_output)
                TokensTracker.update(lm_output.llm_output)
            ret.extend([[g.text for g in gen] for gen in lm_output.generations])
        return ret

    async def agenerate(self, inputs: List[dict], **gen_kwargs) -> List[List[str]]:
        _gkwargs = copy.deepcopy(self.gen_kwargs)
        _gkwargs.update(**gen_kwargs)
        if self.model_type == 'gpt3' and _gkwargs.get('n', 1) > 1:
            _gkwargs['best_of'] = _gkwargs['n']

        assert langchain.llm_cache is not None
        lm = self.lm_class(**_gkwargs)
        chain = LLMChain(llm=lm, prompt=self.prompt)
        tasks = [chain.agenerate([ib]) for ib in inputs]
        ret_list = await asyncio.gather(*tasks)
        for lm_out_i in ret_list:
            logger.info(lm_out_i.llm_output)
            TokensTracker.update(lm_out_i.llm_output, module=type(self).__name__)
            self.total_tokens += lm_out_i.llm_output.get('token_usage', {}).get('total_tokens', 0)
        lm_output = LLMResult(generations=[lm_out_i.generations[0] for lm_out_i in ret_list])

        ret = [[g.text for g in gen] for gen in lm_output.generations]
        # if self.model_type in ['gpt-3.5-turbo-0613', 'chatgpt']:
        #     breakpoint()
        return ret

    def format_print(self, input: Dict, _print: Callable = print):
        _print(self.prompt.format(**input))

    def format_print_to(self, input: Dict, file=None):
        with open(file, 'a+') as f:
            self.format_print(input, _print=lambda x: f.write(str(x) + '\n'))


class SimplePromptOpenAIGenerator(OpenAIGenerator):
    def __init__(self, prompt_template: PromptTemplate, model='chatgpt', debug_openai=False):
        self.debug_openai = debug_openai
        if model in completion_model_map:
            prompt = prompt_template
        elif model in chat_model_map:
            prompt = ChatPromptTemplate.from_messages([
                HumanMessagePromptTemplate(prompt=prompt_template)
            ])
        else:
            raise NotImplementedError
        super().__init__(prompt=prompt, model=model)


class JSONItemGenerator:

    async def postprocess_generation(self, gen: str, expected_items: int = None) -> List[dict]:
        """
        Takes a (potentially multi-line) string and turns it into a list of dicts
        """
        results = []

        for line in gen.split('\n'):
            if not line.strip(): continue
            line = line.strip(', ')
            line = line.strip(".")
            try:
                results.append(ast.literal_eval(line.replace('null', "None")))
            except:
                try:
                    results.append(json.loads(line))
                except:
                    try:
                        fixer = JSONFixer()
                        fixed_json: dict = (await fixer.afix(line))
                        results.append(fixed_json)
                    except:
                        continue

        if expected_items and len(results) != expected_items:
            if len(results) > expected_items:
                results = results[:expected_items]
            else:
                res = [{} for _ in range(expected_items)]
                for r in results:
                    res[r['I'] - 1] = r
                if any(res):
                    results = res
                else:  # final resort
                    results = results + [{} for _ in range(expected_items - len(results))]
        return results


class JSONOpenAIGenerator(SimplePromptOpenAIGenerator, JSONItemGenerator):
    def __init__(self, *args, **kwargs):
        super(JSONOpenAIGenerator, self).__init__(*args, **kwargs)

    def batchify(self, items_to_batch, max_size=None):
        if len(items_to_batch) <= 25:
            _statement_batch_size = len(items_to_batch)
        elif len(items_to_batch) > 25 and len(items_to_batch) <= 50:
            _statement_batch_size = int(len(items_to_batch) / 2) + 1
        elif len(items_to_batch) > 50:
            # _statement_batch_size = min(30, int(len(statements_to_score) / 4) + 1)
            _statement_batch_size = 25
        else:
            raise NotImplementedError()
        if max_size is not None:
            if len(items_to_batch) % max_size == 1:
                _statement_batch_size = max_size - 1
            else:
                _statement_batch_size = max_size

        statement_batches = [items_to_batch[i:i + _statement_batch_size]
                             for i in range(0, len(items_to_batch), _statement_batch_size)]

        return statement_batches

    async def run(self, inputs: List[dict], **kwargs) -> List[List[List[dict]]]:
        generations: List[List[str]] = await self.agenerate(inputs, **kwargs)
        result = [list(await asyncio.gather(*[self.postprocess_generation(gg) for gg in g]))
                  for g in generations]
        return result


class JSONFixer(JSONOpenAIGenerator):
    def __init__(self):
        PROMPT = """You are a system for fixing syntax errors in json items. This includes missing quotes around strings and missing closing brackets. If a key is missing its value, map it to None. Do not add new key/value pairs that are not already there.

Given the following malformed json item, return a serialized, one-line version that can be complied by json.loads() in python.
Your output should be this json item on a single line and nothing else. 

{input}
"""
        super(JSONFixer, self).__init__(prompt_template=PromptTemplate.from_template(PROMPT))

    async def afix(self, input_str) -> dict:
        '''
        takes a malformed json line and tries to fix it with gpt
        :param input_str:
        :return: json loaded item
        '''
        inputs = [dict(input=input_str)]
        ret: str = (await self.agenerate(inputs))[0][0]
        ret = ret.strip("\n").split("\n")[0]
        try:
            ret = json.loads(ret)
        except:
            ret = ast.literal_eval(ret.replace('null', "None"))

        if isinstance(ret, str):
            assert False

        return ret


message_type_to_prompt_class = {
    'human': HumanMessagePromptTemplate,
    'ai': AIMessagePromptTemplate
}


class FollowupPromptOpenAIGenerator(OpenAIGenerator):
    def __init__(self, prompt_template_list: List[Tuple[str, PromptTemplate]], model='gpt3'):

        if model in completion_model_map:
            if any(isinstance(i, FewShotPromptTemplate) for i in prompt_template_list[1:]):
                raise NotImplementedError("cannot handle template lists that have fewshot prompts after the first")
            if isinstance(prompt_template_list[0][1], FewShotPromptTemplate):
                combined_template = '\n\n'.join(template.template for (_, template) in prompt_template_list[1:])
                first_prompt: FewShotPromptTemplate = prompt_template_list[0][1]
                prompt = FewShotPromptTemplate(
                    examples=first_prompt.examples,
                    example_selector=first_prompt.example_selector,
                    example_prompt=first_prompt.example_prompt,
                    suffix=first_prompt.suffix + '\n' + combined_template,
                    input_variables=first_prompt.input_variables + PromptTemplate.from_template(
                        combined_template).input_variables,
                    example_separator=first_prompt.example_separator,
                    prefix=first_prompt.prefix
                )
            else:
                def _get_template(t):
                    if isinstance(t, BaseMessagePromptTemplate):
                        return t
                    else:
                        return t.template

                combined_template = '\n\n'.join(template.template for (_, template) in prompt_template_list)
                prompt = PromptTemplate.from_template(combined_template)
        elif model in chat_model_map:
            prompt = ChatPromptTemplate.from_messages([
                message_type_to_prompt_class[_type](prompt=template) for (_type, template) in prompt_template_list
            ])
        else:
            raise NotImplementedError
        super().__init__(prompt=prompt, model=model)


langchain.schema.LLMResult langchain.prompts.HumanMessagePromptTemplate langchain.LLMChain langchain.PromptTemplate.from_template
[((557, 584), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (574, 584), False, 'import logging\n'), ((2451, 2481), 'copy.deepcopy', 'copy.deepcopy', (['self.gen_kwargs'], {}), '(self.gen_kwargs)\n', (2464, 2481), False, 'import copy\n'), ((2738, 2774), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'lm', 'prompt': 'self.prompt'}), '(llm=lm, prompt=self.prompt)\n', (2746, 2774), False, 'from langchain import LLMChain, PromptTemplate, FewShotPromptTemplate\n'), ((3822, 3852), 'copy.deepcopy', 'copy.deepcopy', (['self.gen_kwargs'], {}), '(self.gen_kwargs)\n', (3835, 3852), False, 'import copy\n'), ((4109, 4145), 'langchain.LLMChain', 'LLMChain', ([], {'llm': 'lm', 'prompt': 'self.prompt'}), '(llm=lm, prompt=self.prompt)\n', (4117, 4145), False, 'from langchain import LLMChain, PromptTemplate, FewShotPromptTemplate\n'), ((4531, 4604), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': '[lm_out_i.generations[0] for lm_out_i in ret_list]'}), '(generations=[lm_out_i.generations[0] for lm_out_i in ret_list])\n', (4540, 4604), False, 'from langchain.schema import LLMResult\n'), ((4228, 4250), 'asyncio.gather', 'asyncio.gather', (['*tasks'], {}), '(*tasks)\n', (4242, 4250), False, 'import asyncio\n'), ((9418, 9433), 'json.loads', 'json.loads', (['ret'], {}), '(ret)\n', (9428, 9433), False, 'import json\n'), ((3575, 3617), 'blangchain.tracking_utils.TokensTracker.update', 'TokensTracker.update', (['lm_output.llm_output'], {}), '(lm_output.llm_output)\n', (3595, 3617), False, 'from blangchain.tracking_utils import TokensTracker\n'), ((9012, 9048), 'langchain.PromptTemplate.from_template', 'PromptTemplate.from_template', (['PROMPT'], {}), '(PROMPT)\n', (9040, 9048), False, 'from langchain import LLMChain, PromptTemplate, FewShotPromptTemplate\n'), ((11334, 11381), 'langchain.PromptTemplate.from_template', 'PromptTemplate.from_template', (['combined_template'], {}), '(combined_template)\n', (11362, 11381), False, 'from langchain import LLMChain, PromptTemplate, FewShotPromptTemplate\n'), ((3315, 3388), 'langchain.schema.LLMResult', 'LLMResult', ([], {'generations': '[lm_out_i.generations[0] for lm_out_i in ret_list]'}), '(generations=[lm_out_i.generations[0] for lm_out_i in ret_list])\n', (3324, 3388), False, 'from langchain.schema import LLMResult\n'), ((3068, 3090), 'asyncio.gather', 'asyncio.gather', (['*tasks'], {}), '(*tasks)\n', (3082, 3090), False, 'import asyncio\n'), ((5450, 5500), 'langchain.prompts.HumanMessagePromptTemplate', 'HumanMessagePromptTemplate', ([], {'prompt': 'prompt_template'}), '(prompt=prompt_template)\n', (5476, 5500), False, 'from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, AIMessagePromptTemplate\n'), ((6195, 6211), 'json.loads', 'json.loads', (['line'], {}), '(line)\n', (6205, 6211), False, 'import json\n'), ((10752, 10799), 'langchain.PromptTemplate.from_template', 'PromptTemplate.from_template', (['combined_template'], {}), '(combined_template)\n', (10780, 10799), False, 'from langchain import LLMChain, PromptTemplate, FewShotPromptTemplate\n')]