/**
 * siq-proxy — zero-dependency streaming OpenAI proxy for the SIQ-1-35B RunPod
 * serverless endpoint, hosted as a Hugging Face Docker Space.
 *
 * The Pi Agent web app (alexwortega.github.io/pi-agent) can't call RunPod
 * directly: the API key would leak into the public bundle, RunPod returns no
 * CORS headers, and its serverless API is async /run+poll. This server holds the
 * key (an HF Space secret), adds CORS, applies SIQ-1 request shaping, and streams
 * the RunPod vLLM worker's OpenAI passthrough (SSE) straight through to the browser.
 *
 * Mirrors the /api/siq route in piagent_web/server/index.js. Secrets / env:
 *   RUNPOD_API_KEY (required)  SIQ_EID (required)  SIQ_MODEL (default "siq")
 *   SIQ_MINTOK (default 2048)  SIQ_LIMIT_HOUR_IP (default 240)  PORT (default 7860)
 */
import http from "node:http";

const PORT = parseInt(process.env.PORT || "7860", 10);
const RUNPOD_API_KEY = process.env.RUNPOD_API_KEY || "";
const SIQ_EID = process.env.SIQ_EID || "";
const SIQ_MODEL = process.env.SIQ_MODEL || "siq";
const SIQ_MINTOK = parseInt(process.env.SIQ_MINTOK || "2048", 10);
const SIQ_LIMIT_HOUR_IP = parseInt(process.env.SIQ_LIMIT_HOUR_IP || "240", 10);
const siqEnabled = !!(RUNPOD_API_KEY && SIQ_EID);

const ALLOWED = new Set([
  "https://alexwortega.github.io",
  "http://localhost:5050",
  "http://localhost:4173",
  "http://localhost:5173",
]);

// in-memory per-IP rate limiter (single-instance Space)
const counters = new Map();
function rateLimit(key, limit, windowMs = 3_600_000) {
  const now = Date.now();
  const arr = (counters.get(key) || []).filter((t) => now - t < windowMs);
  if (arr.length >= limit) return { ok: false, retryInMs: windowMs - (now - arr[0]) };
  arr.push(now);
  counters.set(key, arr);
  return { ok: true };
}
setInterval(() => {
  const now = Date.now();
  for (const [k, arr] of counters) {
    const fresh = arr.filter((t) => now - t < 3_600_000);
    if (fresh.length === 0) counters.delete(k);
    else counters.set(k, fresh);
  }
}, 600_000).unref();

function cors(req, res) {
  const origin = req.headers.origin;
  res.setHeader("Access-Control-Allow-Origin", origin && ALLOWED.has(origin) ? origin : "*");
  res.setHeader("Access-Control-Allow-Methods", "GET,POST,OPTIONS");
  res.setHeader("Access-Control-Allow-Headers", "Content-Type");
  res.setHeader("Access-Control-Max-Age", "86400");
}
function json(res, code, obj) {
  const body = JSON.stringify(obj);
  res.writeHead(code, { "Content-Type": "application/json" });
  res.end(body);
}
function clientIp(req) {
  return ((req.headers["x-forwarded-for"] || "").toString().split(",")[0].trim() ||
    req.socket.remoteAddress || "unknown").slice(0, 64);
}
function readBody(req) {
  return new Promise((resolve) => {
    let buf = "";
    req.on("data", (c) => { buf += c; if (buf.length > 4_000_000) req.destroy(); });
    req.on("end", () => { try { resolve(JSON.parse(buf || "{}")); } catch { resolve(null); } });
    req.on("error", () => resolve(null));
  });
}

const RP = `https://api.runpod.ai/v2/${SIQ_EID}`;
const rpHeaders = { "Content-Type": "application/json", Authorization: `Bearer ${RUNPOD_API_KEY}` };

/**
 * Submit a job via RunPod async /run and poll /status to completion. We use the
 * async API (not /runsync, which times out under concurrency — siq1.md §3.A).
 * The SIQ-1 GGUF worker (dwarfgen/siq1-gguf-worker) takes the OpenAI chat request
 * DIRECTLY as `input` (NOT wrapped in openai_route/openai_input — that envelope is
 * for the vLLM worker, and this worker silently ignores it and answers an empty
 * prompt). Returns the worker's OpenAI-shaped output (or throws).
 */
async function runAndPoll(openaiInput, signal, onStatus, pollMs = 280_000) {
  const r = await fetch(`${RP}/run`, {
    method: "POST",
    headers: rpHeaders,
    signal,
    body: JSON.stringify({ input: openaiInput }),
  });
  const job = await r.json();
  if (job.output) return job.output; // some workers answer synchronously
  const id = job.id;
  if (!id) throw new Error("no job id: " + JSON.stringify(job).slice(0, 200));
  const deadline = Date.now() + pollMs;
  let last = "";
  while (Date.now() < deadline) {
    if (signal?.aborted) throw new Error("aborted");
    await new Promise((res) => setTimeout(res, 1200));
    const st = await fetch(`${RP}/status/${id}`, { headers: rpHeaders, signal }).then((x) => x.json());
    // surface worker lifecycle so the client can show "starting" vs "generating"
    if (st.status && st.status !== last) { last = st.status; onStatus?.(st.status); }
    if (st.status === "COMPLETED") {
      const out = st.output;
      if (Array.isArray(out) && out.length) return out[out.length - 1];
      return out;
    }
    if (["FAILED", "CANCELLED", "TIMED_OUT"].includes(st.status)) {
      throw new Error(`job ${st.status}: ` + JSON.stringify(st).slice(0, 200));
    }
  }
  throw new Error("poll timeout");
}

