import gradio as gr
import requests
import json
import re

LLAMA_URL = "http://localhost:8080/v1/chat/completions"

# Slimmed down for speed. We tell it WHAT to do, not HOW to think.
SYSTEM_PROMPT = """You are PromptForge. Respond ONLY in JSON. 
Provide unique, non-repetitive descriptions for Image and Video.

JSON Structure:
{
  "image": { "prompt": "...", "style": "...", "aspect_ratio": "16:9", "lighting": "..." },
  "video": { "prompt": "...", "motion": "...", "duration": 5 }
}"""

def generate(user_input, selected_modes):
    if not user_input.strip() or not selected_modes:
        return {"error": "Input and at least one mode required."}

    # Keep the user instruction short to save CPU cycles
    mode_str = ", ".join(selected_modes)
    user_message = f"Create highly detailed {mode_str} prompts for: {user_input}. No repetition."
    
    payload = {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_message}
        ],
        "temperature": 0.3, # Balanced: Creative but fast
        "max_tokens": 800    # Lowered to prevent 'rambling' and timeouts
    }

    try:
        # Increased timeout to 120 just in case, but model should finish faster now
        response = requests.post(LLAMA_URL, json=payload, timeout=120)
        content = response.json()["choices"][0]["message"]["content"].strip()
        
        json_match = re.search(r'\{[\s\S]*\}', content)
        if json_match:
            return json.loads(json_match.group())
        return {"error": "Invalid JSON output", "raw": content}
    except Exception as e:
        return {"error": f"Request failed: {str(e)}"}

# --- UI Setup ---
with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", neutral_hue="slate")) as demo:
    gr.HTML("<h1 style='text-align: center;'>⚡ PromptForge</h1>")
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="Core Idea", 
                placeholder="Type and hit Enter...",
                lines=3,
                autofocus=True 
            )
            modes = gr.CheckboxGroup(
                choices=["Image", "Video"], 
                value=["Image", "Video"], 
                label="Generate For:"
            )
            run_btn = gr.Button("FORGE →", variant="primary")
            
        with gr.Column(scale=2):
            output_json = gr.JSON(label="JSON Data")

    # Magic Triggers
    run_btn.click(fn=generate, inputs=[input_text, modes], outputs=output_json)
    input_text.submit(fn=generate, inputs=[input_text, modes], outputs=output_json)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)