import nbformat
from nbformat.v4 import new_notebook, new_markdown_cell, new_code_cell
from nbconvert import HTMLExporter
from huggingface_hub import InferenceClient
from e2b_code_interpreter import Sandbox
from transformers import AutoTokenizer
from traitlets.config import Config

config = Config()
html_exporter = HTMLExporter(config=config, template_name="classic")


with open("llama3_template.jinja", "r") as f:
    llama_template = f.read() 


MAX_TURNS = 4


def parse_exec_result_nb(execution):
    """Convert an E2B Execution object to Jupyter notebook cell output format"""
    outputs = []
    
    if execution.logs.stdout:
        outputs.append({
            'output_type': 'stream',
            'name': 'stdout',
            'text': ''.join(execution.logs.stdout)
        })
    
    if execution.logs.stderr:
        outputs.append({
            'output_type': 'stream',
            'name': 'stderr',
            'text': ''.join(execution.logs.stderr)
        })

    if execution.error:
        outputs.append({
            'output_type': 'error',
            'ename': execution.error.name,
            'evalue': execution.error.value,
            'traceback': [line for line in execution.error.traceback.split('\n')]
        })

    for result in execution.results:
        output = {
            'output_type': 'execute_result' if result.is_main_result else 'display_data',
            'metadata': {},
            'data': {}
        }
        
        if result.text:
            output['data']['text/plain'] = [result.text]  # Array for text/plain
        if result.html:
            output['data']['text/html'] = result.html
        if result.png:
            output['data']['image/png'] = result.png
        if result.svg:
            output['data']['image/svg+xml'] = result.svg
        if result.jpeg:
            output['data']['image/jpeg'] = result.jpeg
        if result.pdf:
            output['data']['application/pdf'] = result.pdf
        if result.latex:
            output['data']['text/latex'] = result.latex
        if result.json:
            output['data']['application/json'] = result.json
        if result.javascript:
            output['data']['application/javascript'] = result.javascript

        if result.is_main_result and execution.execution_count is not None:
            output['execution_count'] = execution.execution_count

        if output['data']:
            outputs.append(output)

    return outputs


system_template = """\
<details>
  <summary style="display: flex; align-items: center;">
    <div class="alert alert-block alert-info" style="margin: 0; width: 100%;">
      <b>System: <span class="arrow">▶</span></b>
    </div>
  </summary>
  <div class="alert alert-block alert-info">
    {}
  </div>
</details>

<style>
details > summary .arrow {{
  display: inline-block;
  transition: transform 0.2s;
}}
details[open] > summary .arrow {{
  transform: rotate(90deg);
}}
</style>
"""

user_template = """<div class="alert alert-block alert-success">
<b>User:</b> {}
</div>
"""

header_message = """<p align="center">
  <img src="https://huggingface.co./spaces/lvwerra/jupyter-agent/resolve/main/jupyter-agent.png" />
</p>


<p style="text-align:center;">Let a LLM agent write and execute code inside a notebook!</p>"""

bad_html_bad = """input[type="file"] {
  display: block;
}"""


def create_base_notebook(messages):
    base_notebook = {
        "metadata": {
            "kernel_info": {"name": "python3"},
            "language_info": {
                "name": "python",
                "version": "3.12",
            },
        },
        "nbformat": 4,
        "nbformat_minor": 0,
        "cells": []
    }
    base_notebook["cells"].append({
            "cell_type": "markdown",
            "metadata": {},
            "source": header_message
            })

    if len(messages)==0:
        base_notebook["cells"].append({
                            "cell_type": "code",
                            "execution_count": None,
                            "metadata": {},
                            "source": "",
                            "outputs": []
                        })

    code_cell_counter = 0
    
    for message in messages:
        if message["role"] == "system":
            text = system_template.format(message["content"].replace('\n', '<br>'))
            base_notebook["cells"].append({
                "cell_type": "markdown",
                "metadata": {},
                "source": text
                })
        elif message["role"] == "user":
            text = user_template.format(message["content"].replace('\n', '<br>'))
            base_notebook["cells"].append({
                "cell_type": "markdown",
                "metadata": {},
                "source": text
                })

        elif message["role"] == "assistant" and "tool_calls" in message:
            base_notebook["cells"].append({
                "cell_type": "code",
                "execution_count": None,
                "metadata": {},
                "source": message["content"],
                "outputs": []
            })

        elif message["role"] == "ipython":
            code_cell_counter +=1
            base_notebook["cells"][-1]["outputs"] = message["nbformat"]
            base_notebook["cells"][-1]["execution_count"] = code_cell_counter

        elif message["role"] == "assistant" and "tool_calls" not in message:
            base_notebook["cells"].append({
                "cell_type": "markdown",
                "metadata": {},
                "source": message["content"]
            })
            
        else:
            raise ValueError(message)
        
    return base_notebook, code_cell_counter

