services/hf_queue_service.py

"""HF Dataset queue: async Mac → HF prediction pipeline.

Architecture (pull-based, Mac never opens a port):
  HF Space → writes "pending" to HF Dataset queue
  Mac worker → polls queue every 5 min → runs prediction → writes "done"
  HF Space → reads fresh result from queue on next request

Queue file (queue.json in HF Dataset repo):
  { "<stock_no>": { "status": "pending|done|error",
                    "requested_at": "ISO",
                    "completed_at": "ISO|null",
                    "source": "hf_space|mac_worker",
                    "result": {...}|null } }
"""
import json
import logging
import os
import tempfile
import time
from datetime import datetime, timedelta, timezone
from pathlib import Path

logger = logging.getLogger(__name__)

DATASET_REPO        = os.getenv("HF_QUEUE_DATASET", "DennisChan0909/stock-predictor-queue")
QUEUE_FILE          = "queue.json"
RESULT_FRESH_HOURS  = 24   # "done" results older than this are re-queued
STALE_REFRESH_HOURS = 2    # Mac refreshes results older than this
_QUEUE_READ_TTL     = 90   # seconds — cache queue.json in-memory to avoid repeated HF Hub downloads

_queue_mem_cache: dict = {}   # {"data": {...}, "ts": float}


def _get_token() -> str | None:
    token = os.getenv("HF_TOKEN")
    if token:
        return token
    token_file = Path(__file__).resolve().parent.parent / ".hf_token"
    if token_file.exists():
        t = token_file.read_text().strip()
        if t and t != "PASTE_YOUR_HF_TOKEN_HERE":
            return t
    return None


def _hf_available() -> bool:
    try:
        import huggingface_hub  # noqa: F401
        return True
    except ImportError:
        return False


def is_on_hf_space() -> bool:
    return bool(os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID") or os.getenv("SPACE_HOST"))


def read_queue(force: bool = False) -> dict:
    if not _hf_available():
        return {}
    token = _get_token()
    if not token:
        return {}

    # In-memory cache — avoids repeated HF Hub downloads on every predict request
    if not force:
        cached = _queue_mem_cache.get("data")
        if cached is not None and time.time() - _queue_mem_cache.get("ts", 0) < _QUEUE_READ_TTL:
            return cached

    try:
        from huggingface_hub import hf_hub_download
        path = hf_hub_download(
            repo_id=DATASET_REPO, filename=QUEUE_FILE,
            repo_type="dataset", token=token, force_download=True,
        )
        with open(path) as f:
            data = json.load(f)
        _queue_mem_cache["data"] = data
        _queue_mem_cache["ts"] = time.time()
        return data
    except Exception as e:
        logger.debug("HF queue read: %s", e)
        return {}


def write_queue(queue: dict) -> bool:
    if not _hf_available():
        return False
    token = _get_token()
    if not token:
        return False
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            json.dump(queue, f, indent=2, default=str)
            tmp = f.name
        api.upload_file(
            path_or_fileobj=tmp,
            path_in_repo=QUEUE_FILE,
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=token,
            commit_message=f"worker {datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M')}",
        )
        os.unlink(tmp)
        # Invalidate in-memory cache so next read gets fresh data
        _queue_mem_cache.clear()
        return True
    except Exception as e:
        logger.warning("HF queue write: %s", e)
        return False


def is_publishable_result(result: object) -> bool:
    """Return True when a queue result is a completed ML-style prediction."""
    if not isinstance(result, dict):
        return False
    if result.get("source") == "quick_rules":
        return False
    if str(result.get("disclaimer", "")).startswith("Quick estimate"):
        return False
    return True


def get_cached_result(stock_no: str, max_age_hours: float | None = RESULT_FRESH_HOURS) -> dict | None:
    """Return a fresh 'done' prediction from the queue, or None."""
    queue = read_queue()
    item = queue.get(stock_no)
    if not item or item.get("status") != "done":
        return None
    result = item.get("result")
    if not is_publishable_result(result):
        return None
    completed = item.get("completed_at")
    if completed and max_age_hours is not None:
        try:
            ts = datetime.fromisoformat(completed.replace("Z", "+00:00"))
            if datetime.now(timezone.utc) - ts > timedelta(hours=max_age_hours):
                return None
        except Exception:
            pass
    return result


def enqueue(stock_no: str) -> bool:
    """Add stock_no to queue as pending. Idempotent — skips if already pending or freshly done."""
    queue = read_queue()
    existing = queue.get(stock_no, {})
    status = existing.get("status")

    if status == "pending":
        return True

    if status == "done":
        completed = existing.get("completed_at")
        if completed:
            try:
                ts = datetime.fromisoformat(completed.replace("Z", "+00:00"))
                if datetime.now(timezone.utc) - ts < timedelta(hours=STALE_REFRESH_HOURS):
                    return True  # still fresh
            except Exception:
                pass

    queue[stock_no] = {
        "status": "pending",
        "requested_at": datetime.now(timezone.utc).isoformat(),
        "completed_at": None,
        "source": "hf_space",
        "result": None,
    }
    return write_queue(queue)


def queue_status(stock_no: str) -> dict:
    """Return the current queue entry for a stock, or {"status": "not_queued"}."""
    queue = read_queue()
    return queue.get(stock_no, {"status": "not_queued"})