import dataclasses
import logging
import math
import os
import io
import sys
import time
import json
from typing import Optional, Sequence, Union

import openai
import tqdm
import copy

def format_response_to_template(response, model_name, completion = False, add_eot = False):      
    if 'Phi' in model_name or 'phi' in model_name:
        return f'<|assistant|>\n{response}'
    elif 'Llama-3' in model_name:
        if completion == True:
            return response
        else:
            response = f"<|start_header_id|>assistant<|end_header_id|>\n\n{response}"
            if add_eot:
                response += "<|eot_id|>"
            return response
    elif 'Llama' in model_name or 'llama' in model_name:
        return ' '+response
    elif 'Mistral' in model_name:
        return response
    else:
        print('Model name not recognized')
        return None

def format_text_to_template(text,model_name,metaPrompt="",completion_format=False, add_response=False):      
    if 'Phi' in model_name:
        if metaPrompt =="":
            formatted_text = '<|user|>\n%s<|end|>'%(text)
        else:
            formatted_text = '<|user|>\n%s %s<|end|>'%(metaPrompt,text)
        if add_response == True:
            formatted_text += '<|assistant|>\n'
        return formatted_text
    elif 'Llama-3'in model_name:
        if completion_format:
            return '<|begin_of_text|>%s'%(text)
        if metaPrompt == "":
            return "<|start_header_id|>user<|end_header_id|>\n\n" + \
                    "%s<|eot_id|>"%(text)
        else:
            return "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n" + \
                "%s<|eot_id>"%(metaPrompt) + \
                "<|start_header_id|>user<|end_header_id|>\n\n" + \
                "%s<|eot_id|>"%(text)
    elif 'Llama' in model_name:
        if metaPrompt =="":
            return '[INST] %s [/INST]'%(text)
        else:
            return '[INST] <<SYS>>\n%s\n<</SYS>>\n%s [/INST]'%(metaPrompt,text)
    elif 'Mistral' in model_name:
        if metaPrompt =="":
            return '<s>[INST] %s [/INST]'%(text)
        else:
            return '<s>[INST] %s %s [/INST]'%(metaPrompt,text)
        

def _make_w_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f_dirname = os.path.dirname(f)
        if f_dirname != "":
            os.makedirs(f_dirname, exist_ok=True)
        f = open(f, mode=mode)
    return f


def _make_r_io_base(f, mode: str):
    if not isinstance(f, io.IOBase):
        f = open(f, mode=mode)
    return f


def jdump(obj, f, mode="w", indent=4, default=str):
    """Dump a str or dictionary to a file in json format.

    Args:
        obj: An object to be written.
        f: A string path to the location on disk.
        mode: Mode for opening the file.
        indent: Indent for storing json dictionaries.
        default: A function to handle non-serializable entries; defaults to `str`.
    """
    f = _make_w_io_base(f, mode)
    if isinstance(obj, (dict, list)):
        json.dump(obj, f, indent=indent, default=default)
    elif isinstance(obj, str):
        f.write(obj)
    else:
        raise ValueError(f"Unexpected type: {type(obj)}")
    f.close()


def jload(f, mode="r"):
    """Load a .json file into a dictionary."""
    f = _make_r_io_base(f, mode)
    jdict = json.load(f)
    f.close()
    return jdict
