# /// script
# dependencies = ["gradio==6.30.0", "spaces==0.51.3", "torch", "transformers==5.20.1"]
# ///
import os
import sys
from pathlib import Path

import gradio as gr
import spaces
import torch
from fastapi import HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from transformers import AutoModelForSequenceClassification, AutoTokenizer

ROOT = Path(__file__).resolve().parent
MODEL = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("PANCODE_MODEL", "tiagozip/pancode")
LOCAL = Path(MODEL).resolve() if Path(MODEL).is_dir() else None
CHUNK_CHARS = 6000
MAX_CHARS = 40000
MAX_CHUNKS = 12

dev = "cuda" if os.environ.get("SPACES_ZERO_GPU") or torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForSequenceClassification.from_pretrained(MODEL, dtype=torch.float32).to(dev).eval()
CONFIG = {
    "model": "local" if LOCAL else MODEL,
    "local": bool(LOCAL),
    "heads": model.config.pancode_heads,
    "thresholds": model.config.pancode_thresholds,
    "eval": model.config.pancode_eval,
}
PAGE = (ROOT / "index.html").read_text()
ASSETS = {p.name for p in (ROOT / "assets").glob("*.woff2")}

app = gr.Server()


@app.api(name="score")
@spaces.GPU(duration=20)
def score(chunks: list[str]) -> list[list[float]]:
    if not 0 < len(chunks) <= MAX_CHUNKS or max(map(len, chunks)) > CHUNK_CHARS or sum(map(len, chunks)) > MAX_CHARS:
        raise gr.Error(f"Send at most {MAX_CHARS:,} characters, in chunks of up to {CHUNK_CHARS:,}.")
    enc = tok(chunks, truncation=True, max_length=CONFIG["thresholds"]["max_len"], padding=True, return_tensors="pt").to(dev)
    with torch.no_grad():
        return model(**enc).logits.float().cpu().tolist()


@app.get("/", response_class=HTMLResponse)
def page():
    return PAGE


@app.get("/config.json")
def config():
    return CONFIG


@app.get("/assets/{name}")
def asset(name: str):
    if name not in ASSETS:
        raise HTTPException(404)
    return FileResponse(ROOT / "assets" / name)


@app.get("/models/local/{path:path}")
def model_file(path: str):
    f = LOCAL and (LOCAL / path).resolve()
    if not f or not f.is_relative_to(LOCAL) or not f.is_file():
        raise HTTPException(404)
    return FileResponse(f)


app.launch()