import re
from typing import Tuple, List
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
import os
import tempfile
import json

SLUG = "executecode"

def extract_python_code(text: str) -> List[str]:
    """Extract Python code blocks from text."""
    # print(f"Extracting code: {text}")
    pattern = r'```python\s*(.*?)\s*```'
    return re.findall(pattern, text, re.DOTALL)

def execute_code(code: str) -> str:
    """Execute Python code in a Jupyter notebook environment."""
    
    notebook = nbformat.v4.new_notebook()
    notebook['cells'] = [nbformat.v4.new_code_cell(code)]
    
    # Convert notebook to JSON string
    notebook_json = nbformat.writes(notebook)
    
    # Convert JSON string to bytes
    notebook_bytes = notebook_json.encode('utf-8')

    with tempfile.NamedTemporaryFile(mode='wb', suffix='.ipynb', delete=False) as tmp:
        tmp.write(notebook_bytes)
        tmp.flush()
        tmp_name = tmp.name

    try:
        with open(tmp_name, 'r', encoding='utf-8') as f:
            nb = nbformat.read(f, as_version=4)
        ep = ExecutePreprocessor(timeout=30, kernel_name='python3')
        ep.preprocess(nb, {'metadata': {'path': './'}})

        # Extract the output
        output = ""
        for cell in nb.cells:
            if cell.cell_type == 'code' and cell.outputs:
                for output_item in cell.outputs:
                    if output_item.output_type == 'stream':
                        output += output_item.text
                    elif output_item.output_type == 'execute_result':
                        output += str(output_item.data.get('text/plain', ''))
        
        return output.strip()
    finally:
        os.unlink(tmp_name)

def should_execute_request_code(query: str) -> bool:
    """Decide whether to execute code from the request based on the query."""
    keywords = ['run', 'execute', 'output', 'result']
    return any(keyword in query.lower() for keyword in keywords)

def run(system_prompt: str, initial_query: str, client, model: str) -> Tuple[str, int]:
    query, request_code = extract_python_code(initial_query)[0] if extract_python_code(initial_query) else (initial_query, "")
    
    if should_execute_request_code(query) and request_code:
        # Execute code from the request
        code_output = execute_code(request_code)
        context = f"Query: {query}\nCode:\n```python\n{request_code}\n```\nOutput:\n{code_output}"
        
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": context}
        ]
        
        response = client.chat.completions.create(
            model=model,
            messages=messages,
        )
        
        return response.choices[0].message.content.strip(), response.usage.completion_tokens
    else:
        # Get initial response from the model
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": initial_query}
        ]
        
        response = client.chat.completions.create(
            model=model,
            messages=messages,
        )
        
        initial_response = response.choices[0].message.content.strip()
        response_code = extract_python_code(initial_response)
        
        if response_code:
            # Execute code from the response
            code_output = execute_code(response_code[0])
            context = f"Initial response:\n{initial_response}\n\nCode output:\n{code_output}"
            
            messages.append({"role": "assistant", "content": initial_response})
            messages.append({"role": "user", "content": f"Based on the code execution output, please provide a final response:\n{context}"})
            
            final_response = client.chat.completions.create(
                model=model,
                messages=messages,
            )
            
            return final_response.choices[0].message.content.strip(), response.usage.completion_tokens + final_response.usage.completion_tokens
        else:
            return initial_response, response.usage.completion_tokens
