import os, time, logging
from fastapi import FastAPI
from pydantic import BaseModel
import uvicorn
import httpx

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

HF_TOKEN = os.getenv("HF_TOKEN", "")
API_URL = "https://api-inference.huggingface.co/models/google/gemma-2-2b-it"

app = FastAPI()

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 2000
    temperature: float = 0.3

@app.post("/generate")
def generate(req: GenerateRequest):
    headers = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}
    
    # Add HF as inference provider header
    headers["x-use-cache"] = "false"
    
    for attempt in range(5):
        try:
            r = httpx.post(API_URL, json={
                "inputs": req.prompt,
                "parameters": {"max_new_tokens": req.max_tokens, "temperature": req.temperature, "return_full_text": False}
            }, headers=headers, timeout=120)
            
            logger.info(f"HF API status: {r.status_code}, attempt: {attempt+1}")
            
            if r.status_code == 200:
                result = r.json()
                if isinstance(result, list) and len(result) > 0:
                    text = result[0].get("generated_text", "")
                elif isinstance(result, dict):
                    text = result.get("generated_text", "")
                else:
                    text = str(result)
                if text:
                    return {"text": text, "model": "gemma-2b"}
            
            if r.status_code == 503 or r.status_code == 429:
                wait = min((attempt + 1) * 5, 25)
                logger.info(f"Waiting {wait}s...")
                time.sleep(wait)
                continue
                
        except Exception as e:
            logger.error(f"Error: {e}")
            time.sleep(3)
            continue
    
    return {"text": "Model warming up, try again in 30 seconds.", "model": "gemma-2b"}

@app.get("/health")
def health():
    return {"status": "ok", "model": "gemma-2b-it", "type": "hf-inference-api"}

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))