"""
✍️ CapGen AI — Caption Generator Indonesia (v2)
Versi improved: few-shot prompting + generate 1 per 1
"""

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"

print(f"⏳ Loading {MODEL_NAME}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True,
)
model.eval()
print("✅ Model siap!")

# ══════════════════════════════════════════════════════
#  CONTOH CAPTION PER GAYA (ini kunci kualitasnya)
# ══════════════════════════════════════════════════════

CONTOH = {
    "Aesthetic": {
        "Instagram": (
            "Contoh caption bagus:\n"
            '"Di antara decitan kerupuk dan aroma rempah, hidup rasanya lengkap 🌾✨\n'
            '#warungcita #hidupsimple #nungguinmatahari"\n'
            '"Senja kali ini bilang, kamu boleh lelah tapi jangan berhenti 🌅\n'
            '#senja #istirahatsejenak #catatanpetang"'
        ),
        "TikTok": (
            "Contoh caption bagus:\n"
            '"Senja ngomong kamu cantik hari ini 🧡"\n'
            '"Yang penting bahagia, sisanya urusan nanti 🌿"'
        ),
        "default": (
            "Contoh caption bagus:\n"
            '"Hal kecil yang bikin senyum tanpa alasan 🌿✨\n'
            '#simplejoy #catatanhariini"\n'
            '"Beberapa hal indah datang tanpa direncanakan 🌙\n'
            '#unexpected #grateful"'
        ),
    },
    "Lucu": {
        "Instagram": (
            "Contoh caption bagus:\n"
            '"Diet mulai Senin... Senin apa ya yang mana 🤔\n'
            '#dietgagal #makanyangpenting #realita"\n'
            '"Hidupku kayak WiFi kantoran, kadang nyambung kadang nggak 📶😅\n'
            '#hidupini #relatable #galau"'
        ),
        "default": (
            "Contoh caption bagus:\n"
            '"Udah usaha maksimal, hasilnya tetap minimal 🤡"\n'
            '"Mau kaya tapi males, gimana tuh 😭"'
        ),
    },
    "Jualan": {
        "Instagram": (
            "Contoh caption bagus:\n"
            '"PROMO TERBATAS! 🔥 Serum glow yang biasanya 150rb, sekarang CUMA 89rb!\n'
            'Stok terbatas, siapa cepat dia dapat 💕\n'
            'Order via DM / klik link di bio!\n'
            '#skincare #serumglow #promoskincare #diskon"\n'
            '"Best seller kami kembali hadir! ✨ Celana linin premium yang nyaman seharian.\n'
            'Harga normal 199rb → NOW 129rb aja!\n'
            'Tersisa 30 pcs saja, buruan!\n'
            '#celanalinen #fashionpria #diskon"'
        ),
        "TikTok": (
            "Contoh caption bagus:\n"
            '"Harga 89rb aja buat kulit glowing! Link di bio 🔗🔥"'
        ),
        "default": (
            "Contoh caption bagus:\n"
            '"Promo hari ini! Dari 150rb jadi 89rb aja 🔥\n'
            'Order langsung via DM ya!\n'
            '#promohariini #diskon"'
        ),
    },
    "Motivasi": {
        "default": (
            "Contoh caption bagus:\n"
            '"Proses itu berat, tapi hasilnya buat kamu bangun bangun pagi 🔥\n'
            '#motivasipagi #hustle #nevergiveup"\n'
            '"Kamu bukan lambat, kamu lagi build fondasi yang kuat 💪\n'
            '#proses #trusttheprocess #growth"'
        ),
    },
    "Romantis": {
        "default": (
            "Contoh caption bagus:\n"
            '"Kamu tuh kayak kopi pagi, tanpa kamu seharian rasanya hambar ☕💕\n'
            '#pacarsayang #romantis #cintalokal"\n'
            '"Senyum kamu itu autofokus, yang lain jadi blur semua 📸💗\n'
            '#jodoh #romance #lovestory"'
        ),
    },
    "Gaul/Jaksel": {
        "default": (
            "Contoh caption bagus:\n"
            '"Literally this is the best day ever bangeet bestie 🫶✨\n'
            '#bestie #jaksel #vibeskekinian"\n'
            '"Makanannya literally bikin crying aaah parah banget sih 🍜😭\n'
            '#foodie #makanenak #gaul"'
        ),
    },
    "Santai": {
        "default": (
            "Contoh caption bagus:\n"
            '"Hari ini ngga ada rencana spesial, tapi ngga apa-apa kok 🌿\n'
            '#chill #santai #flow"\n'
            '"Nongkrong, ngopi, ngobrol ga jelas. Tapi bahagia ☕😊\n'
            '#weekendvibes #hangout #qualitytime"'
        ),
    },
    "Formal": {
        "default": (
            "Contoh caption bagus:\n"
            '"Kami mengucapkan terima kasih atas kehadiran Bapak/Ibu dalam acara seminar nasional ini.\n'
            'Semoga ilmu yang diperoleh dapat bermanfaat.\n'
            '#seminar #profesional #diknas"\n'
            '"Peluncuran program terbaru kami telah berlangsung dengan baik.\n'
            'Kami berharap program ini dapat memberikan dampak positif bagi masyarakat.\n'
            '#programbaru #layananpublik"'
        ),
    },
}

