"""The core package contains the `Chat` class central to this package."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['show', 'Chat', 'print_conversation']

# %% ../nbs/00_core.ipynb 3
import os
import requests
import json
from datetime import datetime
from IPython.display import Markdown, display_markdown
from fastcore.basics import patch
from dotenv import load_dotenv

# %% ../nbs/00_core.ipynb 6
def show(string): return Markdown(string)

# %% ../nbs/00_core.ipynb 7
class Chat():
    def __init__(self, model):
        self.model = model
        self.usage = []
        self.context = []
        self.date = None
        self.time = None
        self.history = []
        self.chatsdir = ''
        self.title = f'conversation-with-{self.model.split("/")[1].replace(":", "-")}'
        
    def __call__(self, prompt):
        self.context.append({"role": "user", "content": prompt})
        resp = requests.post(url="https://openrouter.ai/api/v1/chat/completions", 
                             headers={"Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}"},
                             json={"model": self.model,
                                   "messages": self.context})
        ans = resp.json()['choices'][0]['message']
        self.context.append(ans)
        self.usage.append(resp.json()['usage'])
        self.date = resp.json()['created']
        self.history.append(resp.json())
        self.chatsdir = f"chats/{self.history[-1]['model']}/"
        return self._show(ans['content'])
        
    def _show(self, string): return Markdown(string)
        
    def usage_summary(self):
        intok = [u['prompt_tokens'] for u in self.usage]
        outok = [u['completion_tokens'] for u in self.usage]
        totok = [u['total_tokens'] for u in self.usage]
        print(f"Input tokens: {sum(intok)}\nOutput tokens: {sum(outok)}\nTotal tokens: {sum(totok)}")
        
    def forget_last(self):
        self.context = self.context[:-1]
        return self.context

# %% ../nbs/00_core.ipynb 8
@patch
def save_conversation(self:Chat, generate_title=True):
    conv = json.dumps(self.context, ensure_ascii=False, indent=2)
    self.time = datetime.fromtimestamp(self.date).strftime('%Y%m%d-%H%M%S')
    os.makedirs(self.chatsdir, exist_ok=True)
    if generate_title:
        self.context.append({"role": "user", "content": "Generate an all lower case title for our conversation connecting the words with hyphens (-) instead of spaces. Output nothing but the title. Five words max!"})
        resp = requests.post(url="https://openrouter.ai/api/v1/chat/completions", 
                             headers={"Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}"},
                             json={"model": self.model,
                                   "messages": self.context})
        self.usage.append(resp.json()['usage'])
        self.title = resp.json()['choices'][0]['message']['content'].strip()
        self.forget_last()
    
    file_path = self.chatsdir + self.time + "-" + self.title + ".txt"
    with open(file_path, 'w') as f:
        f.write(conv)
    return "Saving complete!"

# %% ../nbs/00_core.ipynb 19
@patch
def print_conversation(self:Chat, user="Me"):
    assistant = self.model.split("/")[1].split("-")[0].capitalize()
    display_markdown(self._show(f"**{(" ".join(word.capitalize() for word in self.title.split("-")))}**"))
    for i in self.context:
        if i['role'] == 'user':
            display_markdown(self._show(f"**{user}**: {i['content']}"))
        elif i['role'] == 'assistant':
            display_markdown(self._show(f"**{assistant}**: {i['content']}"))

# %% ../nbs/00_core.ipynb 21
def print_conversation(conversation, user="User", assistant="Assistant"):
    for i in conversation:
        if i['role'] == 'user':
            display_markdown(show(f"**{user}**: {i['content']}"))
        elif i['role'] == 'assistant':
            display_markdown(show(f"**{assistant}**: {i['content']}"))
