import os
from dotenv import load_dotenv
import anthropic
import gradio as gr

load_dotenv()

os.environ['ANTHROPIC_API_KEY'] = os.getenv('ANTHROPIC_API_KEY')
claude = anthropic.Anthropic()
CLAUDE_MODEL = "claude-3-5-sonnet-20241022"

system_message = "You are an assistant that reimplements Python code in high performance C++. \
    Respond only with C++ code; use comments sparingly and do not provide any explanation other than occasional comments. \
    The C++ response needs to produce an identical output in the fastest possible time."

def user_prompt_for(python):
    user_prompt = f"Rewrite this Python code in C++ with the fastest possible implementation that produces identical output in the least time. \
        Respond only with C++ code; do not explain your work other than a few comments. \
        Pay attention to number types to ensure no int overflows. Remember to #include all necessary C++ packages such as iomanip.\
        \n\n{python}"
    
    return user_prompt

def rewrite(python):
    result = claude.messages.stream(
        model=CLAUDE_MODEL,
        max_tokens=2000,
        system=system_message,
        messages=[{"role": "user", "content": user_prompt_for(python)}],
    )

    reply = ""
    with result as stream:
        for text in stream.text_stream:
            reply += text
            yield reply.replace('```cpp\n','').replace('```','')

demo = gr.Interface(
            fn=rewrite,
            inputs=gr.Code(label="Python code:", language="python", lines=10),
            outputs=gr.Code(label="C++ code:", language="cpp", lines=10)
        )
demo.launch()