app.py

import os
import time
from typing import List, Optional, Any
from uuid import uuid4

import google.generativeai as genai
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

API_KEY = os.getenv("GEMINI_API_KEY", "").strip()
if not API_KEY:
    raise RuntimeError("Defina GEMINI_API_KEY no Hugging Face Space.")

genai.configure(api_key=API_KEY)

_model = None
_model_name = None

def _clean_model_name(name: str) -> str:
    name = (name or "").strip()
    return name[7:] if name.startswith("models/") else name

def _available_generate_models() -> List[str]:
    names = []
    try:
        for item in genai.list_models():
            methods = getattr(item, "supported_generation_methods", None) or []
            if "generateContent" in methods:
                name = _clean_model_name(getattr(item, "name", ""))
                if name:
                    names.append(name)
    except Exception:
        pass
    return names

def get_model():
    global _model, _model_name
    if _model is not None:
        return _model

    requested = _clean_model_name(os.getenv("GEMINI_MODEL", ""))
    available = _available_generate_models()

    if requested:
        chosen = requested
    else:
        preferred = [
            "gemini-2.5-flash",
            "gemini-2.0-flash",
            "gemini-flash-latest",
            "gemini-pro-latest",
        ]
        chosen = next((x for x in preferred if x in available), None)
        if not chosen and available:
            chosen = available[0]
        if not chosen:
            chosen = _clean_model_name(
                os.getenv("GEMINI_FALLBACK_MODEL", "gemini-2.0-flash")
            )

    _model_name = chosen
    _model = genai.GenerativeModel(chosen)
    return _model

def current_model_name() -> str:
    return _model_name or _clean_model_name(os.getenv("GEMINI_MODEL", "")) or "auto"

app = FastAPI(title="Gemini OpenAI-Compatible Endpoint", version="2.0.0")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)

class PromptRequest(BaseModel):
    prompt: str
    max_new_tokens: int = Field(default=300, ge=1, le=8192)
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    history: List[str] = Field(default_factory=list)

class GenerateSoalRequest(BaseModel):
    topic: str
    level: int = Field(default=1, ge=1, le=5)
    tipe_pertanyaan: Optional[str] = None

class ChatMessage(BaseModel):
    role: str
    content: Any

class ChatCompletionRequest(BaseModel):
    model: Optional[str] = None
    messages: List[ChatMessage]
    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
    max_tokens: int = Field(default=512, ge=1, le=8192)
    stream: bool = False

def normalize_content(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and "text" in item:
                parts.append(str(item["text"]))
            else:
                parts.append(str(item))
        return "\n".join(parts)
    return "" if content is None else str(content)

def chat_messages_to_gemini(messages: List[ChatMessage]):
    result = []
    system = []

    for msg in messages:
        role = (msg.role or "user").lower()
        text = normalize_content(msg.content)

        if role == "system":
            system.append(text)
            continue

        result.append({
            "role": "model" if role == "assistant" else "user",
            "parts": [{"text": text}],
        })

    if system:
        prefix = "INSTRUÇÕES DO SISTEMA:\n" + "\n\n".join(system) + "\n\n"
        if result and result[0]["role"] == "user":
            result[0]["parts"][0]["text"] = prefix + result[0]["parts"][0]["text"]
        else:
            result.insert(0, {"role": "user", "parts": [{"text": prefix}]})

    return result or [{"role": "user", "parts": [{"text": "Olá"}]}]

def extract_text(result) -> str:
    try:
        return result.text
    except Exception:
        return str(result)

def format_prompt_soal(topic, tipe=None, level=1):
    tipe_instrucao = f"Tipe soal solicitado: {tipe}." if tipe else "Tipo de questão livre."
    return f"""
Crie uma questão de prática de gramática inglesa sobre: "{topic}".
{tipe_instrucao}
Nível: {level}/5.
Responda SOMENTE com JSON válido com as chaves:
text_pertanyaan, tipe_pertanyaan, opsi, jawaban_benar, penjelasan.
""".strip()

@app.get("/")
async def root():
    return {
        "status": "online",
        "service": "Gemini endpoint",
        "model": current_model_name(),
        "endpoints": {
            "generate": "POST /generate",
            "chat": "POST /v1/chat/completions",
            "models": "GET /v1/models",
            "health": "GET /health",
        },
    }

@app.head("/")
async def root_head():
    return {}

@app.get("/health")
async def health():
    return {"status": "ok", "model": current_model_name()}

@app.get("/v1/models")
async def models():
    available = _available_generate_models()
    selected = current_model_name()
    if selected == "auto" and available:
        selected = available[0]
    return {
        "object": "list",
        "data": [{
            "id": selected,
            "object": "model",
            "created": int(time.time()),
            "owned_by": "google",
        }],
    }

@app.post("/generate")
async def generate_text(req: PromptRequest):
    try:
        messages = []
        for msg in req.history:
            if msg.startswith("Bot:") or msg.startswith("Assistant:"):
                role = "model"
                text = msg.split(":", 1)[1].strip() if ":" in msg else msg
            else:
                role = "user"
                text = msg.split(":", 1)[1].strip() if msg.startswith("User:") else msg
            messages.append({"role": role, "parts": [{"text": text}]})

        messages.append({"role": "user", "parts": [{"text": req.prompt}]})

        result = get_model().generate_content(
            messages,
            generation_config={
                "temperature": req.temperature,
                "max_output_tokens": req.max_new_tokens,
            },
        )
        return {"generated_text": extract_text(result), "model": current_model_name()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/v1/chat/completions")
async def chat_completions(req: ChatCompletionRequest):
    if req.stream:
        raise HTTPException(status_code=400, detail="stream=true ainda não é suportado.")

    try:
        result = get_model().generate_content(
            chat_messages_to_gemini(req.messages),
            generation_config={
                "temperature": req.temperature,
                "max_output_tokens": req.max_tokens,
            },
        )
        answer = extract_text(result)
        return {
            "id": "chatcmpl-" + uuid4().hex,
            "object": "chat.completion",
            "created": int(time.time()),
            "model": current_model_name(),
            "choices": [{
                "index": 0,
                "message": {"role": "assistant", "content": answer},
                "finish_reason": "stop",
            }],
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/generate-soal")
async def generate_soal(req: GenerateSoalRequest):
    try:
        import json
        prompt = format_prompt_soal(req.topic, req.tipe_pertanyaan, req.level)
        result = get_model().generate_content(
            prompt,
            generation_config={"temperature": 0.8, "max_output_tokens": 512},
        )
        text = extract_text(result).strip()

        if text.startswith("```"):
            lines = text.splitlines()[1:]
            if lines and lines[-1].strip().startswith("```"):
                lines = lines[:-1]
            text = "\n".join(lines).strip()
        if text.lower().startswith("json"):
            text = text[4:].lstrip()

        soal_data = json.loads(text)
        soal_data["id"] = str(uuid4())
        soal_data["topic_id"] = req.topic
        soal_data["level"] = req.level
        soal_data["is_ai_generated"] = True
        return {"soal": soal_data}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Gagal generate soal: {e}")