/** Pull {content, reasoning} out of whatever shape the worker returned. */
function extractMessage(out) {
  const msg = out?.choices?.[0]?.message || out?.choices?.[0]?.delta || {};
  return {
    content: msg.content ?? (typeof out === "string" ? out : ""),
    reasoning: msg.reasoning_content ?? "",
  };
}

async function handleChat(req, res) {
  if (!siqEnabled) return json(res, 503, { error: { message: "SIQ proxy not configured (RUNPOD_API_KEY / SIQ_EID)", type: "config" } });
  const body = await readBody(req);
  if (!body || !Array.isArray(body.messages) || body.messages.length === 0) {
    return json(res, 400, { error: { message: "messages required", type: "invalid_request" } });
  }
  const ip = clientIp(req);
  const lim = rateLimit("ip:" + ip, SIQ_LIMIT_HOUR_IP);
  if (!lim.ok) return json(res, 429, { error: { message: `rate limit: ${SIQ_LIMIT_HOUR_IP}/hour per IP`, type: "rate_limit", retry_in_s: Math.ceil(lim.retryInMs / 1000) } });

  // ---- SIQ-1 request shaping (siq1.md §3): thinking toggle, effort, min tokens.
  const ctk = { ...(body.chat_template_kwargs || {}) };
  const thinking = ctk.enable_thinking !== false;
  const messages = body.messages.map((m) => ({ role: m.role, content: m.content }));
  const effort = typeof body.effort === "string" ? body.effort : "";
  if (thinking && effort) {
    // Merge the effort directive INTO the existing system message — this GGUF
    // worker returns an empty completion if it gets a second system message.
    const line = `Reasoning effort: ${effort}`;
    const si = messages.findIndex((m) => m.role === "system");
    if (si >= 0) {
      if (!/Reasoning effort:/i.test(String(messages[si].content || ""))) {
        messages[si] = { ...messages[si], content: `${line}\n\n${messages[si].content}` };
      }
    } else {
      messages.unshift({ role: "system", content: line });
    }
  }
  const openaiInput = {
    model: SIQ_MODEL,
    messages,
    temperature: typeof body.temperature === "number" ? body.temperature : 0.7,
    max_tokens: Math.max(parseInt(body.max_tokens, 10) || 0, SIQ_MINTOK),
    top_p: typeof body.top_p === "number" ? body.top_p : 0.95,
    chat_template_kwargs: { enable_thinking: thinking },
  };
  if (typeof body.top_k === "number") openaiInput.top_k = body.top_k;
  if (typeof body.presence_penalty === "number") openaiInput.presence_penalty = body.presence_penalty;

  const ac = new AbortController();
  let finished = false;
  // real client disconnect = res 'close' before we finish (req 'close' can fire
  // early once the body is buffered, which is NOT a disconnect).
  res.on("close", () => { if (!finished) ac.abort(); });

  // SSE framing — the client engine accumulates delta.reasoning_content (folded
  // into <think>…</think>) and delta.content. RunPod async isn't token-streaming,
  // so we emit the completed result as one reasoning delta + one content delta.
  res.writeHead(200, {
    "Content-Type": "text/event-stream; charset=utf-8",
    "Cache-Control": "no-cache, no-transform",
    Connection: "keep-alive",
  });
  const sse = (obj) => res.write(`data: ${JSON.stringify(obj)}\n\n`);
  // keepalive comments so the connection isn't reaped during a long generation
  const ka = setInterval(() => res.write(": keepalive\n\n"), 15_000);
  try {
    const out = await runAndPoll(openaiInput, ac.signal, (status) => sse({ siq_status: status }));
    const { content, reasoning } = extractMessage(out);
    if (reasoning) sse({ choices: [{ index: 0, delta: { reasoning_content: reasoning } }] });
    sse({ choices: [{ index: 0, delta: { content }, finish_reason: "stop" }] });
    res.write("data: [DONE]\n\n");
  } catch (e) {
    sse({ error: { message: "siq upstream: " + (e?.message || e), type: "upstream" } });
  } finally {
    finished = true;
    clearInterval(ka);
    res.end();
  }
}

const server = http.createServer(async (req, res) => {
  cors(req, res);
  if (req.method === "OPTIONS") { res.writeHead(204); return res.end(); }
  const path = (req.url || "/").split("?")[0].replace(/\/+$/, "") || "/";

  if (req.method === "GET" && (path === "/" || path === "/api/health")) {
    return json(res, 200, { ok: true, siq: siqEnabled, model: SIQ_MODEL });
  }
  if (req.method === "GET" && path === "/api/siq/v1/models") {
    return json(res, 200, { object: "list", data: [{ id: SIQ_MODEL, object: "model", owned_by: "runpod-serverless" }] });
  }
  if (req.method === "POST" && path === "/api/siq/v1/chat/completions") {
    return handleChat(req, res);
  }
  json(res, 404, { error: { message: "not found", type: "not_found" } });
});

server.listen(PORT, "0.0.0.0", () => console.log(`siq-proxy on :${PORT} (siq=${siqEnabled}, eid=${SIQ_EID || "unset"}, model=${SIQ_MODEL})`));