import os
import math
import time
import logging
from collections import deque
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import InferenceClient

# ── Logging ──────────────────────────────────────────────────────────────────
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# ── App setup ─────────────────────────────────────────────────────────────────
app = FastAPI(title="NB4170 LLM Proxy")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Binder + local notebooks need this
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── HF client (token lives only here, in a Space Secret) ──────────────────────
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN environment variable is not set")

client = InferenceClient(token=HF_TOKEN)

# ── Allowed models ─────────────────────────────────────────────────────────────
ALLOWED_MODELS = {
    "openai-community/gpt2":          "openai-community/gpt2",
    "llama-1b":                       "meta-llama/Llama-3.2-1B-Instruct",
    "meta-llama/Llama-3.2-1B-Instruct": "meta-llama/Llama-3.2-1B-Instruct",
    "llama-8b":                       "meta-llama/Llama-3.1-8B-Instruct",
    "meta-llama/Llama-3.1-8B-Instruct": "meta-llama/Llama-3.1-8B-Instruct",
}

# ── Simple rate limiter: max 60 requests per minute across all students ────────
request_times = deque()
RATE_LIMIT = 60
RATE_WINDOW = 60  # seconds

def check_rate_limit():
    now = time.time()
    while request_times and request_times[0] < now - RATE_WINDOW:
        request_times.popleft()
    if len(request_times) >= RATE_LIMIT:
        raise HTTPException(
            status_code=429,
            detail=f"Rate limit reached ({RATE_LIMIT} requests/min). Wait a moment and try again."
        )
    request_times.append(now)

# ── Request schema ─────────────────────────────────────────────────────────────
class GenerateRequest(BaseModel):
    system_prompt: str = "" # added system prompt
    prompt: str
    model: str = "gpt2"
    max_tokens: int = 50

# ── Health check ───────────────────────────────────────────────────────────────
@app.get("/")
def health():
    return {"status": "ok", "message": "NB4170 LLM Proxy is running"}

# ── Main endpoint ──────────────────────────────────────────────────────────────
@app.post("/generate")
def generate(req: GenerateRequest):
    # Rate limit
    check_rate_limit()

    # Validate model
    model_id = ALLOWED_MODELS.get(req.model)
    if not model_id:
        raise HTTPException(
            status_code=400,
            detail=f"Model '{req.model}' not allowed. Choose from: {list(ALLOWED_MODELS.keys())}"
        )

    # Clamp max_tokens to avoid runaway costs
    max_tokens = min(req.max_tokens, 1000)

    logger.info(f"Request: model={model_id}, max_tokens={max_tokens}, prompt_len={len(req.prompt)}")

    try:
        messages = []
        # add system prompt to messages if provided
        if req.system_prompt:
            messages.append({"role": "system", "content": req.system_prompt})
        messages.append({"role": "user", "content": req.prompt})
        
        response = client.chat.completions.create(
            model=model_id,
            messages=messages,
            max_tokens=max_tokens,
            logprobs=True,
            top_logprobs=1,
        )
    except Exception as e:
        logger.error(f"HF API error: {e}")
        raise HTTPException(status_code=502, detail=f"HF Inference API error: {str(e)}")

        # ── Post-process response ──────────────────────────────────────────────────
    try:
        answer = response.choices[0].message.content

        # Extract logprobs — same logic as your original notebook
        logprobs_content = response.choices[0].logprobs.content
        logprobs_dict = {x.token: x.logprob for x in logprobs_content}
        token_probs_dict = {token: math.exp(lp) for token, lp in logprobs_dict.items()}

    except Exception as e:
        logger.error(traceback.format_exc())
        raise HTTPException(status_code=500, detail=f"Post-processing error: {e}")


    return {
            "answer": answer,
            "logprobs": logprobs_dict,
            "token_probs": token_probs_dict,
            "logprobs_content": logprobs_content,
        }