api/job_queue.py

import os
import json
import time
import asyncio
import logging
from typing import Optional, List, Dict, Any
from fastapi import APIRouter, HTTPException, Request
from pydantic import BaseModel
import httpx

_logger = logging.getLogger("agente_ai.jq")
router = APIRouter(prefix="/jq", tags=["job_queue"])

# ── Configurazione ────────────────────────────────────────────────────────────
# S42: Grid Orchestrator — Cluster 4x4 (A, B, C, D)
# SPACE_ROLE: brain | hands | memory | audit
_SPACE_ROLE  = os.getenv("SPACE_ROLE", "unknown")
_JQ_ENABLED  = os.getenv("JQ_ENABLED", "0") == "1"
_INTERNAL_TOKEN = os.getenv("INTERNAL_TOKEN", "")

# Redis keys (Upstash)
_K_PENDING  = "jq:pending"
_K_WAKE     = "jq:wake"
_K_CONSUMER = f"jq:consumer:{_SPACE_ROLE}"
_K_RESULT   = lambda tid: f"jq:result:{tid}"
_K_EVENTS   = lambda tid: f"jq:events:{tid}"
_K_LOAD     = lambda role: f"jq:load:{role}"

class JobPayload(BaseModel):
    taskId: str
    goal: str
    context: Optional[Dict[str, Any]] = None
    priority: int = 1

# ── Helper Redis (via HTTP REST per stabilità mobile/serverless) ──────────────
async def _rcmd(cmd: List[Any]) -> Optional[Dict[str, Any]]:
    url = os.getenv("UPSTASH_REDIS_REST_URL")
    tok = os.getenv("UPSTASH_REDIS_REST_TOKEN")
    if not url or not tok:
        return None
    try:
        async with httpx.AsyncClient() as client:
            r = await client.post(
                url,
                headers={"Authorization": f"Bearer {tok}"},
                json=cmd,
                timeout=5.0
            )
            return r.json()
    except Exception as e:
        _logger.error("[jq] redis error: %s", e)
        return None

def _redis_ok() -> bool:
    return bool(os.getenv("UPSTASH_REDIS_REST_URL") and os.getenv("UPSTASH_REDIS_REST_TOKEN"))

async def _llen(key: str) -> int:
    res = await _rcmd(["LLEN", key])
    return int(res.get("result", 0)) if res else 0

async def _lrange(key: str, start: int, end: int = -1) -> List[str]:
    res = await _rcmd(["LRANGE", key, start, end])
    return res.get("result", []) if res else []

# ── Core Logic ────────────────────────────────────────────────────────────────
async def publish_load_metrics():
    """Pubblica il carico corrente su Redis per il bilanciamento."""
    if not _redis_ok(): return
    data = {
        "role": _SPACE_ROLE,
        "ts": int(time.time() * 1000),
        "active_tasks": len(asyncio.all_tasks()) - 5, # Stima approssimativa
        "cpu": 0, # Placeholder per metrics reali
        "mem": 0
    }
    await _rcmd(["SET", _K_LOAD(_SPACE_ROLE), json.dumps(data), "EX", "60"])

async def _load_publisher_loop():
    """Loop periodico per aggiornare lo stato del nodo."""
    while True:
        try:
            await publish_load_metrics()
        except Exception as e:
            _logger.debug("[jq] load publisher error: %s", e)
        await asyncio.sleep(30)

async def _hands_consumer_loop():
    """Loop consumer per nodi HANDS, MEMORY e AUDIT."""
    _logger.info("[jq] consumer loop avviato per ruolo: %s", _SPACE_ROLE)
    while True:
        try:
            # S429: Tool Success Contract — verifica heartbeat consumer
            await _rcmd(["SET", _K_CONSUMER, "1", "EX", "15"])
            
            # TODO: Implementazione effettiva del prelievo job e esecuzione
            # Per ora è un segnaposto per la struttura UltraVSS v7.0
            
        except Exception as e:
            _logger.debug("[jq] consumer loop error: %s", e)
        await asyncio.sleep(5)

async def start_job_queue_consumer() -> None:
    """Punto di ingresso per main.py _on_startup()."""
    if not _redis_ok():
        _logger.warning("[jq] Redis non configurato — job queue disabilitato")
        return

    asyncio.create_task(_load_publisher_loop())

    # Grid Orchestrator: Consumer specializzati (HANDS, MEMORY, AUDIT)
    if _SPACE_ROLE in ("hands", "memory", "audit", "unknown"):
        asyncio.create_task(_hands_consumer_loop())
    else:
        _logger.info("[jq] SPACE_ROLE=%s — consumer non avviato (solo load publisher)", _SPACE_ROLE)

# ── FastAPI endpoints ──────────────────────────────────────────────────────────
@router.get("/status")
async def jq_status():
    result = {
        "space_role":       _SPACE_ROLE,
        "jq_enabled":       _JQ_ENABLED,
        "redis_configured": _redis_ok(),
        "ts":               int(time.time() * 1000),
    }
    return result

@router.get("/load/{role}")
async def jq_load(role: str):
    if role not in ("brain", "hands", "memory", "audit"):
        raise HTTPException(400, "role non valido")
    res = await _rcmd(["GET", _K_LOAD(role)])
    if not res or not res.get("result"):
        raise HTTPException(404, f"Metriche {role} non disponibili")
    return json.loads(res["result"])

@router.post("/submit")
async def jq_submit(job: JobPayload, request: Request):
    if _INTERNAL_TOKEN and request.headers.get("X-Internal-Token") != _INTERNAL_TOKEN:
        raise HTTPException(401, "Unauthorized")
    # Logica di sottomissione job...
    return {"taskId": job.taskId, "status": "queued"}