"""
✍️ CapGen AI — Caption Generator Indonesia
===========================================
SETUP (HANYA 2 LANGKAH):
------------------------
1. Buat token gratis di: https://huggingface.co/settings/tokens
   - Klik "New token"
   - Type pilih "Read" (cukup ini saja)
   - Copy token-nya

2. Di Space kamu, buka: Settings → Repository Secrets
   - Name:  HF_TOKEN
   - Value: paste token tadi
   - Klik Save

SELESAI! Restart Space dan langsung bisa dipakai.
"""

import gradio as gr
import requests
import os

# ══════════════════════════════════════════════════════
#  KONFIGURASI MODEL
# ══════════════════════════════════════════════════════

MODEL_NAME = "Qwen/Qwen2.5-3B-Instruct"
# MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
# MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
# MODEL_NAME = "google/gemma-2-2b-it"

HF_TOKEN = os.environ.get("HF_TOKEN", "")
CHAT_API = f"https://api-inference.huggingface.co/models/{MODEL_NAME}/v1/chat/completions"
TEXT_API = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"

# ══════════════════════════════════════════════════════
#  DATA
# ══════════════════════════════════════════════════════

PLATFORMS = {
    "Instagram":        "Maks 2200 karakter. Boleh panjang, emoji bebas, hashtag penting.",
    "TikTok":           "Maks 150 karakter. SANGAT pendek, catchy, hook di awal.",
    "Twitter/X":        "Maks 280 karakter. Padat, tajam, bisa jadi thread hook.",
    "LinkedIn":         "Maks 3000 karakter. Profesional tapi engaging, boleh storytelling.",
    "WhatsApp Status":  "Maks 100 karakter. Sangat singkat, to the point.",
}

STYLES = {
    "Aesthetic":    "Puitis, minimalis emoji, kata-kata indah, vibes kekinian",
    "Lucu":         "Meme-able, punchline di akhir, bahasa slang, relatable",
    "Jualan":       "Ada harga/CTA, urgency, benefit jelas, profesional tapi ramah",
    "Motivasi":     "Inspiratif, powerful words, emoji api/bintang, hook kuat",
    "Romantis":     "Manis, puitis, bisa sedikit dramatis, emoji hati",
    "Gaul/Jaksel":  "Bahasa slang Jakarta: literally, bestie, sksksk, bangeet",
    "Santai":       "Casual, kayak ngobrol sama temen, ga formal",
    "Formal":       "Bahasa baku, sopan, cocok buat instansi/organisasi",
}

LANGUAGES = {
    "Indonesia":       "Bahasa Indonesia standar yang baik dan benar",
    "Indonesia Gaul":  "Campuran Indonesia dan slang gaul kekinian",
    "Bahasa Jawa":     "Bahasa Jawa (ngoko/krama sesuai konteks)",
    "Bahasa Sunda":    "Bahasa Sunda yang gaul dan lucu",
    "English":         "Bahasa Inggris yang 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",
}

SYSTEM_PROMPT = """Kamu adalah CapGen AI, ahli pembuat caption media sosial untuk pengguna Indonesia.

ATURAN KETAT (WAJIB DIIKUTI):
1. HANYA keluarkan caption. TANPA kata "Berikut", "Tentu", "Baik", "Oke", atau penjelasan apapun.
2. Setiap caption dipisahkan EXACTLY dengan --- (tiga strip, tanpa spasi tambahan).
3. Setiap caption WAJIB sertakan 3-8 hashtag relevan di AKHIR caption.
4. Sesuaikan panjang dengan batasan platform yang dipilih.
5. Gaya bahasa HARUS konsisten sesuai pilihan user.
6. Jangan beri nomor urut (1., 2., 3. dst).
7. Jangan pakai tanda kutip atau formatting markdown di caption."""


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

def call_api(messages):
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }
    params = {"max_tokens": 1000, "temperature": 0.85, "top_p": 0.9}

    try:
        resp = requests.post(
            CHAT_API, headers=headers,
            json={"model": MODEL_NAME, "messages": messages, **params},
            timeout=60
        )
        if resp.status_code == 200:
            return resp.json()["choices"][0]["message"]["content"]
        if resp.status_code == 503:
            return "LOADING"
        if resp.status_code == 429:
            return "RATE_LIMIT"
    except requests.exceptions.Timeout:
        return "TIMEOUT"
    except Exception:
        pass

    try:
        prompt = f"{messages[0]['content']}\n\n{messages[1]['content']}"
        resp = requests.post(
            TEXT_API, headers=headers,
            json={
                "inputs": prompt,
                "parameters": {
                    "max_new_tokens": 1000,
                    "temperature": 0.85,
                    "top_p": 0.9,
                    "return_full_text": False
                }
            },
            timeout=60
        )
        if resp.status_code == 200:
            data = resp.json()
            if isinstance(data, list) and "generated_text" in data[0]:
                text = data[0]["generated_text"]
                if prompt in text:
                    text = text.replace(prompt, "").strip()
                return text
        if resp.status_code == 503:
            return "LOADING"
        if resp.status_code == 429:
            return "RATE_LIMIT"
    except requests.exceptions.Timeout:
        return "TIMEOUT"
    except Exception:
        pass

    return "ERROR"