# Fallback kalau gaya/platform ga ada di contoh
def get_contoh(gaya, platform):
    if gaya in CONTOH:
        if platform in CONTOH[gaya]:
            return CONTOH[gaya][platform]
        return CONTOH[gaya]["default"]
    return CONTOH["Aesthetic"]["default"]


PLATFORMS = {
    "Instagram":        "Maks 2200 karakter. Boleh panjang, emoji bebas, hashtag penting.",
    "TikTok":           "Maks 150 karakter. SANGAT pendek, catchy.",
    "Twitter/X":        "Maks 280 karakter. Padat dan tajam.",
    "LinkedIn":         "Maks 3000 karakter. Profesional tapi engaging.",
    "WhatsApp Status":  "Maks 100 karakter. Sangat singkat.",
}

STYLES = {
    "Aesthetic":    "Puitis, minimalis emoji, kata-kata indah",
    "Lucu":         "Punchline di akhir, bahasa slang, relatable",
    "Jualan":       "Ada harga/CTA, urgency, benefit jelas",
    "Motivasi":     "Inspiratif, powerful words, emoji api",
    "Romantis":     "Manis, puitis, emoji hati",
    "Gaul/Jaksel":  "Slang Jakarta: literally, bestie, bangeet",
    "Santai":       "Casual, kayak ngobrol sama temen",
    "Formal":       "Bahasa baku, sopan",
}

LANGUAGES = {
    "Indonesia":       "Bahasa Indonesia standar",
    "Indonesia Gaul":  "Slang gaul kekinian",
    "Bahasa Jawa":     "Bahasa Jawa (ngoko/krama)",
    "Bahasa Sunda":    "Bahasa Sunda gaul",
    "English":         "Bahasa Inggris natural",
    "Bilingual":       "Campuran Indonesia dan English",
}

PRESETS = {
    "🍜 Makanan":    "Foto makan nasi goreng spesial di warung pinggir jalan, porsi besar, harga cuma 15rb, ada telur dan kerupuk",
    "🌅 Sunset":     "Foto sunset di pantai Bali, langit oranye kemerahan, pegang minuman kelapa, rambut diterpa angin",
    "🛒 Jualan":     "Foto produk skincare baru, botol kaca bening, tekstur gel, harga normal 150rb promo jadi 89rb",
    "🏢 Kantor":     "Foto bareng tim kantor setelah selesai project besar, semua senang, pakai seragam kompak",
    "🐱 Kucing":     "Foto kucing tidur miring di bantal pink, lucu banget, bulunya lembut dan fluffy",
    "🛵 Jalan-jalan": "Foto naik motor melalui jalan Malioboro malam hari, lampu warna-warni, habis belanja batik",
    "🏋️ Gym":        "Foto habis workout di gym, badan berkeringat, baju sport, terlihat bugar dan proud",
    "📚 Belajar":    "Foto meja belajar yang rapi, ada buku dan kopi, lampu belajar warm light, vibes produktif",
}

