# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_magic_line.ipynb.

# %% auto 0
__all__ = ['parse_cell', 'fr_line']

# %% ../nbs/01_magic_line.ipynb 5
import json
from typing import List, Dict, Tuple
import textwrap
import claudette
from IPython import get_ipython
from IPython.display import display, clear_output, Markdown, Javascript

# %% ../nbs/01_magic_line.ipynb 6
# A single cell can contain multiple messages.
# A message is either a user message (starts with %fr) or a bot message (starts with #).
# Both can be multiline.
def parse_cell(
    cell: str # The raw body of the cell
) -> Tuple[List[Dict[str, str]], int]:
    """
    A single cell can contain multiple messages.
    A message is either a user message (starts with %fr) or a bot message (starts with #).
    Both can be multiline.

    Returns: a list of messages (with 'role' and 'content') and the number of %fr magics in the cell
    """
    parsed_lines = []
    num_magic = 0
    for line in cell.split('\n'):
        if line.startswith('%fr'):
            message = {'role': 'user', 'content': line[3:].strip()}
            num_magic += 1
        elif line.strip().startswith('#'):
            message = {'role': 'assistant', 'content': line[1:].strip()}
        else: continue

        if not parsed_lines or parsed_lines[-1]['role'] != message['role']:
            parsed_lines.append(message)
        else:
            parsed_lines[-1]['content'] += ("\n" + message['content'])

    return parsed_lines, num_magic

# %% ../nbs/01_magic_line.ipynb 7
models = [
    'claude-3-opus-20240229',
    'claude-3-5-sonnet-20240620',
    'claude-3-haiku-20240307',
]
chat_client = claudette.Client(model=models[1])

magic_count = 0
messages = []

# %% ../nbs/01_magic_line.ipynb 8
def fr_line(line: str):
    """The magic function for the %fr magic command."""
    global magic_count, messages
    ip = get_ipython()
    # raw_cell = ip.history_manager.input_hist_raw[-1]
    raw_cell = ip.get_parent()["content"]["code"]

    # The cell might have multiple %lm magics, but we only want to process the last one.
    # Presumably, the previous ones would have been processed already.
    if magic_count <= 0:
        messages, magic_count = parse_cell(raw_cell)

    # This is the last %lm magic invocation of the cell.
    # But we ignore cells that don't have a user message as the last message.
    if magic_count == 1 and len(messages) > 0 and messages[-1]['role'] == 'user' and messages[-1]['content'].strip():
        reply = ""
        display_id = display(Markdown("🚀..."), display_id=True)
        try:
            r = chat_client([m['content'] for m in messages], stream=True)
            for token in r:
                reply += token
                display_id.update(Markdown(reply))

            if reply:
                reply = textwrap.fill(text=reply, width=100, initial_indent="# ", subsequent_indent="# ")
                raw_cell += f"\n{reply}\n\n%fr "
                ip.set_next_input(raw_cell, replace=True)

            clear_output()

        except BaseException as e:
            display_id.update(Markdown(f"🚫 {repr(e)}"))


    magic_count -= 1