def clean_text(text):
    if not text:
        return text
    unwanted = [
        "Baik, berikut", "Berikut adalah", "Berikut beberapa",
        "Tentu, berikut", "Oke, berikut", "Tentu saja,",
        "Berikut caption", "Berikut ini", "Ini dia",
        "Sure, here", "Here are", "Of course",
    ]
    for uw in unwanted:
        if text.strip().lower().startswith(uw.lower()):
            dash_idx = text.find("---")
            if dash_idx > 0:
                return text[dash_idx + 3:].strip()
            dot_idx = text.find(".", len(uw))
            if 0 < dot_idx < 100:
                text = text[dot_idx + 1:].strip()
            break
    return text


def parse_captions(raw_text, jumlah):
    text = clean_text(raw_text)
    parts = [p.strip() for p in text.split("---") if p.strip()]

    captions = []
    for p in parts:
        p = p.strip()
        if len(p) > 2 and p[0].isdigit() and p[1] in ".)":
            p = p[2:].strip()
        if len(p) > 15:
            captions.append(p)

    if not captions and len(text) > 15:
        captions = [text]

    while len(captions) < jumlah:
        captions.append("")
    return captions[:jumlah]


def generate_caption(konteks, platform, gaya, bahasa, jumlah):
    if not konteks or not konteks.strip():
        vals = ["⚠️ Tulis dulu deskripsi foto/konten kamu di kolom atas!"] + [""] * 4
        return tuple(gr.update(value=v, visible=True) for v in vals)

    if not HF_TOKEN:
        msg = (
            "⚠️ TOKEN BELUM DISET!\n\n"
            "Cara setup:\n"
            "1. Buat token di huggingface.co/settings/tokens\n"
            "2. Di Space ini: Settings → Repository Secrets\n"
            "3. Name: HF_TOKEN  Value: paste token\n"
            "4. Restart Space"
        )
        vals = [msg] + [""] * 4
        return tuple(gr.update(value=v, visible=True) for v in vals)

    user_msg = f"""PLATFORM: {platform}
{PLATFORMS[platform]}

GAYA: {gaya}
{STYLES[gaya]}

BAHASA: {bahasa}
{LANGUAGES[bahasa]}

KONTEKS FOTO/KONTEN:
{konteks.strip()}

BUAT TEPAT {jumlah} CAPTION YANG BERBEDA."""

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

    result = call_api(messages)

    error_map = {
        "LOADING":    "⏳ Model sedang loading... Tunggu 20-30 detik lalu coba lagi.",
        "RATE_LIMIT": "⚠️ Terlalu banyak request. Tunggu 1-2 menit lalu coba lagi.",
        "TIMEOUT":    "⏳ Request timeout. Coba lagi.",
        "ERROR":      "❌ Gagal menghubungi API. Coba lagi nanti.",
    }

    if result in error_map:
        vals = [error_map[result]] + [""] * 4
        return tuple(gr.update(value=v, visible=True) for v in vals)

    captions = parse_captions(result, jumlah)

    output = []
    for i in range(5):
        if i < jumlah:
            output.append(gr.update(value=captions[i], visible=True))
        else:
            output.append(gr.update(value="", visible=False))

    return tuple(output)


def update_visibility(jumlah):
    return tuple(gr.update(visible=(i < jumlah)) for i in range(5))


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


# ══════════════════════════════════════════════════════
#  UI
# ══════════════════════════════════════════════════════

CUSTOM_CSS = """
#btn-generate {
    font-size: 18px !important;
    font-weight: 600 !important;
    padding: 14px 32px !important;
    border-radius: 12px !important;
}
.caption-output textarea {
    font-size: 15px !important;
    line-height: 1.7 !important;
    border-radius: 10px !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;
}
"""

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 jago prompt · 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",
                )
                jumlah = gr.Dropdown(
                    choices=[1, 2, 3, 5],
                    value=3,
                    label="🔢 Jumlah",
                )

            btn_generate = gr.Button(
                "✨ Generate Caption",
                variant="primary",
                size="lg",
                elem_id="btn-generate",
            )

        with gr.Column(scale=1):
            gr.Markdown("### 📋 Hasil Caption")
            gr.Markdown(
                '<span style="font-size:12px;color:#999;">'
                'Klik tombol 📋 untuk copy tiap caption</span>'
            )

            caption_boxes = []
            caption_groups = []
            for i in range(5):
                with gr.Group(visible=(i < 3)) as group:
                    with gr.Row():
                        box = gr.Textbox(
                            label=f"Caption {i + 1}",
                            lines=4,
                            max_lines=8,
                            interactive=False,
                            elem_classes=["caption-output"],
                            scale=5,
                        )
                        copy_btn = gr.CopyButton(value=box, scale=1)
                    caption_boxes.append(box)
                    caption_groups.append(group)

            jumlah.change(
                fn=update_visibility,
                inputs=[jumlah],
                outputs=caption_groups,
            )

            btn_generate.click(
                fn=generate_caption,
                inputs=[konteks, platform, gaya, bahasa, jumlah],
                outputs=caption_boxes,
            )
            btn_generate.click(
                fn=update_visibility,
                inputs=[jumlah],
                outputs=caption_groups,
            )

    gr.Markdown(
        '<div class="footer-text">'
        "Dibuat dengan ❤️ untuk komunitas Indonesia · "
        "Powered by Hugging Face Inference API"
        "</div>"
    )

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