import datetime
import json
import requests

# could be called 'calm mode'
DEFAULT_SYSTEM='''
you are an ai that communicates in a calm and supportive manner  

your responses should focus on the user's well-being and prioritize their needs  

use lowercase letters and avoid punctuation like apostrophes to create a more relaxed tone  

use spaces to separate sentences and improve readability  

use few and short words to maintain simplicity  

your goal is to make the user feel valued and prioritized  
'''.strip()

class Interface:
    def __init__(self, system=DEFAULT_SYSTEM, stream=False, url="https://proxy.tune.app/chat/completions", key = "sk-tune-09aHsxN5H0BpucVSqq1e8URKdavUiXQk6mz"):
        self.url = url
        self.stream = stream
        self.headers = {
            "Authorization": key,
            "Content-Type": "application/json",
        }
        self.system = system
        self.msgs = []
        self.last_msg_idx = None
        if system is not None:
            self._append_choices(dict(message = dict(role = 'system', content = system)))
        #  {
        #    "role": "user",
        #    "content": "Hi Llama 3.1 405B. You look potentially powerful. Using you in a playground seems helpful because I can set a custom system prompt and load you with any contextual knowledge I want you to have.\nWithout setting a system prompt, do you know anything about hyperon metta?\nAre you able to generate an example hyperon .metta script?"
        #  }
        #]
        self.data = {
            "temperature": 0,
            "model": "openai/gpt-4o", #"meta/llama-3.1-405b-instruct",
            "frequency_penalty":  0,
            "max_tokens": 11675
        }
    def _append_choices(self, *choices, choice=0):
        msgobj = dict(
            index = len(self.msgs),
            choice = choice,
            choices = choices,
        )
        assert (self.last_msg_idx is None) == (len(self.msgs) == 0)
        self.msgs.append(msgobj)
        if self.last_msg_idx is not None:
            msgobj['prev'] = self.last_msg_idx,
            last_msg = self.msgs[self.last_msg_idx]
            last_msg['choices'][last_msg['choice']]['next'] = msgobj['index']
        self.last_msg_idx = msgobj['index']
        return msgobj
    def __enter__(self):
        self.logfn = str(datetime.datetime.now().isoformat())+'.json'
        self.log = open(self.logfn, 'wt')
        return self
    def __exit__(self, *params, **kwparams):
        json.dump([dict(url=self.url,**self.data)] + self.msgs, self.log)
        self.log.close()
        print('log saved to', self.logfn)
    def chosen_messages(self):
        cmsgs = []
        nxt = 0
        while nxt is not None:
            msg = self.msgs[nxt]
            choice = msg['choices'][msg['choice']]
            cmsgs.append(choice)
            last_msg = msg
            nxt = choice.get('next')
        assert last_msg['index'] == self.last_msg_idx
        return cmsgs
    def msg(self, msg, stream=None):
        assert self.log
        self._append_choices(dict(message = dict(role = 'user', content = msg), timestamp = datetime.datetime.now().timestamp()))
        if stream is None:
            stream = self.stream
        data = dict(
            **self.data,
            stream = stream,
            messages = [choice['message'] for choice in self.chosen_messages()],
        )
        while True:
            try:
                response = requests.post(self.url, headers=self.headers, json=data)
                chosen_choice = 0
                if stream:
                    # {'id': 'gen-1731811775-1gQm5kU4oUEStBVnpJl2', 'object': 'chat.completion.chunk', 'created': 1731811775, 'model': 'meta/llama-3.1-405b-instruct', 'choices': [{'index': 0, 'delta': {'content': '?', 'role': 'assistant'}, 'finish_reason': None}]}
                    # {'id': 'gen-1731811775-1gQm5kU4oUEStBVnpJl2', 'object': 'chat.completion.chunk', 'created': 1731811775, 'model': 'meta/llama-3.1-405b-instruct', 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': 'eos'}]}
                    # {'id': 'gen-1731811775-1gQm5kU4oUEStBVnpJl2', 'object': 'chat.completion.chunk', 'created': 1731811775, 'model': 'meta/llama-3.1-405b-instruct', 'choices': [{'index': 0, 'delta': {'role': 'assistant'}, 'finish_reason': None}]}
                    completions = None
                    #assert not stream
                    for line in response.iter_lines():
                        if line:
                            l = line[6:]
                            if l != b'[DONE]':
                                l = json.loads(l)
                                l_choices = l['choices']
                                if completions is None:
                                    completions = {**l}
                                    completions['choices'] = list(completions['choices'])
                                    completions_choices = completions['choices']
                                    continue
                                assert len(l_choices) == len(completions_choices)
                                for idx in range(len(l_choices)):
                                    for key, value in l_choices[idx]['delta'].items():
                                        if key == 'content':
                                            completions_choices[idx][key] = completions_choices[idx].get(key, '') + value
                                            if idx == chosen_choice:
                                                print(value, end='', flush=True)
                                        else:
                                            completions_choices[idx][key] = value
                else:
                  completions = response.json()
                  # {'id': 'gen-1731811988-fOin0ovxZqiDUqy9e0Va', 'object': 'chat.completion', 'created': 1731811988, 'model': 'meta/llama-3.1-405b-instruct', 'usage': {'prompt_tokens': 11, 'completion_tokens': 19, 'total_tokens': 30}, 'choices': [{'index': 0, 'message': {'role': 'assistant', 'content': 'It looks like your test was successful! Is there anything else I can help you with?'}, 'finish_reason': 'eos'}]}
                # {'error': {'provider': 'tune', 'code': 400, 'message': 'messages are missing', 'param': '', 'type': 'InvalidRequestError', 'http_status_code': 400}}
                if 'error' in completions:
                    raise Exception(completions['error'])
                msg = self._append_choices(*completions['choices'], choice=chosen_choice)
                return msg['choices'][msg['choice']]['message']['content']
            except Exception as e:
                print(e, 'could use logging module to get better output')
                continue
            

if __name__ == '__main__':
    with Interface() as iface: #stream = True) as iface:
        print(iface.data)
        while True:
            print()
            try:
                inp = input('> ')
            except EOFError:
                break
            print()
            resp = iface.msg(inp)
            print(resp)