import gradio as gr
from inference_script import generate_outreach_message, load_model_and_tokenizer
import os

HF_TOKEN = os.environ.get("HF_TOKEN")

try:
    model, tokenizer, device = load_model_and_tokenizer(
        "KirubaLS/Outreach_model_6",   
        "google/gemma-2b",             
        hf_token=HF_TOKEN
    )
except Exception as e:
    print(f"Error loading model and tokenizer: {e}")
    model = None
    tokenizer = None
    device = None


# Main prediction function
def predict_message(industry, company, prospect, designation, client_info, my_info):
    if model is None or tokenizer is None:
        return "Error: Model not loaded. Please check the logs."
    return generate_outreach_message(industry, company, prospect, designation, client_info, my_info)


# Build the Gradio app with blocks
with gr.Blocks(title="Gemma Outreach Generator") as demo:
    gr.Markdown("## Gemma Outreach Message Generator")
    gr.Markdown(
        "Generate a personalized first-level outreach message using your fine-tuned Gemma model."
    )

    with gr.Row():
        with gr.Column():
            industry    = gr.Textbox(label="Industry", placeholder="e.g. FinTech")
            company     = gr.Textbox(label="Company", placeholder="e.g. Stripe")
            prospect    = gr.Textbox(label="Prospect Name", placeholder="e.g. John Doe")
            designation = gr.Textbox(label="Designation", placeholder="e.g. Senior Developer")
            client_info = gr.Textbox(label="Client Info", placeholder="e.g. Expertise in payment APIs")
            my_info     = gr.Textbox(label="My Info", placeholder="e.g. I'm a growth strategist...")

        with gr.Column():
            output_box = gr.Textbox(label="Generated Outreach Message", lines=12)

    with gr.Row():
        generate_btn = gr.Button(" Generate Message")
        clear_btn    = gr.Button(" Clear Inputs")

    # Actions
    generate_btn.click(
        predict_message,
        inputs=[industry, company, prospect, designation, client_info, my_info],
        outputs=output_box,
        show_progress=True
    )

    clear_btn.click(
        fn=lambda: [""] * 6,
        inputs=[],
        outputs=[industry, company, prospect, designation, client_info, my_info]
    )

demo.launch()