"""
Call for Tasks — a Gradio `gr.Server` app with a custom frontend.

A community board where people propose tasks an `ml-intern` harness should be
able to do. Visitors sign in with Hugging Face, describe a task, optionally
attach relevant HF datasets / models, and submit. As they type, similar
existing tasks surface so we avoid duplicates. A second tab lets people upvote.

Architecture
------------
* `gr.Server`           -> FastAPI + Gradio's API engine (queue, JS client).
* HF OAuth              -> "Sign in with HF" using the env vars HF Spaces injects.
                           The session id rides back in the redirect URL and is then
                           sent as a request header (no third-party iframe cookies).
* Storage = a mounted HF Storage Bucket (read-write) at $DATA_DIR (default /data).
  One JSON file per task; one empty marker file per (task, voter) so votes are
  idempotent and writes never conflict between users.
* Semantic dedup        -> sentence-transformers embeddings (lazy, cached) with a
  pure-stdlib lexical fallback so the app works even if the model can't load.
"""

from __future__ import annotations

import base64
import hashlib
import json
import os
import re
import secrets
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import gradio as gr
import httpx
from fastapi import Request
from fastapi.responses import HTMLResponse, JSONResponse, RedirectResponse
from urllib.parse import urlencode

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #

HERE = Path(__file__).resolve().parent

# The bucket is mounted read-write here in the Space; locally we fall back to a
# folder next to the app so it can be run and tested without any Hub access.
def _usable_dir(p: Path) -> bool:
    """True only if we can actually create and write inside `p`."""
    try:
        p.mkdir(parents=True, exist_ok=True)
        probe = p / ".write_probe"
        probe.write_text("ok", encoding="utf-8")
        probe.unlink()
        return True
    except OSError:
        return False


DATA_DIR = Path(os.getenv("DATA_DIR", "/data"))
if not _usable_dir(DATA_DIR):
    DATA_DIR = HERE / "local_data"
print(f"[storage] using DATA_DIR={DATA_DIR}")
TASKS_DIR = DATA_DIR / "tasks"
VOTES_DIR = DATA_DIR / "votes"
for _d in (TASKS_DIR, VOTES_DIR):
    _d.mkdir(parents=True, exist_ok=True)

# OAuth env vars are injected by HF Spaces when `hf_oauth: true` is set.
OAUTH_CLIENT_ID = os.getenv("OAUTH_CLIENT_ID")
OAUTH_CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET")
OPENID_PROVIDER_URL = os.getenv("OPENID_PROVIDER_URL", "https://huggingface.co")
OAUTH_SCOPES = os.getenv("OAUTH_SCOPES", "openid profile")
SPACE_HOST = os.getenv("SPACE_HOST")  # e.g. abidlabs-call-for-tasks.hf.space
SESSION_SECRET = os.getenv("SESSION_SECRET") or secrets.token_hex(32)
OAUTH_ENABLED = bool(OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET)

EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
# Semantic cosine scores sit higher than the lexical fallback's blended score,
# so each mode gets its own cutoff.
SEMANTIC_THRESHOLD = float(os.getenv("SIMILARITY_THRESHOLD", "0.45"))
LEXICAL_THRESHOLD = float(os.getenv("LEXICAL_THRESHOLD", "0.22"))

# --------------------------------------------------------------------------- #
# Storage layer  (operates directly on the mounted bucket filesystem)
# --------------------------------------------------------------------------- #

_io_lock = threading.Lock()


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()


def _slug(text: str, n: int = 48) -> str:
    s = re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")
    return (s[:n] or "task").strip("-")


def _task_path(task_id: str) -> Path:
    return TASKS_DIR / f"{task_id}.json"


def save_task(task: dict) -> dict:
    with _io_lock:
        _task_path(task["id"]).write_text(json.dumps(task, indent=2), encoding="utf-8")
    return task


