import json, os, time, uuid
from collections import defaultdict, deque

import torch
from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import JSONResponse
from huggingface_hub import HfApi
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from pydantic import BaseModel

# ---------- Config ----------
MODEL_ID = "wolethereader/STORM-OS-MT-3B-BIDIRECTIONAL"
ORG_NAME = "wolethereader"
EN = "eng_Latn"
LANG_CODES = {"yo": "yor_Latn", "ha": "hau_Latn", "ig": "ibo_Latn", "pcm": "pcm_Latn"}
VALID_LANGS = set(LANG_CODES.keys())
MAX_TEXT_CHARS = 2000

app = FastAPI(title="STORM-OS Bidirectional MT API")

def log_event(event, **fields):
    print(json.dumps({"event": event, "ts": time.time(), **fields}))

# ---------- Auth: HF token AND must belong to the org ----------
security = HTTPBearer()
hf_api = HfApi()
_token_cache = {}
TOKEN_CACHE_TTL = 300
EXTERNAL_ACCESS_TOKEN = os.environ.get("EXTERNAL_ACCESS_TOKEN")

async def verify_org_token(creds: HTTPAuthorizationCredentials = Depends(security)):
    token = creds.credentials

    if EXTERNAL_ACCESS_TOKEN and token == EXTERNAL_ACCESS_TOKEN:
        log_event("auth_external_token_used")
        return "external-collaborator"

    now = time.time()
    cached = _token_cache.get(token)
    if cached and cached[1] > now:
        return cached[0]
    try:
        info = hf_api.whoami(token=token)
    except Exception:
        log_event("auth_failed_invalid_token")
        raise HTTPException(status_code=401, detail="Invalid or expired Hugging Face token")
    username = info.get("name", "unknown")
    user_orgs = [o.get("name") for o in info.get("orgs", [])]
    if ORG_NAME not in user_orgs:
        log_event("auth_failed_not_org_member", user=username, orgs=user_orgs)
        raise HTTPException(status_code=403, detail=f"Token does not belong to a member of '{ORG_NAME}'")
    _token_cache[token] = (username, now + TOKEN_CACHE_TTL)
    return username

# ---------- Rate limiting ----------
_rate_state = defaultdict(deque)
RATE_LIMIT_PER_MIN = 30

def check_rate_limit(username: str):
    now = time.time()
    q = _rate_state[username]
    while q and q[0] < now - 60:
        q.popleft()
    if len(q) >= RATE_LIMIT_PER_MIN:
        raise HTTPException(status_code=429, detail="Rate limit exceeded, try again shortly")
    q.append(now)

# ---------- Model ----------
tokenizer = None
model = None
HF_TOKEN = os.environ.get("HF_TOKEN")  # repo is private

@app.on_event("startup")
async def startup():
    global tokenizer, model
    log_event("loading_model", model=MODEL_ID)
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16, token=HF_TOKEN)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    model.eval()
    log_event("model_loaded_ok", device=device)

class TranslateRequest(BaseModel):
    text: str
    direction: str  # "forward" (local -> English) or "reverse" (English -> local)
    lang: str        # the local language code, regardless of direction
    max_new_tokens: int = 128

@app.get("/")
def root():
    return {"status": "ok", "languages": sorted(VALID_LANGS), "directions": ["forward", "reverse"], "engine": MODEL_ID}

@app.get("/health")
def health():
    return {"status": "ok" if model is not None else "loading"}

@app.post("/translate")
async def translate(req: TranslateRequest, username: str = Depends(verify_org_token)):
    check_rate_limit(username)

    if req.lang not in VALID_LANGS:
        raise HTTPException(status_code=400, detail=f"lang must be one of {sorted(VALID_LANGS)}")
    if req.direction not in ("forward", "reverse"):
        raise HTTPException(status_code=400, detail="direction must be 'forward' or 'reverse'")
    if not req.text or not req.text.strip():
        raise HTTPException(status_code=400, detail="text must not be empty")
    if len(req.text) > MAX_TEXT_CHARS:
        raise HTTPException(status_code=400, detail=f"text exceeds {MAX_TEXT_CHARS} character limit")

    request_id = str(uuid.uuid4())
    start = time.time()

    if req.direction == "forward":
        src_lang, tgt_lang = LANG_CODES[req.lang], EN
    else:
        src_lang, tgt_lang = EN, LANG_CODES[req.lang]

    tokenizer.src_lang = src_lang
    inputs = tokenizer(req.text, return_tensors="pt", truncation=True, max_length=128).to(model.device)
    tgt_id = tokenizer.convert_tokens_to_ids(tgt_lang)

    with torch.no_grad():
        out = model.generate(**inputs, forced_bos_token_id=tgt_id, max_new_tokens=req.max_new_tokens, max_length=None)

    translated = tokenizer.decode(out[0], skip_special_tokens=True)
    elapsed_s = round(time.time() - start, 2)

    log_event("translate_ok", request_id=request_id, user=username, direction=req.direction, lang=req.lang, elapsed_s=elapsed_s)

    return {
        "request_id": request_id,
        "direction": req.direction,
        "lang": req.lang,
        "translated_text": translated,
        "elapsed_s": elapsed_s,
    }

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    log_event("request_error", path=str(request.url.path), status_code=exc.status_code, detail=exc.detail)
    return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})