def execute_code(sbx, code):
    execution = sbx.run_code(code, on_stdout=lambda data: print('stdout:', data))
    output = ""
    if len(execution.logs.stdout) > 0:
        output += "\n".join(execution.logs.stdout)
    if len(execution.logs.stderr) > 0:
        output += "\n".join(execution.logs.stderr)
    if execution.error is not None:
        output += execution.error.traceback
    return output, execution


def parse_exec_result_llm(execution):
    output = ""
    if len(execution.logs.stdout) > 0:
        output += "\n".join(execution.logs.stdout)
    if len(execution.logs.stderr) > 0:
        output += "\n".join(execution.logs.stderr)
    if execution.error is not None:
        output += execution.error.traceback
    return output
    
    
def update_notebook_display(notebook_data):
    notebook = nbformat.from_dict(notebook_data)
    notebook_body, _ = html_exporter.from_notebook_node(notebook)
    notebook_body = notebook_body.replace(bad_html_bad, "")
    return notebook_body

def run_interactive_notebook(client, model, tokenizer, messages, sbx, max_new_tokens=512):
    notebook_data, code_cell_counter = create_base_notebook(messages)
    turns = 0
    
    #code_cell_counter = 0
    while turns <= MAX_TURNS:
        turns += 1
        input_tokens = tokenizer.apply_chat_template(
            messages,
            chat_template=llama_template,
            builtin_tools=["code_interpreter"], 
            add_generation_prompt=True
        )
        model_input = tokenizer.decode(input_tokens)

        print(f"Model input:\n{model_input}\n{'='*80}")
        
        response_stream = client.text_generation(
            model=model,
            prompt=model_input,
            details=True,
            stream=True,
            do_sample=True,
            repetition_penalty=1.1,
            temperature=0.8,
            max_new_tokens=max_new_tokens,
        )
        
        assistant_response = ""
        tokens = []
        
        code_cell = False
        for i, chunk in enumerate(response_stream):
            if not chunk.token.special:
                content = chunk.token.text
            else:
                content = ""
            tokens.append(chunk.token.text)                
            assistant_response += content

            if len(tokens)==1:
                create_cell=True
                code_cell = "<|python_tag|>" in tokens[0]
                if code_cell:
                    code_cell_counter +=1
            else:
                create_cell = False
            
            # Update notebook in real-time
            if create_cell:
                if "<|python_tag|>" in tokens[0]:
                    notebook_data["cells"].append({
                        "cell_type": "code",
                        "execution_count": None,
                        "metadata": {},
                        "source": assistant_response,
                        "outputs": []
                    })
                else:
                    notebook_data["cells"].append({
                        "cell_type": "markdown",
                        "metadata": {},
                        "source": assistant_response
                    })
            else:
                notebook_data["cells"][-1]["source"] = assistant_response
            if i%16 == 0:
                yield update_notebook_display(notebook_data), notebook_data, messages
        yield update_notebook_display(notebook_data), notebook_data, messages


        # Handle code execution
        if code_cell:
            notebook_data["cells"][-1]["execution_count"] = code_cell_counter

            
            exec_result, execution = execute_code(sbx, assistant_response)
            messages.append({
                "role": "assistant",
                "content": assistant_response,
                "tool_calls": [{
                    "type": "function",
                    "function": {
                        "name": "code_interpreter",
                        "arguments": {"code": assistant_response}
                    }
                }]
            })
            messages.append({"role": "ipython", "content": parse_exec_result_llm(execution), "nbformat": parse_exec_result_nb(execution)})
            
            # Update the last code cell with execution results
            notebook_data["cells"][-1]["outputs"] = parse_exec_result_nb(execution)
            update_notebook_display(notebook_data)
        else:
            messages.append({"role": "assistant", "content": assistant_response})
            if tokens[-1] == "<|eot_id|>":
                break
    
    yield update_notebook_display(notebook_data), notebook_data, messages