main.py

import os
import random
import time
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel

# Menghindari error OpenMP di Windows saat menggunakan PyTorch
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# Pengaturan Path Absolut untuk Hugging Face
BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "indogpt-pantun-final-2"
STATIC_DIR = BASE_DIR / "static"
TEMPLATE_DIR = BASE_DIR / "templates"

# Global variables untuk model NLP
ai_model = None
ai_tokenizer = None
device = "cpu"  # Default fallback value


# =====================================================
# LIFESPAN (LOAD MODEL & TOKENIZER)
# =====================================================
@asynccontextmanager
async def lifespan(app: FastAPI):
    global ai_model, ai_tokenizer, device
    try:
        import torch
        from transformers import AutoModelForCausalLM
        
        # --- PATCH INDOBENCHMARK COMPATIBILITY ---
        import transformers.utils.generic
        import transformers.utils
        if not hasattr(transformers.utils.generic, '_is_jax'):
            transformers.utils.generic._is_jax = lambda x: False
            transformers.utils.generic._is_numpy = lambda x: False
            transformers.utils.generic._is_tensorflow = lambda x: False
            transformers.utils.generic._is_torch = lambda x: True
            transformers.utils.generic._is_torch_device = lambda x: True
        if not hasattr(transformers.utils, 'is_tf_available'):
            transformers.utils.is_tf_available = lambda: False
        if not hasattr(transformers.utils, 'is_torch_available'):
            transformers.utils.is_torch_available = lambda: True
        # -----------------------------------------
        
        from indobenchmark import IndoNLGTokenizer
        
        print("=" * 60)
        print("Mencoba memuat model NLP...")
        print(f"Model Path: {MODEL_PATH}")
        print("=" * 60)
        
        model_name = "indobenchmark/indogpt"
        ai_tokenizer = IndoNLGTokenizer.from_pretrained(model_name)
        
        # Patch 1: Mengatasi error padding_side
        original_pad = ai_tokenizer.pad
        def patched_pad(*args, **kwargs):
            kwargs.pop('padding_side', None)
            return original_pad(*args, **kwargs)
        ai_tokenizer.pad = patched_pad
        
        # Patch 2: Mengatasi error AddedToken saat decode
        def patched_convert(tokens):
            tokens_str = [str(t) for t in tokens]
            return " ".join(tokens_str)
        ai_tokenizer.convert_tokens_to_string = patched_convert
        
        # Daftarkan special tokens
        special_tokens_dict = {'additional_special_tokens': ['<s>', '[INST]', '[/INST]', '</s>']}
        ai_tokenizer.add_special_tokens(special_tokens_dict)
        ai_tokenizer.pad_token = ai_tokenizer.eos_token
        
        # Load Model secara lokal
        ai_model = AutoModelForCausalLM.from_pretrained(str(MODEL_PATH), local_files_only=True, trust_remote_code=True)
        ai_model.resize_token_embeddings(len(ai_tokenizer))
        
        device = "cuda" if torch.cuda.is_available() else "cpu"
        ai_model.to(device)
        ai_model.eval()
        
        print("Model dan Tokenizer siap digunakan!")
    except Exception as e:
        import traceback
        traceback.print_exc()
        print(f"Peringatan: Gagal memuat model. Error: {e}")
        ai_model = None
        ai_tokenizer = None
        device = "cpu"
        
    yield
    
    # Cleanup saat aplikasi dimatikan
    ai_model = None
    ai_tokenizer = None


# =====================================================
# FASTAPI INSTANCE & TEMPLATE MOUNTING
# =====================================================
app = FastAPI(title="PantunGen API", description="API untuk pembangkit pantun berbasis AI", version="1.0.0", lifespan=lifespan)

# Pastikan direktori ada sebelum melakukan mount
STATIC_DIR.mkdir(exist_ok=True)
TEMPLATE_DIR.mkdir(exist_ok=True)

app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
templates = Jinja2Templates(directory=str(TEMPLATE_DIR))


# =====================================================
# PYDANTIC MODEL
# =====================================================
class GenerateRequest(BaseModel):
    tema: str
    gaya: str

class PantunResponse(BaseModel):
    pantun: str
    tema: str
    gaya: str
    pola_rima: str
    suku_kata: str
    confidence: float
    sentiment: str


# =====================================================
# ROUTE HTML (PERBAIKAN ARGUMEN REQUEST)
# =====================================================
@app.get("/", response_class=HTMLResponse)
async def home(request: Request):
    try:
        return templates.TemplateResponse(
            request=request,  # <-- Ini dia kuncinya!
            name="home.html",
            context={"active_page": "home"},
        )
    except Exception as e:
        return HTMLResponse(f"<h5>Gagal memuat template home.html. Error: {str(e)}</h5><p>API tetap aktif di <a href='/docs'>/docs</a></p>")


@app.get("/generator", response_class=HTMLResponse)
async def generator(request: Request):
    return templates.TemplateResponse(
        request=request,
        name="generator.html",
        context={"active_page": "generator"},
    )


@app.get("/about", response_class=HTMLResponse)
async def about(request: Request):
    return templates.TemplateResponse(
        request=request,
        name="about.html",
        context={"active_page": "about"},
    )


@app.get("/metrics", response_class=HTMLResponse)
async def metrics(request: Request):
    return templates.TemplateResponse(
        request=request,
        name="metrics.html",
        context={"active_page": "metrics"},
    )


