"""
deploy/tgi_space/app.py — Smart inference proxy for persona generation.

Priority chain (first available wins):
  1. GROQ_API_KEY → Groq Cloud (fast, free: 14,400 req/day with llama-3.1-8b-instant)
  2. GEMINI_API_KEY → Gemini Flash (generous free: 1,500 req/day, 1M tok/day)
  3. Local CPU → transformers (slow fallback, only for smoke-testing)

To activate a fast provider, set the env var in the HF Space settings:
  - Groq:   GROQ_API_KEY  = gsk_...  (free at https://console.groq.com)
  - Gemini: GEMINI_API_KEY = AIza...  (free at https://aistudio.google.com)
"""
from __future__ import annotations

import asyncio
import json
import os
import time
import uuid
from typing import Any, Dict, List, Optional

import httpx
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel

app = FastAPI(title="Persona Inference Proxy")

# ---------------------------------------------------------------------------
# Config (env vars)
# ---------------------------------------------------------------------------

MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-0.5B-Instruct")
PORT = int(os.environ.get("PORT", 7860))
MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", 600))

GROQ_API_KEY = os.environ.get("GROQ_API_KEY", "")
GROQ_MODEL = os.environ.get("GROQ_MODEL", "llama-3.1-8b-instant")

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
GEMINI_MODEL = os.environ.get("GEMINI_MODEL", "gemini-1.5-flash")

# Local model (loaded lazily — only if no fast provider)
_tokenizer = None
_model = None
_local_loaded = False


def _infer_mode() -> str:
    if GROQ_API_KEY:
        return "groq"
    if GEMINI_API_KEY:
        return "gemini"
    return "local-cpu"


def _active_model() -> str:
    mode = _infer_mode()
    if mode == "groq":
        return GROQ_MODEL
    if mode == "gemini":
        return GEMINI_MODEL
    return MODEL_ID


# ---------------------------------------------------------------------------
# Request/response schemas (OpenAI-compatible)
# ---------------------------------------------------------------------------

class Message(BaseModel):
    role: str
    content: str


class ChatRequest(BaseModel):
    model: str = MODEL_ID
    messages: List[Message]
    temperature: float = 0.9
    max_tokens: int = MAX_NEW_TOKENS
    top_p: float = 0.95


# ---------------------------------------------------------------------------
# Provider implementations
# ---------------------------------------------------------------------------

async def _call_groq(req: ChatRequest) -> dict:
    payload = {
        "model": GROQ_MODEL,
        "messages": [{"role": m.role, "content": m.content} for m in req.messages],
        "max_tokens": min(req.max_tokens, MAX_NEW_TOKENS),
        "temperature": req.temperature,
        "top_p": req.top_p,
    }
    async with httpx.AsyncClient(timeout=90.0) as client:
        r = await client.post(
            "https://api.groq.com/openai/v1/chat/completions",
            json=payload,
            headers={
                "Authorization": f"Bearer {GROQ_API_KEY}",
                "Content-Type": "application/json",
            },
        )
    if r.status_code != 200:
        raise HTTPException(status_code=r.status_code, detail=f"Groq error: {r.text[:200]}")
    return r.json()


async def _call_gemini(req: ChatRequest) -> dict:
    """Call Gemini via its OpenAI-compatible endpoint."""
    payload = {
        "model": GEMINI_MODEL,
        "messages": [{"role": m.role, "content": m.content} for m in req.messages],
        "max_tokens": min(req.max_tokens, MAX_NEW_TOKENS),
        "temperature": req.temperature,
    }
    async with httpx.AsyncClient(timeout=90.0) as client:
        r = await client.post(
            "https://generativelanguage.googleapis.com/v1beta/openai/chat/completions",
            json=payload,
            headers={
                "Authorization": f"Bearer {GEMINI_API_KEY}",
                "Content-Type": "application/json",
            },
        )
    if r.status_code != 200:
        raise HTTPException(status_code=r.status_code, detail=f"Gemini error: {r.text[:200]}")
    return r.json()


def _load_local_model():
    global _tokenizer, _model, _local_loaded
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    print(f"Loading local model {MODEL_ID} on CPU ...")
    _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    _model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        torch_dtype=torch.float32,
        device_map="auto",
        trust_remote_code=True,
    )
    _model.eval()
    _local_loaded = True
    print(f"Local model loaded: {MODEL_ID}")


async def _call_local(req: ChatRequest) -> dict:
    global _local_loaded
    if not _local_loaded:
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, _load_local_model)

    import torch

    msgs = [{"role": m.role, "content": m.content} for m in req.messages]

    def _generate():
        text = _tokenizer.apply_chat_template(
            msgs, tokenize=False, add_generation_prompt=True
        )
        inputs = _tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = _model.generate(
                **inputs,
                max_new_tokens=min(req.max_tokens, MAX_NEW_TOKENS),
                temperature=max(req.temperature, 0.01),
                top_p=req.top_p,
                do_sample=req.temperature > 0,
                pad_token_id=_tokenizer.eos_token_id,
            )
        new_tokens = outputs[0][inputs.input_ids.shape[1]:]
        return _tokenizer.decode(new_tokens, skip_special_tokens=True)

    loop = asyncio.get_event_loop()
    content = await loop.run_in_executor(None, _generate)

    return {
        "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
        "object": "chat.completion",
        "created": int(time.time()),
        "model": MODEL_ID,
        "choices": [{"index": 0, "message": {"role": "assistant", "content": content}, "finish_reason": "stop"}],
        "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
    }


# ---------------------------------------------------------------------------
# Startup: eagerly load local model only if no fast provider configured
# ---------------------------------------------------------------------------

@app.on_event("startup")
async def startup():
    mode = _infer_mode()
    print(f"Inference mode: {mode}  active_model: {_active_model()}")
    if mode == "local-cpu":
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, _load_local_model)


# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------

@app.get("/health")
async def health():
    mode = _infer_mode()
    loaded = True if mode != "local-cpu" else _local_loaded
    return {
        "status": "ok",
        "mode": mode,
        "model": _active_model(),
        "loaded": loaded,
    }


@app.post("/v1/chat/completions")
async def chat_completions(req: ChatRequest):
    mode = _infer_mode()
    if mode == "groq":
        return await _call_groq(req)
    if mode == "gemini":
        return await _call_gemini(req)
    return await _call_local(req)


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=PORT)