import os
import json
import time
import requests
import openai
import copy
import math

from loguru import logger


DEBUG = int(os.environ.get('DEBUG', '0'))
VLLM=int(os.environ.get('VLLM', 0))
STOP_TOKEN=os.environ.get('STOP_TOKEN', None)
MAX_NNN = 2

def remove_repeated_chars(s, n):
    if n < 1:
        return ""

    result = []
    count = 1

    for i in range(1, len(s) + 1):
        if i < len(s) and s[i] == s[i - 1]:
            count += 1
        else:
            if count <= n:
                result.extend([s[i - 1]] * count)
            count = 1

    return ''.join(result)


def generate_together(
    model, messages, n=1, logprobs=0,
    max_tokens=2048, temperature=0.7,base_instruction=None
):

    if logprobs:
        assert n == 1 and max_tokens == 1
    cur_logprobs = None
    inputs = []
    outputs = []
    halfed=0

    # for sleep_time in [1, 2, 4, 8, 8, 8, 16, 32, 2, 2]:
    sleep_time = 1

    MAX_n = MAX_NNN
    if "microsoft/WizardLM-2-8x22B" in model:
        MAX_n = 1
    number_of_requests = math.ceil(n / MAX_n)
    current_n = min(n, MAX_n)
    for i in range(number_of_requests):
        cur_outputs = None
        cur_inputs = None
        while cur_outputs is None:
            returned = None
            try:
                if VLLM:
                    endpoint = f"{os.environ['VLLMURL']}/chat/completions"
                else:
                    endpoint = 'https://api.together.xyz/v1/chat/completions'
                if DEBUG:
                    logger.debug(f"Sending messages ({len(messages)}) (last message: `{messages[-1]['content'][:20]}...`) to `{model}`.")
                
                if VLLM:
                    res = requests.post(endpoint, json={
                        "model": model,
                        "max_tokens": max_tokens,
                        "temperature": (temperature if temperature > 1e-4 else 0),
                        "messages": messages,
                        "n": current_n,
                        "logprobs": logprobs,
                        "stop_token_ids": [int(STOP_TOKEN)] if STOP_TOKEN else None,
                    })
                else:
                    res = requests.post(endpoint, json={
                        "model": model,
                        "max_tokens": max_tokens,
                        "temperature": (temperature if temperature > 1e-4 else 0),
                        "messages": messages,
                        "n": current_n,
                        "logprobs": logprobs,
                    }, headers={
                        "Authorization": f"Bearer {os.environ['OPENSOURCE_API_KEY']}",
                })

                returned = res.json()
                cur_logprobs = res.json()['choices'][0]['logprobs']
                cur_inputs = [messages for x in range(current_n)] # list of list of dicts
                cur_outputs = [choice['message']['content'] for choice in res.json()['choices']]
                break
            except Exception as e:
                logger.error(e)
                logger.error(returned)
                if returned is not None and ('error' in returned or "object" in returned):
                    too_long = False
                    if "error" not in returned and returned['object'] == 'error':
                        too_long = True
                    elif "type" in returned['error'] and (returned['error']['type'] == 'invalid_request_error' or returned['error']['type'] == 'BadRequestError'):
                        too_long = True
                    if too_long:
                        logger.info("Input + output is longer than max_position_id.")
                        if halfed<4:
                            max_tokens = max_tokens/2
                            halfed+=1
                        else:
                            logger.info("fall back to default mode")
                            messages = base_instruction

                logger.info(f"Retry in {sleep_time}s.. with max_tokens: {max_tokens}")
                time.sleep(sleep_time)
                sleep_time = min(sleep_time * 2, 32)
        inputs.extend(cur_inputs)
        outputs.extend(cur_outputs)
        current_n = min(n - (i+1)*MAX_n, MAX_n)

    if len(outputs) == 0:
        logger.debug(f"Output None...")
        return outputs
        
    outputs = [output.strip() for output in outputs]

    if DEBUG:
        logger.debug(f"{len(outputs)} outputs. Sample Output: `{outputs[0][:20]}...`.")

    return {"outputs":outputs,"inputs":inputs, "logprobs":cur_logprobs}


def generate_openai(
    model, messages,  n=1,
    max_tokens=2048, temperature=0.7,base_instruction=None,logprobs=False
):

    client = openai.OpenAI(
        api_key = API_KEY,
    )
    cur_logprobs = None
    outputs = None
    for sleep_time in [16, 1, 2, 4, 8,16,16,16, 32]:
        try:

            if DEBUG:
                logger.debug(f"Sending messages ({len(messages)}) (last message: `{messages[-1]['content'][:20]}`) to `{model}`.")

            completion = client.chat.completions.create(
                model=model,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                n=n,
                logprobs=logprobs
            )
            outputs = [x.message.content for x in completion.choices]
            cur_logprobs = completion.choices[0].logprobs
            break
        
        except Exception as e:
            logger.error(e)
            logger.error(messages)
            messages[0]['content'] = remove_repeated_chars(messages[0]['content'], 100)
            logger.error(messages)
            logger.info(f"Retry in {sleep_time}s..")
            time.sleep(sleep_time)

    inputs = [messages for x in range(len(outputs))]
    outputs = [output.strip() for output in outputs]

    return {"outputs":outputs,"inputs":inputs, "logprobs":cur_logprobs}


def inject_references_to_messages(
    messages, references, mt_bench=False,
):

    messages = copy.deepcopy(messages)
    
    system = f'''You have been provided with a set of responses from various open-source models to the latest user query. Your task is to synthesize these responses into a single, high-quality response. It is crucial to critically evaluate the information provided in these responses, recognizing that some of it may be biased or incorrect. Your response should not simply replicate the given answers but should offer a refined, accurate, and comprehensive reply to the instruction. Ensure your response is well-structured, coherent, and adheres to the highest standards of accuracy and reliability.

Responses from models:'''

    for i, reference in enumerate(references):
        
        system += f"\n{i+1}. {reference}"

    if messages[0]['role'] == 'system':

        messages[0]['content'] += '\n\n' + system

    else:
        if mt_bench:
            messages = messages[:-1] + [{'role': 'user', 'content': system + "\n\nUser Query: "+ messages[-1]["content"]}] 
        else:
            messages = [
                {'role': 'system', 'content': system}
            ] + messages

    return messages


def generate_with_references(
    model, messages, references=[],
    max_tokens=2048, temperature=0.7, n=1,
    generate_fn=generate_together, aggPrompt=None, mt_bench=False,
):
    base_instruction = messages.copy()
    if len(references) > 0:
        messages = inject_references_to_messages(messages, references,mt_bench=mt_bench)
    if "gpt" in model and "/" not in model:
        generate_fn = generate_openai
    return generate_fn(model=model, messages=messages, temperature=temperature, max_tokens=max_tokens, n=n, base_instruction = base_instruction)
    