main.py

import os
import sqlite3
import datetime
import time
import base64
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from gradio_client import Client

app = FastAPI()

# --- CORS (Agar Web/HP tidak memblokir) ---
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], 
    allow_credentials=False, 
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- CONFIG ---
TARGET_SPACE = "black-forest-labs/FLUX.1-dev"
DB_PATH = "/app/data/users.db"

# --- TOKEN LIST (Smart Bypass) ---
_prefix = "hf_" 
RAW_TOKENS = [
    "PiRCDDtPcPFMLWkTkVaZmzoleHOunXnLIA", "BHvZXGICstaktSwycmwNmzHGrTNmKxnlRZ",
    "ZdgawyTPzXIpwhnRYIteUKSMsWnEDtGKtM", "nMiFYAFsINxAJWPwiCQlaunmdgmrcxKoaT",
    "PccpUIbTckCiafwErDLkRlsvqhgtfZaBHL", "faGyXBPfBkaHXDMUSJtxEggonhhZbomFIz",
    "SndsPaRWsevDXCgZcSjTUlBYUJqOkSfFmn", "CqobFdUpeVCeuhUaiuXwvdczBUmoUHXRGa",
    "JKCQYUhhHPPkpucegqkNSyureLdXpmeXRF", "tBYfslUwHNiNMufzwAYIlrDVovEWmOQulC",
    "LKLdrdUxyUyKODSUthmqHXqDMfHrQueera", "ivSBboJYQVcifWkCNcOTOnxUQrZOtOglnU"
]
HF_TOKENS = [_prefix + t for t in RAW_TOKENS]
BAD_TOKENS = set()

# --- DATABASE ---
def init_db():
    conn = sqlite3.connect(DB_PATH)
    conn.execute('CREATE TABLE IF NOT EXISTS users (username TEXT PRIMARY KEY, password TEXT, credits INTEGER, last_reset_date TEXT)')
    conn.execute('CREATE TABLE IF NOT EXISTS gallery (id INTEGER PRIMARY KEY, username TEXT, image_base64 TEXT, prompt TEXT, created_at TEXT)')
    conn.commit()
    conn.close()

init_db()

# --- MODELS ---
class UserAuth(BaseModel):
    username: str
    password: str

class GenerateReq(BaseModel):
    username: str
    password: str
    prompt: str
    width: int = 1024
    height: int = 1024
    guidance_scale: float = 3.5
    num_inference_steps: int = 28

# --- ENDPOINTS ---
@app.get("/")
def home():
    return {"status": "Online", "mode": "Docker Force Update"}

@app.get("/ping")
def ping():
    return {"status": "alive", "time": time.time()}

@app.post("/auth/register")
def register(user: UserAuth):
    with sqlite3.connect(DB_PATH) as conn:
        try:
            today = datetime.date.today().isoformat()
            conn.execute("INSERT INTO users VALUES (?, ?, 10, ?)", (user.username, user.password, today))
            return {"msg": "OK", "credits": 10}
        except:
            raise HTTPException(400, "Username taken")

@app.post("/auth/login")
def login_endpoint(user: UserAuth):
    with sqlite3.connect(DB_PATH) as conn:
        cur = conn.execute("SELECT credits, last_reset_date FROM users WHERE username=? AND password=?", (user.username, user.password))
        res = cur.fetchone()
        if not res: raise HTTPException(401, "Login gagal")
        creds, last_date = res
        today = datetime.date.today().isoformat()
        if last_date != today:
            conn.execute("UPDATE users SET credits=10, last_reset_date=? WHERE username=?", (today, user.username))
            creds = 10
        return {"msg": "OK", "credits": creds}

@app.post("/my-gallery")
def my_gallery(user: UserAuth):
    with sqlite3.connect(DB_PATH) as conn:
        cur = conn.execute("SELECT image_base64 FROM gallery WHERE username=? ORDER BY id DESC LIMIT 20", (user.username,))
        return {"gallery": [{"image": r[0]} for r in cur.fetchall()]}

@app.post("/generate")
def generate(req: GenerateReq):
    global BAD_TOKENS
    
    # 1. Cek User & Kredit
    with sqlite3.connect(DB_PATH) as conn:
        cur = conn.execute("SELECT credits FROM users WHERE username=? AND password=?", (req.username, req.password))
        res = cur.fetchone()
        if not res: raise HTTPException(401, "Auth gagal")
        if res[0] <= 0: raise HTTPException(402, "Kredit habis")
        conn.execute("UPDATE users SET credits = credits - 1 WHERE username=?", (req.username,))

    # 2. Smart Loop Token
    success = False
    result_b64 = None
    last_err = ""

    # Reset jika semua token dianggap mati
    if len(BAD_TOKENS) >= len(HF_TOKENS):
        BAD_TOKENS.clear()

    for token in HF_TOKENS:
        if token in BAD_TOKENS: continue
        
        try:
            print(f"Mencoba Token: ...{token[-4:]}")
            # Dockerfile menjamin library ini support hf_token
            client = Client(TARGET_SPACE, hf_token=token)
            
            res = client.predict(
                prompt=req.prompt,
                seed=0,
                randomize_seed=True,
                width=req.width,
                height=req.height,
                guidance_scale=req.guidance_scale,
                num_inference_steps=req.num_inference_steps,
                api_name="/infer"
            )
            
            with open(res[0], "rb") as f:
                result_b64 = base64.b64encode(f.read()).decode('utf-8')
            success = True
            break
            
        except Exception as e:
            print(f"Gagal ({token[-4:]}): {e}")
            last_err = str(e)
            BAD_TOKENS.add(token)

    if success:
        with sqlite3.connect(DB_PATH) as conn:
            conn.execute("INSERT INTO gallery (username, image_base64, prompt, created_at) VALUES (?, ?, ?, ?)", 
                         (req.username, result_b64, req.prompt, datetime.datetime.now().isoformat()))
            rem = conn.execute("SELECT credits FROM users WHERE username=?", (req.username,)).fetchone()[0]
        return {"status": "success", "image_base64": result_b64, "remaining_credits": rem}
    else:
        # Refund Kredit
        with sqlite3.connect(DB_PATH) as conn:
            conn.execute("UPDATE users SET credits = credits + 1 WHERE username=?", (req.username,))
        
        # Pesan Error Detail
        if "429" in last_err or "quota" in last_err.lower():
            raise HTTPException(503, "Semua Token Limit (Quota Exceeded). Coba lagi nanti.")
        else:
            raise HTTPException(500, f"Server Error: {last_err}")