import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import torch
import json
import re

MODEL_ID = "cahya/gpt2-small-indonesian-522M"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, device_map={"": "cpu"})

generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

def clean_and_extract_json(text):
    json_start = text.find('{')
    if json_start >= 0:
        text = text[json_start:]
    text = re.sub(r"```json|```", "", text).strip()
    return text

def generate_ui(prompt):
    out = generator(prompt, max_new_tokens=180, do_sample=True, top_p=0.95, top_k=40, temperature=0.8)
    raw_text = out[0]["generated_text"]

    cleaned_text = clean_and_extract_json(raw_text)
    try:
        parsed = json.loads(cleaned_text)
        script = parsed.get("script", "").strip()
        caption = parsed.get("caption", "").strip()
        hashtags = parsed.get("hashtags", [])
        if not script or not caption or not isinstance(hashtags, list):
            raise ValueError("Missing or invalid fields")
        return {
            "script": script,
            "caption": caption,
            "hashtags": hashtags
        }
    except Exception:
        fallback_script = raw_text.strip()
        if prompt in fallback_script:
            fallback_script = fallback_script.replace(prompt, "").strip()
        fallback_caption = fallback_script.split(".")[0].strip() or "Simak info menarik ini!"
        return {
            "script": fallback_script,
            "caption": fallback_caption,
            "hashtags": ["#indonesia", "#fyp"]
        }

# Use gr.Blocks to enable enable_api properly
with gr.Blocks() as demo:
    prompt_input = gr.Textbox(lines=6, placeholder="Masukkan prompt...")
    output_json = gr.JSON()
    generate_btn = gr.Button("Generate")

    generate_btn.click(fn=generate_ui, inputs=prompt_input, outputs=output_json)

    gr.Markdown("# TikTok Script Generator")
    gr.Markdown("Generates short Indonesian scripts (JSON).")

if __name__ == "__main__":
    # IMPORTANT: Use `api_open=True` instead of enable_api (since Gradio 3.40+)
    demo.launch(server_name="0.0.0.0", port=7860, api_open=True)