# ══════════════════════════════════════════════════════
#  LOGIC
# ══════════════════════════════════════════════════════

def generate_text(messages):
    text = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=300,
            temperature=0.7,
            top_p=0.85,
            top_k=50,
            do_sample=True,
            repetition_penalty=1.2,
        )
    new_tokens = outputs[0][inputs["input_ids"].shape[-1]:]
    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()


def build_prompt(konteks, platform, gaya, bahasa):
    contoh = get_contoh(gaya, platform)
    return (
        f"Buat SATU caption untuk:\n\n"
        f"Platform: {platform} ({PLATFORMS[platform]})\n"
        f"Gaya: {gaya} ({STYLES[gaya]})\n"
        f"Bahasa: {bahasa} ({LANGUAGES[bahasa]})\n"
        f"Konteks: {konteks.strip()}\n\n"
        f"{contoh}\n\n"
        f"Sekarang buat SATU caption baru yang BERBEDA dari contoh di atas, "
        f"untuk konteks yang saya berikan. Langsung tulis caption-nya saja, "
        f"sertakan 3-6 hashtag di akhir."
    )


def clean_output(text):
    """Bersihkan output dari kalimat pembuka yang tidak perlu"""
    lines_to_remove = [
        "Baik, berikut captionnya:", "Berikut captionnya:",
        "Berikut:", "Tentu, berikut:", "Oke, berikut:",
        "Baik,", "Tentu,", "Oke,", "Here is", "Sure, here",
        "Berikut adalah caption yang kamu minta:",
        "Ini captionnya:", "Nih captionnya:",
    ]
    for line in lines_to_remove:
        if text.strip().startswith(line):
            text = text.strip()[len(line):].strip()
            if text.startswith(":"):
                text = text[1:].strip()
    return text.strip()


def generate_caption(konteks, platform, gaya, bahasa):
    if not konteks or not konteks.strip():
        return "⚠️ Tulis dulu deskripsi foto/konten kamu!"

    system_msg = (
        "Kamu adalah penulis caption media sosial profesional bahasa Indonesia. "
        "Kamu HANYA mengeluarkan caption. Tidak ada penjelasan, tidak ada pembuka, "
        "tidak ada kata 'Berikut'. Langsung tulis caption saja."
    )
    user_msg = build_prompt(konteks, platform, gaya, bahasa)

    messages = [
        {"role": "system", "content": system_msg},
        {"role": "user", "content": user_msg}
    ]

    try:
        result = generate_text(messages)
    except Exception as e:
        return f"❌ Error: {str(e)[:200]}"

    return clean_output(result)


def fill_preset(name):
    return PRESETS.get(name, "")


# ══════════════════════════════════════════════════════
#  UI — DIUBAH: Generate 1 caption, bisa klik berkali-kali
# ══════════════════════════════════════════════════════

CUSTOM_CSS = """
#btn-generate {
    font-size: 18px !important;
    font-weight: 600 !important;
    padding: 14px 32px !important;
    border-radius: 12px !important;
}
.output-box textarea {
    font-size: 16px !important;
    line-height: 1.8 !important;
    border-radius: 10px !important;
    min-height: 120px !important;
}
.preset-row button {
    min-width: fit-content !important;
    font-size: 13px !important;
}
.footer-text {
    text-align: center;
    color: #aaa;
    font-size: 12px;
    margin-top: 16px;
}
.hint-text {
    font-size: 13px;
    color: #888;
    margin-top: 4px;
}
"""

