# server.py
import os
import io
import json
import time
import random
from typing import Any, Dict, List, Optional
from fastapi import FastAPI, Query, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, FileResponse, PlainTextResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from datetime import datetime, timezone
from huggingface_hub import HfApi, upload_file, hf_hub_download, create_repo
from huggingface_hub.utils import HfHubHTTPError

# === Config ===
REPO_ID = os.environ.get("STORE_REPO", "YanAdjeNole/asset-annotator-store")
HF_TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=HF_TOKEN)

app = FastAPI()

# CORS（同域访问其实不需要，但保留更宽松）
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ---------- Hub helpers ----------
_repo_ready = False

def ensure_repo():
    # create_repo is idempotent but still an API call; only do it once per process.
    global _repo_ready
    if _repo_ready:
        return
    try:
        create_repo(REPO_ID, repo_type="dataset", exist_ok=True, token=HF_TOKEN)
        _repo_ready = True
    except Exception:
        pass

def upload_json(path_in_repo: str, obj: Any, message: str, max_retries: int = 4):
    # HF dataset repos are git-backed and rate-limit commits per repo. Under
    # concurrent load the commit endpoint returns 429 ("concurrency queue") or
    # transient 500s, so retry with exponential backoff + jitter.
    ensure_repo()
    data = json.dumps(obj, ensure_ascii=False, indent=2).encode("utf-8")
    last_err: Optional[Exception] = None
    for attempt in range(max_retries):
        try:
            upload_file(
                path_or_fileobj=io.BytesIO(data),
                path_in_repo=path_in_repo,
                repo_id=REPO_ID,
                repo_type="dataset",
                token=HF_TOKEN,
                commit_message=message,
            )
            return
        except HfHubHTTPError as e:
            last_err = e
            status = getattr(getattr(e, "response", None), "status_code", None)
            if status in (429, 500, 502, 503) and attempt < max_retries - 1:
                time.sleep(min(8.0, 2 ** attempt) + random.random())
                continue
            raise
    if last_err:
        raise last_err

def download_json(path_in_repo: str) -> Optional[Any]:
    try:
        local = hf_hub_download(repo_id=REPO_ID, repo_type="dataset", filename=path_in_repo, token=HF_TOKEN)
        with open(local, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None

# ---------- API schemas ----------
class Dataset(BaseModel):
    dataset_id: str
    name: str
    data: Dict[str, List[Dict[str, Any]]]
    assets: List[str]
    dates: List[str]

class Annotation(BaseModel):
    dataset_id: str
    user_id: str
    selections: List[Dict[str, Any]]
    step: int
    window_len: int

# ---------- API routes ----------
@app.get("/api/health")
def health():
    dist_exists = os.path.exists("dist/index.html")
    return {"ok": True, "repo": REPO_ID, "dist": dist_exists}

@app.post("/api/dataset/upsert")
def upsert_dataset(ds: Dataset):
    payload = ds.dict()
    payload["id"] = ds.dataset_id  # 兼容前端
    payload["saved_at"] = datetime.now(timezone.utc).isoformat()
    upload_json("datasets/latest.json", payload, f"upsert dataset {ds.dataset_id}")
    return {"ok": True}


@app.get("/api/dataset/latest")
def get_latest():
    obj = download_json("datasets/latest.json")
    if obj is None:
        raise HTTPException(status_code=404, detail="No dataset yet.")
    return obj

@app.get("/api/annotation/get")
def get_annotation(dataset_id: str = Query(...), user_id: str = Query(...)):
    obj = download_json(f"annotations/{dataset_id}/{user_id}.json")
    if obj is None:
        return {"user_id": user_id, "dataset_id": dataset_id, "selections": [], "step": 1, "window_len": 7}
    return obj

@app.post("/api/annotation/upsert")
def upsert_annotation(ann: Annotation):
    payload = ann.dict()
    payload["updated_at"] = datetime.now(timezone.utc).isoformat()
    upload_json(f"annotations/{ann.dataset_id}/{ann.user_id}.json", payload, f"upsert annotation {ann.dataset_id}/{ann.user_id}")
    return {"ok": True}

# ---------- Static mount + SPA fallback ----------
# index.html 不能被浏览器缓存，否则用户会一直加载旧的 hashed JS bundle。
# hashed 的 /assets/*.js 本身可长期缓存（文件名带 hash，内容变了名字就变）。
NO_CACHE_HEADERS = {"Cache-Control": "no-cache, no-store, must-revalidate"}

# 1) 如果 dist 存在，挂载静态资源
if os.path.isdir("dist"):
    app.mount("/assets", StaticFiles(directory="dist/assets", html=False), name="assets")

# 2) 根路径：返回 index.html 或错误提示
@app.get("/")
def index():
    if os.path.exists("dist/index.html"):
        return FileResponse("dist/index.html", headers=NO_CACHE_HEADERS)
    return PlainTextResponse(
        "Frontend not built or not found. Missing dist/index.html.\n"
        "Check Dockerfile build stage and ensure '@vitejs/plugin-react' is installed.",
        status_code=500,
    )

# 3) 其它所有路径 → SPA fallback 到 index.html（不覆盖 /api/*）
@app.get("/{full_path:path}")
def spa(full_path: str):
    if full_path.startswith("api/"):
        raise HTTPException(status_code=404, detail="Not found")
    if os.path.exists("dist/index.html"):
        return FileResponse("dist/index.html", headers=NO_CACHE_HEADERS)
    return PlainTextResponse("Frontend not built.", status_code=500)