import groq
import gradio as gr
import re
import json
import os


# Setup Groq
client = groq.Client(api_key=os.environ.get("GROQ_API_KEY"))  # Correct initialization
model = "llama-3.3-70b-versatile"

def handle_groq_error(e, model_name):
    error_data = e.args[0]
    
    if isinstance(error_data, str):
        # Use regex to extract the JSON part of the string
        json_match = re.search(r'(\{.*\})', error_data)
        if json_match:
            json_str = json_match.group(1)
            # Ensure the JSON string is well-formed
            json_str = json_str.replace("'", '"')  # Replace single quotes with double quotes
            error_data = json.loads(json_str)

    if isinstance(e, groq.RateLimitError):
        if isinstance(error_data, dict) and 'error' in error_data and 'message' in error_data['error']:
            error_message = error_data['error']['message']
            raise gr.Error(error_message)
    else:
        raise gr.Error(f"Error during Groq API call: {e}")



def calculate_and_roast(expression):
    try:
        # Basic calculator functionality (limited for safety)
        allowed_chars = set("0123456789+-*/(). ")
        if not all(char in allowed_chars for char in expression):
            return "Invalid input: Only numbers and basic operators (+, -, *, /) are allowed.", None

        result = eval(expression)

        # Groq API call
        try:
            prompt = f"Roast me for calculating something as simple as '{expression}'. Make it funny."
            response = client.request(model, {"prompt": prompt})
            roast = response["text"][0]
            return result, roast

        except groq.GroqApiException as e:
             handle_groq_error(e, model)

    except (ValueError, TypeError, SyntaxError, ZeroDivisionError) as e:
        return f"Calculation error: {e}", None



with gr.Blocks() as interface:
    input_expression = gr.Textbox(label="Enter an expression")
    calculate_button = gr.Button("Calculate and Roast")
    output_result = gr.Textbox(label="Result")
    output_roast = gr.Textbox(label="Roast")

    calculate_button.click(
        fn=calculate_and_roast,
        inputs=input_expression,
        outputs=[output_result, output_roast],
    )


interface.launch(share=True)