import gradio as gr
from huggingface_hub import InferenceClient
import os
import json

# Using the exact model you specified
MODEL_ID = "google/gemma-3n-e4b-it"

client = InferenceClient(api_key=os.environ.get("HF_TOKEN"))

def generate(current_text, tone, length, context, recent_steps):
    prompt = f"""
    You are a precise senior copy editor. Rewrite the provided copy using the requested controls.
    Create exactly three meaningfully different options for the user to choose from.
    Each option must follow the same selected tone and length.
    Return valid JSON exactly like this: {{"options":["first rewrite","second rewrite","third rewrite"]}}
    
    Tone: {tone}
    Length: {length}
    Context: {context}
    Recent Steps: {recent_steps}
    
    Current copy:
    {current_text}
    """
    
    try:
        completion = client.chat_completion(
            model=MODEL_ID,
            messages=[
                {"role": "system", "content": "Return valid JSON only. No commentary."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=500,
            temperature=0.7
        )
        return completion.choices[0].message.content
    except Exception as e:
        return json.dumps({"error": str(e)})

# Expose as a Gradio API
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="current_text"),
        gr.Textbox(label="tone"),
        gr.Textbox(label="length"),
        gr.Textbox(label="context"),
        gr.Textbox(label="recent_steps")
    ],
    outputs=gr.JSON(),
    api_name="predict"
)

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