# server.py
import os
import io
import json
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

# === Config ===
REPO_ID = os.environ.get("STORE_REPO", "TheFinAI/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 ----------
def ensure_repo():
    try:
        create_repo(REPO_ID, repo_type="dataset", exist_ok=True, token=HF_TOKEN)
    except Exception:
        pass

def upload_json(path_in_repo: str, obj: Any, message: str):
    ensure_repo()
    buf = io.BytesIO(json.dumps(obj, ensure_ascii=False, indent=2).encode("utf-8"))
    upload_file(
        path_or_fileobj=buf,
        path_in_repo=path_in_repo,
        repo_id=REPO_ID,
        repo_type="dataset",
        token=HF_TOKEN,
        commit_message=message,
    )

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 ----------
# 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")
    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")
    return PlainTextResponse("Frontend not built.", status_code=500)