import gradio as gr
from model import TinyAdd2LLM, Add2ModelError

# Load the model
try:
    model = TinyAdd2LLM.load("model/add2_model.bin")
except Exception as e:
    model = None
    print(f"Error loading model: {e}")

def calculate(prompt):
    if model is None:
        return "Error: Model not loaded."
    try:
        ans = model.answer(prompt)
        return ans
    except Add2ModelError as e:
        return f"Error: {e}"
    except Exception as e:
        return f"Unexpected Error: {e}"

# Build custom theme / UI
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="indigo", secondary_hue="slate"),
    title="TinyAdd2LLM - Addition Calculator"
) as demo:
    gr.Markdown(
        """
        # 🧮 TinyAdd2LLM
        
        This is a lightweight 10MB look-up-table "model" serving addition requests instantly. It supports addition of numbers from **0** to **2200**.
        
        ### Supported Prompt Formats:
        - `12 + 30`
        - `12 plus 30`
        - `add 12 and 30`
        - `what is 12 + 30?`
        - `what is 12 plus 30?`
        """
    )
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Your Prompt", 
                placeholder="e.g. 12 + 30, what is 999 plus 999?, add 15 and 7", 
                lines=1
            )
            submit_btn = gr.Button("Calculate", variant="primary")
        
        with gr.Column():
            output_text = gr.Textbox(label="Result", interactive=False)
    
    submit_btn.click(fn=calculate, inputs=input_text, outputs=output_text)
    input_text.submit(fn=calculate, inputs=input_text, outputs=output_text)
    
    gr.Examples(
        examples=[
            ["12 + 30"],
            ["add 15 and 7"],
            ["100 plus 55"],
            ["what is 999 + 999?"],
            ["2199 + 2199"]
        ],
        inputs=input_text
    )

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