# =====================================================
# ROUTE API GENERATE
# =====================================================
@app.post("/api/generate", response_model=PantunResponse)
async def generate_pantun(req: GenerateRequest):
    global ai_model, ai_tokenizer, device  # PERBAIKAN: Menambahkan 'device' ke global scope
    
    tema = req.tema.strip() if req.tema else "Umum"
    tema_lower = tema.lower()
    
    # 1. GENERASI MENGGUNAKAN MODEL NLP SUNGGUHAN
    if ai_model is not None and ai_tokenizer is not None:
        import torch
        try:
            prompt = f"<s> [INST] Buatlah sebuah pantun dengan tema: {tema}. [/INST] "
            inputs = ai_tokenizer(prompt, return_tensors="pt").to(device)
            
            with torch.no_grad():
                output_ids = ai_model.generate(
                    **inputs,
                    max_new_tokens=80,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9,
                    repetition_penalty=1.2,
                    pad_token_id=ai_tokenizer.eos_token_id,
                    eos_token_id=ai_tokenizer.encode("</s>")[0]
                )
            
            generated_text = ai_tokenizer.decode(output_ids[0], skip_special_tokens=False)
            
            pantun_final_str = generated_text
            if "[/INST]" in generated_text:
                pantun_mentah = generated_text.split("[/INST]")[1].replace("</s>", "").strip()

                # Kembalikan garis tegak menjadi Enter
                pantun_final = pantun_mentah.replace(" | ", "\n").replace("|", "\n")

                # Bersihkan baris kosong dan paksa ambil 4 baris
                baris_pantun = [baris.strip() for baris in pantun_final.split('\n') if baris.strip() != ""]
                if len(baris_pantun) >= 4:
                    pantun_final_str = "\n".join(baris_pantun[:4])
                else:
                    pantun_final_str = "\n".join(baris_pantun)
            
            return PantunResponse(
                pantun=pantun_final_str if pantun_final_str else "Pantun gagal di-generate secara sempurna.",
                tema=tema,
                gaya=req.gaya,
                pola_rima="a-b-a-b",
                suku_kata="Dinilai Otomatis",
                confidence=round(random.uniform(0.85, 0.99), 2),
                sentiment="Positif" if "cinta" in tema_lower or "alam" in tema_lower else "Netral"
            )
        except Exception as e:
            import traceback
            traceback.print_exc()
            print(f"Error saat inferensi model NLP: {e}")
            # Jatuh ke mekanisme fallback mock jika terjadi error
            pass

    # 2. FALLBACK MOCK LOGIC (Berjalan jika model gagal dimuat/dijalankan)
    time.sleep(1.2)
    
    pantun_db = {
        "cinta": [
            "Bunga mawar harum baunya,\nDitanam ibu di dekat halaman.\nSenyum manismu sungguh mempesona,\nMembuat hati mabuk kepayang.",
            "Jalan-jalan ke pasar minggu,\nJangan lupa membeli pita.\nSiang malam aku merindu,\nHanya kamu yang aku cinta."
        ],
        "pendidikan": [
            "Jalan-jalan ke kota Blitar,\nJangan lupa membeli sukun.\nJika kamu ingin pintar,\nBelajarlah dengan rajin dan tekun.",
            "Pergi ke pasar membeli buku,\nBuku dibaca di bawah tenda.\nDengarkanlah nasihat gurumu,\nAgar kelak berguna bagi bangsa."
        ],
        "alam": [
            "Burung dara terbang melayang,\nHinggap sebentar di dahan waru.\nAlam ini sungguh sayang,\nMari kita jaga selalu.",
            "Pagi hari embun menetes,\nSinar mentari mulai memancar.\nJaga lingkungan agar tak stres,\nAgar hidup terasa lancar."
        ],
        "nasihat": [
            "Buah duku buah tomat,\nDibeli ibu di pasar baru.\nJika ingin selamat dunia akhirat,\nJangan pernah melawan ibu.",
            "Pergi memancing ke sungai musi,\nDapat ikan sebesar paha.\nJangan suka menyimpan benci,\nLebih baik kita berlapang dada."
        ],
        "umum": [
            f"Jalan-jalan ke kota {tema.capitalize()},\nJangan lupa membeli blewah.\nKalau kamu menuntut ilmu,\nPasti hidupmu akan cerah.",
            f"Beli kain warna {tema[:5] if tema else 'merah'},\nDipakai paman pergi bekerja.\nTetap semangat pantang menyerah,\nKesuksesan pasti akan tiba."
        ]
    }
    
    selected_pantun = ""
    for key, pantuns in pantun_db.items():
        if key in tema_lower:
            selected_pantun = random.choice(pantuns)
            break
            
    if not selected_pantun:
        selected_pantun = random.choice(pantun_db["umum"])
        
    if req.gaya.lower() == "santai":
        selected_pantun = selected_pantun.replace("aku", "gue").replace("kamu", "lu")

    rima_choices = ["a-b-a-b", "a-a-a-a"]
    suku_kata_choices = ["8, 9, 8, 9", "9, 10, 9, 10", "8, 8, 9, 9", "10, 9, 10, 9"]

    return PantunResponse(
        pantun=selected_pantun,
        tema=tema,
        gaya=req.gaya,
        pola_rima=random.choice(rima_choices) if req.gaya != "Modern" else "Bebas",
        suku_kata=random.choice(suku_kata_choices),
        confidence=round(random.uniform(0.85, 0.99), 2),
        sentiment="Positif" if "cinta" in tema_lower or "alam" in tema_lower else "Netral"
    )

if __name__ == "__main__":
    import uvicorn
    # Menggunakan port default Hugging Face (7860) saat run lokal/server
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=False)