import anthropic
import base64
import gradio as gr
import mimetypes

def generate_playground_link(code):
    """
    Generate a Gradio app link from the code.
    """
    base64_code = base64.b64encode(code.encode()).decode()
    return f"https://www.gradio.app/playground?code={base64_code}&reqs=&demo=Blank"

def extract_code(message):
    """
    Extract the code from a message that is formatted in Markdown with triple backticks.
    Include the triple backticks in the return value.
    """
    return message.split('```python')[1].split('```')[0]


def get_response(data, history):
    """
    Get a response from the Anthropic API based on the image and user message.
    """
    client = anthropic.Anthropic()
    response = []
    if data["files"]:
        response.append("Generating Gradio app based on the image...")
        yield response
        image = data["files"][0]  # str filepath
        image_data = base64.standard_b64encode(open(image, "rb").read()).decode("utf-8")
        image_media_type = mimetypes.guess_type(image)[0]
        message = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=1024,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image",
                            "source": {
                                "type": "base64",
                                "media_type": image_media_type,
                                "data": image_data,
                            },
                        },
                        {
                            "type": "text",
                            "text": "Based on the screenshot, write the Python code to generate this Gradio app. Try to recreate the Gradio UI, not necessarily the underlying data or machine learning model. Respond in Markdown enclosing the code in triple backticks."
                        }
                    ],
                }
            ],
        )
    else:
        response.append("Generating Gradio app based on the description...")
        yield response
        request = data["text"]
        message = client.messages.create(
            model="claude-3-5-sonnet-20241022",
            max_tokens=1024,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": request
                        },
                        {
                            "type": "text",
                            "text": "Based on the user's request, write the Python code to generate this Gradio app. Focus on the Gradio UI, not necessarily the underlying data or machine learning model. Respond in Markdown enclosing the code in triple backticks."
                        }
                    ],
                }
            ],
        )
    code = extract_code(message.content[0].text)
    formatted_code = f"```python\n{code}\n```"
    response.append(formatted_code)
    yield response
    response.append("[Try it out in the Gradio playground](" + generate_playground_link(code) + ")")
    yield response


demo = gr.ChatInterface(
    get_response, 
    type="messages", 
    multimodal=True,
    title="Gradio Playground Bot",
)

if __name__ == "__main__":
    demo.launch()