with gr.Blocks() as demo:

    gr.Markdown("""
    <div style="text-align:center; padding:8px 0 4px 0;">
        <h1 style="margin:0;">✍️ CapGen AI</h1>
        <p style="font-size:17px; color:#555; margin:4px 0 0 0;">
            Caption Generator Khusus Bahasa Indonesia
        </p>
        <p style="font-size:13px; color:#999; margin:2px 0 0 0;">
            Gak perlu login · Gak perlu token · Langsung kepakai
        </p>
    </div>
    """)

    with gr.Row(equal_height=False):

        with gr.Column(scale=1):

            gr.Markdown("### 📝 Ceritakan Foto/Konten Kamu")
            konteks = gr.Textbox(
                placeholder=(
                    "Contoh: Foto makan nasi goreng di warung pinggir jalan, "
                    "porsi besar, harga 15rb, ada telur mata sapi..."
                ),
                lines=4,
                max_lines=8,
                show_label=False,
            )

            gr.Markdown("### ⚡ Contoh Cepat (klik = langsung isi)")
            preset_names = list(PRESETS.keys())
            for row_start in range(0, len(preset_names), 4):
                with gr.Row(elem_classes=["preset-row"]):
                    for name in preset_names[row_start:row_start + 4]:
                        gr.Button(name, size="sm", variant="secondary").click(
                            fn=lambda n=name: fill_preset(n),
                            outputs=konteks,
                        )

            gr.Markdown("### ⚙️ Pengaturan")
            with gr.Row():
                platform = gr.Dropdown(
                    choices=list(PLATFORMS.keys()),
                    value="Instagram",
                    label="📱 Platform",
                )
                gaya = gr.Dropdown(
                    choices=list(STYLES.keys()),
                    value="Aesthetic",
                    label="🎨 Gaya",
                )

            with gr.Row():
                bahasa = gr.Dropdown(
                    choices=list(LANGUAGES.keys()),
                    value="Indonesia",
                    label="🌐 Bahasa",
                )

            btn_generate = gr.Button(
                "✨ Generate Caption",
                variant="primary",
                size="lg",
                elem_id="btn-generate",
            )
            gr.Markdown(
                '<p class="hint-text">💡 Klik berkali-kali untuk hasil berbeda-beda</p>'
            )

        with gr.Column(scale=1):
            gr.Markdown("### 📋 Caption Kamu")

            output_caption = gr.Textbox(
                label="",
                lines=5,
                max_lines=10,
                interactive=False,
                show_label=False,
                placeholder="Hasil caption akan muncul di sini...",
                elem_classes=["output-box"],
            )

            gr.Markdown(
                '<p class="hint-text">Blok teks di atas lalu Ctrl+C / Cmd+C untuk copy</p>'
            )

            with gr.Accordion("📜 Riwayat Caption", open=False):
                history = gr.Textbox(
                    label="",
                    lines=12,
                    interactive=False,
                    show_label=False,
                    placeholder="Caption-caption sebelumnya akan muncul di sini...",
                )

    # State untuk riwayat
    history_state = gr.State("")

    def generate_with_history(konteks, platform, gaya, bahasa, prev_history):
        caption = generate_caption(konteks, platform, gaya, bahasa)

        if not caption.startswith("⚠️") and not caption.startswith("❌"):
            separator = "\n" + "─" * 40 + "\n"
            new_history = f"[{platform} · {gaya} · {bahasa}]\n{caption}{separator}{prev_history}"
            return caption, new_history
        return caption, prev_history

    btn_generate.click(
        fn=generate_with_history,
        inputs=[konteks, platform, gaya, bahasa, history_state],
        outputs=[output_caption, history_state],
    )
    btn_generate.click(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[history],
    )

    gr.Markdown(
        '<div class="footer-text">'
        "Dibuat dengan ❤️ untuk komunitas Indonesia · "
        f"Model: {MODEL_NAME}"
        "</div>"
    )

demo.launch(
    theme=gr.themes.Soft(
        primary_hue="orange",
        secondary_hue="amber",
        neutral_hue="stone",
    ),
    css=CUSTOM_CSS,
)