def load_task(task_id: str) -> dict | None:
    p = _task_path(task_id)
    if not p.exists():
        return None
    try:
        return json.loads(p.read_text(encoding="utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def all_tasks() -> list[dict]:
    out: list[dict] = []
    for p in TASKS_DIR.glob("*.json"):
        try:
            out.append(json.loads(p.read_text(encoding="utf-8")))
        except (json.JSONDecodeError, OSError):
            continue
    return out


def vote_count(task_id: str) -> int:
    d = VOTES_DIR / task_id
    return sum(1 for _ in d.iterdir()) if d.is_dir() else 0


def has_voted(task_id: str, username: str) -> bool:
    return (VOTES_DIR / task_id / _slug(username, 80)).exists()


def toggle_vote(task_id: str, username: str) -> tuple[bool, int]:
    """Returns (now_voted, new_count). Idempotent and conflict-free per user."""
    d = VOTES_DIR / task_id
    d.mkdir(parents=True, exist_ok=True)
    marker = d / _slug(username, 80)
    with _io_lock:
        if marker.exists():
            marker.unlink(missing_ok=True)
            voted = False
        else:
            marker.write_text(_now_iso(), encoding="utf-8")
            voted = True
    return voted, vote_count(task_id)


def task_public_view(task: dict, viewer: str | None) -> dict:
    return {
        "id": task["id"],
        "title": task.get("title") or task["description"][:80],
        "description": task["description"],
        "rubric": task.get("rubric", ""),
        "author": task.get("author"),
        "author_name": task.get("author_name") or task.get("author"),
        "author_avatar": task.get("author_avatar"),
        "artifacts": task.get("artifacts", []),
        "created_at": task.get("created_at"),
        "votes": vote_count(task["id"]),
        "voted": bool(viewer and has_voted(task["id"], viewer)),
    }


# --------------------------------------------------------------------------- #
# Similarity engine  (semantic with a lexical fallback)
# --------------------------------------------------------------------------- #

class Similarity:
    """Lazy embedding index over task text. Falls back to lexical matching."""

    def __init__(self) -> None:
        self._model = None
        self._tried = False
        self._np = None
        self._vecs: dict[str, Any] = {}   # task_id -> np.ndarray
        self._text: dict[str, str] = {}   # task_id -> text used for the vector
        self._lock = threading.Lock()

    def _ensure_model(self):
        if self._tried:
            return self._model
        self._tried = True
        try:
            import numpy as np
            from sentence_transformers import SentenceTransformer

            self._np = np
            self._model = SentenceTransformer(EMBED_MODEL)
        except Exception as exc:  # pragma: no cover - environment dependent
            print(f"[similarity] semantic model unavailable, using lexical fallback: {exc}")
            self._model = None
        return self._model

    @staticmethod
    def _task_text(task: dict) -> str:
        arts = " ".join(a.get("repo_id", "") for a in task.get("artifacts", []))
        return f"{task.get('title', '')}. {task.get('description', '')} {task.get('rubric', '')} {arts}".strip()

    def _embed(self, text: str):
        v = self._model.encode([text], normalize_embeddings=True)[0]
        return self._np.asarray(v, dtype="float32")

    def index(self, tasks: list[dict]) -> None:
        if self._ensure_model() is None:
            return
        with self._lock:
            live = {t["id"] for t in tasks}
            for stale in set(self._vecs) - live:
                self._vecs.pop(stale, None)
                self._text.pop(stale, None)
            for t in tasks:
                text = self._task_text(t)
                if self._text.get(t["id"]) != text:
                    self._vecs[t["id"]] = self._embed(text)
                    self._text[t["id"]] = text

    def search(self, query: str, tasks: list[dict], top_k: int = 5) -> list[dict]:
        query = (query or "").strip()
        if len(query) < 6 or not tasks:
            return []
        by_id = {t["id"]: t for t in tasks}

        if self._ensure_model() is not None:
            self.index(tasks)
            qv = self._embed(query)
            scored = [
                (float(self._np.dot(qv, self._vecs[tid])), tid)
                for tid in self._vecs
                if tid in by_id
            ]
            threshold = SEMANTIC_THRESHOLD
        else:
            scored = [
                (self._lexical(query, self._task_text(t)), t["id"]) for t in tasks
            ]
            threshold = LEXICAL_THRESHOLD

        scored.sort(reverse=True)
        results = []
        for score, tid in scored[:top_k]:
            if score < threshold:
                continue
            results.append({**task_public_view(by_id[tid], None), "score": round(score, 3)})
        return results

    @staticmethod
    def _lexical(a: str, b: str) -> float:
        from difflib import SequenceMatcher

        wa, wb = set(re.findall(r"\w+", a.lower())), set(re.findall(r"\w+", b.lower()))
        jacc = len(wa & wb) / len(wa | wb) if wa and wb else 0.0
        ratio = SequenceMatcher(None, a.lower(), b.lower()).ratio()
        return 0.6 * jacc + 0.4 * ratio


similarity = Similarity()


# --------------------------------------------------------------------------- #
# Hugging Face artifact validation (best-effort, never blocks submission)
# --------------------------------------------------------------------------- #

def validate_artifact(kind: str, repo_id: str) -> dict:
    repo_id = (repo_id or "").strip().removeprefix("https://huggingface.co/")
    repo_id = repo_id.strip("/")
    if kind == "dataset":
        repo_id = repo_id.removeprefix("datasets/")
    info = {"type": kind, "repo_id": repo_id, "exists": None}
    if not repo_id:
        return info
    try:
        from huggingface_hub import HfApi

        api = HfApi()
        if kind == "dataset":
            d = api.dataset_info(repo_id)
            info.update(exists=True, likes=getattr(d, "likes", None))
        elif kind == "model":
            m = api.model_info(repo_id)
            info.update(
                exists=True,
                likes=getattr(m, "likes", None),
                downloads=getattr(m, "downloads", None),
                pipeline_tag=getattr(m, "pipeline_tag", None),
            )
        else:
            api.space_info(repo_id)
            info.update(exists=True)
    except Exception:
        info["exists"] = False
    return info


# --------------------------------------------------------------------------- #
# The Server  (FastAPI + Gradio API engine)
# --------------------------------------------------------------------------- #

app = gr.Server(title="Call for Workflows", docs_url=None, redoc_url=None)

# ---- Auth: server-side sessions, token passed via URL param + request header --
# We deliberately avoid relying on a session *cookie*. The Space runs inside a
# cross-origin iframe on huggingface.co, where third-party cookies are widely
# blocked — so a cookie-based login won't reflect without a manual refresh.
# Instead the OAuth callback redirects to "/?oauth_session=<id>"; the frontend
# stores that id and sends it as the `x-cft-session` header on every request.
# (Same approach as trackio.)

OAUTH_BASE = OPENID_PROVIDER_URL.rstrip("/")
SESSION_TTL = 86400 * 30
_oauth_states: dict[str, float] = {}
_sessions: dict[str, dict] = {}


def _evict_expired() -> None:
    now = time.time()
    for k in [k for k, t in _oauth_states.items() if now - t > 3600]:
        _oauth_states.pop(k, None)
    for k in [k for k, s in _sessions.items() if now - s["created"] > SESSION_TTL]:
        _sessions.pop(k, None)


def _new_session(user: dict, token: str | None = None) -> str:
    sid = secrets.token_urlsafe(24)
    _sessions[sid] = {"user": user, "token": token, "created": time.time()}
    return sid


def _sid_from_request(request: Request) -> str | None:
    sid = request.headers.get("x-cft-session")
    if not sid:
        auth = request.headers.get("authorization", "")
        if auth.lower().startswith("bearer "):
            sid = auth[7:].strip()
    return sid or request.cookies.get("cft_session")


def current_user(request: Request) -> dict | None:
    sid = _sid_from_request(request)
    s = _sessions.get(sid or "")
    if s and time.time() - s["created"] <= SESSION_TTL:
        return s["user"]
    if sid:
        _sessions.pop(sid, None)
    return None


def require_user(request: Request) -> dict | None:
    return current_user(request)


# ---- Gradio API endpoint: live "find similar" (goes through Gradio's queue) -- #

@app.api(name="find_similar", concurrency_limit=4)
def find_similar(query: str, top_k: int = 5) -> list:
    """Return existing tasks similar to a draft description.

    Exposed as a Gradio API endpoint so the frontend can call it through the
    Gradio JS client (queued + concurrency-controlled) and so it's reachable
    from `gradio_client` / MCP too.
    """
    return similarity.search(query, all_tasks(), top_k=top_k)


# ---- Auth routes ----------------------------------------------------------- #

def _redirect_uri(request: Request) -> str:
    if SPACE_HOST:
        return f"https://{SPACE_HOST.split(',')[0]}/auth/callback"
    return str(request.base_url).rstrip("/") + "/auth/callback"


def _session_redirect(sid: str) -> RedirectResponse:
    # The whole page navigates here, so login reflects on load with no refresh.
    resp = RedirectResponse(f"/?oauth_session={sid}", status_code=302)
    # Cookie is a best-effort fallback only; the header is the source of truth.
    resp.set_cookie("cft_session", sid, max_age=SESSION_TTL, httponly=True,
                    samesite="none", secure=True, path="/")
    return resp


@app.get("/login")
async def login(request: Request):
    if not OAUTH_ENABLED:
        # Local dev: sign in as a throwaway user so the flow is testable.
        sid = _new_session({
            "username": "dev-user", "name": "Local Dev",
            "avatar": "https://huggingface.co/avatars/placeholder.svg",
        })
        return _session_redirect(sid)
    _evict_expired()
    state = secrets.token_urlsafe(24)
    _oauth_states[state] = time.time()
    url = f"{OAUTH_BASE}/oauth/authorize?" + urlencode({
        "client_id": OAUTH_CLIENT_ID,
        "redirect_uri": _redirect_uri(request),
        "response_type": "code",
        "scope": OAUTH_SCOPES,
        "state": state,
    })
    return RedirectResponse(url, status_code=302)


@app.get("/auth/callback", name="auth_callback")
async def auth_callback(request: Request):
    code = request.query_params.get("code")
    state = request.query_params.get("state")
    if not code or not state or state not in _oauth_states:
        return RedirectResponse("/?oauth_error=1", status_code=302)
    _oauth_states.pop(state, None)
    basic = base64.b64encode(f"{OAUTH_CLIENT_ID}:{OAUTH_CLIENT_SECRET}".encode()).decode()
    try:
        async with httpx.AsyncClient(timeout=10.0) as client:
            tok = await client.post(
                f"{OAUTH_BASE}/oauth/token",
                headers={"Authorization": f"Basic {basic}"},
                data={"grant_type": "authorization_code", "code": code,
                      "redirect_uri": _redirect_uri(request), "client_id": OAUTH_CLIENT_ID},
            )
            tok.raise_for_status()
            access_token = tok.json()["access_token"]
            ui = await client.get(
                f"{OAUTH_BASE}/oauth/userinfo",
                headers={"Authorization": f"Bearer {access_token}"},
            )
            ui.raise_for_status()
            info = ui.json()
    except (httpx.HTTPError, KeyError, ValueError) as exc:
        print(f"[oauth] callback failed: {exc}")
        return RedirectResponse("/?oauth_error=1", status_code=302)
    user = {
        "username": info.get("preferred_username") or info.get("name"),
        "name": info.get("name") or info.get("preferred_username"),
        "avatar": info.get("picture"),
    }
    return _session_redirect(_new_session(user, access_token))


@app.get("/logout")
async def logout(request: Request):
    sid = _sid_from_request(request)
    if sid:
        _sessions.pop(sid, None)
    resp = JSONResponse({"ok": True})
    resp.delete_cookie("cft_session", path="/", samesite="none", secure=True)
    return resp


@app.get("/api/me")
async def api_me(request: Request):
    return JSONResponse({"user": current_user(request), "oauth": OAUTH_ENABLED})


# ---- Data routes ----------------------------------------------------------- #

@app.get("/api/tasks")
async def api_tasks(request: Request, sort: str = "top"):
    viewer = (current_user(request) or {}).get("username")
    tasks = [task_public_view(t, viewer) for t in all_tasks()]
    if sort == "new":
        tasks.sort(key=lambda t: t.get("created_at") or "", reverse=True)
    else:
        tasks.sort(key=lambda t: (t["votes"], t.get("created_at") or ""), reverse=True)
    return JSONResponse({"tasks": tasks, "count": len(tasks)})


@app.post("/api/validate-artifact")
async def api_validate(request: Request):
    body = await request.json()
    return JSONResponse(validate_artifact(body.get("type", "model"), body.get("repo_id", "")))


@app.get("/api/search-artifacts")
async def api_search_artifacts(q: str = "", limit: int = 6):
    """Proxy Hugging Face's quicksearch (the API behind hf.co's own search bar).

    The browser can't call it directly (its CORS is locked to huggingface.co),
    so we proxy here and normalise the response into a single ranked list where
    each item already carries its type — no type picker needed on the frontend.
    """
    q = (q or "").strip()
    if len(q) < 2:
        return JSONResponse({"results": []})
    try:
        async with httpx.AsyncClient(timeout=4.0) as client:
            r = await client.get(
                "https://huggingface.co/api/quicksearch",
                params={"q": q, "limit": limit},
                headers={"User-Agent": "call-for-tasks-space"},
            )
            r.raise_for_status()
            data = r.json()
    except (httpx.HTTPError, ValueError):
        return JSONResponse({"results": []})

    # Interleave the three categories so the dropdown shows a varied top slice,
    # mirroring how HF's own search blends models / datasets / Spaces.
    buckets = [
        ("model", data.get("models", [])),
        ("dataset", data.get("datasets", [])),
        ("space", data.get("spaces", [])),
    ]
    results: list[dict] = []
    for i in range(limit):
        for kind, items in buckets:
            if i < len(items):
                rid = items[i].get("id")
                if rid:
                    results.append({"type": kind, "repo_id": rid})
    return JSONResponse({"results": results[: limit * 2]})


@app.post("/api/submit")
async def api_submit(request: Request):
    user = require_user(request)
    if not user:
        return JSONResponse({"error": "You must sign in with Hugging Face to submit."}, status_code=401)

    body = await request.json()
    description = (body.get("description") or "").strip()
    title = (body.get("title") or "").strip()
    rubric = (body.get("rubric") or "").strip()[:2000]
    if len(description) < 15:
        return JSONResponse({"error": "Please describe the workflow in a little more detail (15+ chars)."}, status_code=400)

    artifacts = []
    for a in body.get("artifacts", [])[:10]:
        repo_id = (a.get("repo_id") or "").strip()
        if repo_id:
            artifacts.append({"type": a.get("type", "model"), "repo_id": repo_id})

    task_id = f"{int(time.time())}-{_slug(title or description, 32)}-{secrets.token_hex(2)}"
    task = {
        "id": task_id,
        "title": title or description[:80],
        "description": description,
        "rubric": rubric,
        "artifacts": artifacts,
        "author": user.get("username"),
        "author_name": user.get("name"),
        "author_avatar": user.get("avatar"),
        "created_at": _now_iso(),
    }
    save_task(task)
    # Author implicitly upvotes their own task.
    toggle_vote(task_id, user.get("username"))
    return JSONResponse({"ok": True, "task": task_public_view(task, user.get("username"))})


@app.post("/api/vote")
async def api_vote(request: Request):
    user = require_user(request)
    if not user:
        return JSONResponse({"error": "Sign in to upvote."}, status_code=401)
    body = await request.json()
    task_id = body.get("task_id", "")
    if not load_task(task_id):
        return JSONResponse({"error": "Unknown task."}, status_code=404)
    voted, count = toggle_vote(task_id, user.get("username"))
    return JSONResponse({"ok": True, "voted": voted, "votes": count})


# ---- Frontend -------------------------------------------------------------- #

_INDEX_HTML = (HERE / "index.html").read_text(encoding="utf-8")


@app.get("/", response_class=HTMLResponse)
async def homepage():
    return _INDEX_HTML


@app.get("/healthz")
async def healthz():
    return {"status": "ok", "tasks": len(list(TASKS_DIR.glob("*.json"))), "oauth": OAUTH_ENABLED}


if __name__ == "__main__":
    host = os.getenv("GRADIO_SERVER_NAME", "0.0.0.0")
    port = int(os.getenv("GRADIO_SERVER_PORT") or os.getenv("PORT") or 7860)
    app.launch(server_name=host, server_port=port)