/**
 * server.ts — Bun HTTP server: owns the agent loop, broadcasts state via SSE.
 *
 * One AI agent runs server-side against the runpod llama.cpp endpoint with
 * structured tool calling, KV-cache reuse, and LLM-summarization compaction.
 * Connected clients are pure viewers — they GET / for the viewer HTML and
 * subscribe to /stream for state updates. No client-side inference, no auth.
 */

import { serve } from 'bun';
import { readFile } from 'fs/promises';
import { SinglePlayerGame, DEFAULTS } from './engine.js';
import {
  PICK_FIRST_TOOLS,
  PICK_FIRST_SYSTEM_PROMPT,
  buildPickFirstPrompt,
  AgentMemory,
} from './agent-brain.js';

// ─── Config ──────────────────────────────────────────────────────────────────

const ENDPOINT = (process.env.ENDPOINT ?? 'https://zopt7f90d1ydgy-8080.proxy.runpod.net').replace(/\/$/, '');
const REMOTE_MODEL = process.env.MODEL ?? 'default';
const PORT = parseInt(process.env.PORT ?? '7860', 10);
const MAX_TOKENS = 100;
const TEMPERATURE = 0.6;

// ─── Singleton state ─────────────────────────────────────────────────────────

const game = new SinglePlayerGame();
const memory = new AgentMemory(PICK_FIRST_SYSTEM_PROMPT, {
  tokenBudget: 4000,
  triggerCompactTokens: 2000,
  keepRecentTokens: 700,
});
const actionLog: Array<{ action: string; detail: string; success: boolean }> = [];
let fallbackCount = 0;
let lastUsage: any = null;
let lastTotalMs = 0;
let serverCtx = 8192;
let remoteModelId: string | null = null;
let status = 'connecting';
let paused = process.env.START_PAUSED === '1';
let lastThinking: string | null = null;

const subscribers = new Set<ReadableStreamDefaultController<Uint8Array>>();

// ─── Trace events (pi-mono `--mode json` shape) ──────────────────────────────
// One JSON object per line on stdout. Event types mirror pi-mono's
// AgentSessionEvent vocabulary; extra app-specific events (fallback_action,
// viewer_connect, world_reset) follow the same discriminated-union pattern.

type TraceEvent = Record<string, unknown>;

function trace(event: TraceEvent) {
  process.stdout.write(JSON.stringify({ ts: new Date().toISOString(), ...event }) + '\n');
}

// ─── State serialization + broadcast ─────────────────────────────────────────

function buildState() {
  return {
    viewport: game.getViewport(),
    inventory: game.inventory,
    stats: game.stats,
    actions: actionLog.slice(-30),
    status,
    thinking: lastThinking,
    summary: memory.summary,
    debug: {
      backend: 'Remote (llama.cpp)',
      endpoint: ENDPOINT.replace(/^https?:\/\//, ''),
      model: remoteModelId,
      serverCtx,
      promptTokens: lastUsage?.prompt_tokens ?? 0,
      cachedTokens: lastUsage?.prompt_tokens_details?.cached_tokens ?? 0,
      completionTokens: lastUsage?.completion_tokens ?? 0,
      totalMs: lastTotalMs,
      memTokens: memory.estimateTokens(),
      memMessages: memory.messages.length,
      fallbacks: fallbackCount,
      compactions: memory.compactions,
      summaryLen: memory.summary?.length ?? 0,
      turn: actionLog.length,
      viewers: subscribers.size,
      cacheRetries,
      cacheCooldown,
      paused,
    },
  };
}

const enc = new TextEncoder();
function broadcast() {
  const payload = enc.encode(`data: ${JSON.stringify(buildState())}\n\n`);
  for (const ctrl of subscribers) {
    try {
      ctrl.enqueue(payload);
    } catch {
      subscribers.delete(ctrl);
    }
  }
}

// ─── Connect: probe endpoint, derive memory budgets from n_ctx ───────────────

async function connect() {
  status = 'connecting';
  // Retry the health check forever — if the upstream pod is restarting, we
  // want to wait it out, not crash the Space and fall into RUNTIME_ERROR.
  while (true) {
    const health = await fetch(`${ENDPOINT}/health`).then(r => r.json()).catch(() => null);
    if (health && health.status === 'ok') break;
    trace({ type: 'connect_waiting', endpoint: ENDPOINT });
    await new Promise(r => setTimeout(r, 5000));
  }

  const models = await fetch(`${ENDPOINT}/v1/models`).then(r => r.json()).catch(() => null);
  remoteModelId = models?.data?.[0]?.id || REMOTE_MODEL;

  const props = await fetch(`${ENDPOINT}/props`).then(r => r.json()).catch(() => null);
  serverCtx = props?.default_generation_settings?.n_ctx ?? 8192;
  // Reserve room for: completion buffer (MAX_TOKENS), tool/template overhead (~400 tok),
  // and a 1.7x margin since chars/4 undercounts real tokens (measured ratio with
  // structured tool_calls + Qwen3.6 chat template was ~1.6).
  const safe = Math.max(1500, Math.floor((serverCtx - MAX_TOKENS - 400) / 1.7));
  memory.tokenBudget = safe;
  memory.triggerCompactTokens = Math.floor(safe * 0.6);
  memory.keepRecentTokens = Math.floor(safe * 0.25);
  trace({
    type: 'agent_start',
    endpoint: ENDPOINT,
    model: remoteModelId,
    n_ctx: serverCtx,
    memoryConfig: {
      tokenBudget: memory.tokenBudget,
      triggerCompactTokens: memory.triggerCompactTokens,
      keepRecentTokens: memory.keepRecentTokens,
    },
  });
  status = 'running';
}

// ─── Inference ───────────────────────────────────────────────────────────────

// llama.cpp build b9038-2ca1161bd has a regression where cache_prompt:true
// causes slots to accumulate state past n_ctx and then 500 every request to
// that slot permanently. The cooldown logic below is a vestige; in practice
// we leave caching disabled (USE_CACHE=false) and pay the re-tokenize cost
// every turn. Set ENV USE_CACHE=1 to opt back in if the build is fixed upstream.
const USE_CACHE = process.env.USE_CACHE === '1';
let cacheCooldown = USE_CACHE ? 0 : Infinity;
const COOLDOWN_TURNS = 30;
const RETRY_ATTEMPTS = 3;
const RETRY_DELAY_MS = 200;
let cacheRetries = 0;

async function callModel() {
  const t0 = performance.now();
  const useCache = cacheCooldown === 0;
  const body: any = {
    model: remoteModelId,
    messages: memory.getMessages(),
    tools: PICK_FIRST_TOOLS,
    tool_choice: 'auto',
    temperature: TEMPERATURE,
    max_tokens: MAX_TOKENS,
    chat_template_kwargs: { enable_thinking: false },
  };
  body.cache_prompt = useCache; // explicit true/false; omitting falls back to broken default

  trace({ type: 'message_start' });

  let res: Response;
  let attempt = 0;
  while (true) {
    res = await fetch(`${ENDPOINT}/v1/chat/completions`, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify(body),
    });
    if (res.ok) break;

    const text = await res.text();
    const isContextErr = res.status === 500 && /Context size/i.test(text);

    if (!isContextErr || attempt >= RETRY_ATTEMPTS) {
      if (attempt > 0) {
        trace({
          type: 'auto_retry_end',
          success: false,
          attempt,
          finalError: `HTTP ${res.status}: ${text}`,
        });
      }
      throw new Error(`HTTP ${res.status}: ${text}`);
    }

    attempt++;
    // On first context-size error, explicitly disable prompt caching for this
    // and future turns until cooldown expires. Setting cache_prompt:false is
    // not the same as omitting the key — the server default routes back to
    // the poisoned cached path on this build (b9038-2ca1161bd).
    if (body.cache_prompt !== false) {
      if (body.cache_prompt === true) {
        cacheRetries++;
        cacheCooldown = COOLDOWN_TURNS;
      }
      body.cache_prompt = false;
    }
    trace({
      type: 'auto_retry_start',
      attempt,
      maxAttempts: RETRY_ATTEMPTS,
      delayMs: RETRY_DELAY_MS,
      errorMessage: text,
      reason: 'slot_saturation',
    });
    await new Promise((r) => setTimeout(r, RETRY_DELAY_MS));
  }

  if (attempt > 0) {
    trace({ type: 'auto_retry_end', success: true, attempt });
  }
  if (cacheCooldown > 0) cacheCooldown--;

  const data = await res.json();
  const totalMs = Math.round(performance.now() - t0);
  const message = data.choices?.[0]?.message ?? {};
  const u = data.usage ?? {};
  // Embed usage on the message in pi-mono shape (input/output/cacheRead/cacheWrite).
  const usageOnMessage = {
    input: u.prompt_tokens ?? 0,
    output: u.completion_tokens ?? 0,
    cacheRead: u.prompt_tokens_details?.cached_tokens ?? 0,
    cacheWrite: 0,
  };
  trace({
    type: 'message_end',
    message: { ...message, usage: usageOnMessage },
    durationMs: totalMs,
  });
  return { message, usage: data.usage, totalMs };
}

// ─── Summarizer ──────────────────────────────────────────────────────────────

const SUMMARY_SYS = `You compress 2D-grid-world game histories into terse, factual notes for the agent's future planning. Output 5-10 bullet points, under 200 words. Capture: directions tried and outcomes, diamonds collected (count + approximate area), obstacles encountered, exploration trajectory, and any failed actions worth remembering — especially repeated failures (e.g. "moved up 5 times, every attempt failed against an obstacle — direction up is blocked, try another"). No prose, no preamble.`;

async function callSummarizer(turnsLog: string, previousSummary: string | null) {
  const userPrompt = previousSummary
    ? `Previous summary:\n${previousSummary}\n\nNew events to merge in:\n${turnsLog}\n\nProduce an updated unified summary.`
    : `Game history so far:\n${turnsLog}\n\nProduce a summary.`;
  const t0 = performance.now();
  // Thinking is enabled for the summarizer (and only the summarizer) so the
  // model can actually reason about what happened — failed-direction loops,
  // unreachable diamonds, etc. — and produce a richer summary that helps the
  // (thinking-off) turn loop break out of behavioural locks. max_tokens is
  // large enough for the reasoning chain plus the bullet output.
  const body: any = {
    model: remoteModelId,
    messages: [
      { role: 'system', content: SUMMARY_SYS },
      { role: 'user', content: userPrompt },
    ],
    temperature: 0.2,
    max_tokens: 2048,
    chat_template_kwargs: { enable_thinking: true },
  };
  // Explicit cache_prompt; omitting it falls back to the broken default on this build.
  body.cache_prompt = cacheCooldown === 0;

  const res = await fetch(`${ENDPOINT}/v1/chat/completions`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(body),
  });
  if (!res.ok) throw new Error(`summarizer HTTP ${res.status}: ${await res.text()}`);
  const data = await res.json();
  const totalMs = Math.round(performance.now() - t0);
  const out = (data.choices?.[0]?.message?.content || '').trim();
  // Surface the summarizer call as a nested message_end so consumers see
  // the same shape as a normal turn's inference.
  const u = data.usage ?? {};
  trace({
    type: 'message_end',
    source: 'summarizer',
    message: {
      role: 'assistant',
      content: out,
      usage: {
        input: u.prompt_tokens ?? 0,
        output: u.completion_tokens ?? 0,
        cacheRead: u.prompt_tokens_details?.cached_tokens ?? 0,
        cacheWrite: 0,
      },
    },
    durationMs: totalMs,
  });
  return out;
}

// ─── Agent loop ──────────────────────────────────────────────────────────────

async function runAgentLoop() {
  while (true) {
    // Honor pause: idle in 500ms ticks, no model calls.
    if (paused) {
      if (status !== 'paused') {
        status = 'paused';
        trace({ type: 'agent_paused' });
        broadcast();
      }
      await new Promise(r => setTimeout(r, 500));
      continue;
    } else if (status === 'paused') {
      status = 'running';
      trace({ type: 'agent_resumed' });
      broadcast();
    }

    const turn = actionLog.length + 1;
    trace({ type: 'turn_start', turn });
    const viewport = game.getViewport();

    // Compact before adding the next user turn so the model sees the summary
    // as established prior context rather than mid-turn.
    if (memory.shouldCompact()) {
      status = 'compacting';
      broadcast();
      const tokensBefore = memory.estimateTokens();
      trace({ type: 'compaction_start', reason: 'threshold', tokensBefore });
      try {
        const result = await memory.compactWith(callSummarizer);
        if (result) {
          trace({
            type: 'compaction_end',
            reason: 'threshold',
            result: {
              summary: memory.summary,
              tokensBefore,
              tokensAfter: memory.estimateTokens(),
              droppedMessages: result.compacted,
            },
          });
        } else {
          trace({ type: 'compaction_end', reason: 'threshold', aborted: true });
        }
      } catch (err: any) {
        trace({
          type: 'compaction_end',
          reason: 'threshold',
          aborted: false,
          willRetry: false,
          errorMessage: String(err?.message ?? err),
        });
      }
    }

    // Snapshot memory length so we can roll back cleanly if this turn no-ops.
    const memoryDepthBefore = memory.messages.length;
    memory.addUserText(buildPickFirstPrompt(viewport, game.inventory));

    let action: { action: string; detail: string; tool_call_id: string } | null = null;
    status = 'thinking';
    broadcast();

    try {
      const { message, usage, totalMs } = await callModel();
      lastUsage = usage;
      lastTotalMs = totalMs;
      memory.addAssistant(message);
      lastThinking = message.reasoning_content ?? null;

      const tc = message.tool_calls?.[0];
      if (tc) {
        let args: any = {};
        try { args = JSON.parse(tc.function.arguments); } catch {}
        if (
          (tc.function.name === 'move' || tc.function.name === 'pick') &&
          ['up', 'down', 'left', 'right'].includes(args.direction)
        ) {
          action = { action: tc.function.name, detail: args.direction, tool_call_id: tc.id };
        }
      }
      status = 'running';
    } catch (err: any) {
      status = 'fallback';
      trace({
        type: 'error',
        source: 'inference',
        message: String(err?.message ?? err),
      });
    }

    if (!action) {
      // No usable tool_call. Don't move the agent, don't pollute actionLog.
      // Roll memory back to pre-turn state so the next attempt isn't biased
      // by an empty assistant turn or an orphaned user message.
      while (memory.messages.length > memoryDepthBefore) memory.popLast();
      fallbackCount++;
      trace({ type: 'fallback_action', turn, totalFallbacks: fallbackCount });
      trace({ type: 'turn_end', turn });
      broadcast();
      await new Promise((r) => setTimeout(r, DEFAULTS.LOOP_DELAY_MS));
      continue;
    }

    const argsObj = { direction: action.detail };
    trace({
      type: 'tool_execution_start',
      toolCallId: action.tool_call_id,
      toolName: action.action,
      args: argsObj,
    });

    let toolResult: string;
    let success = true;
    try {
      const result = game.executeAction(action) as { pickedUp?: string | null } | undefined;
      const picked = result?.pickedUp ? ` — picked up ${result.pickedUp}` : '';
      toolResult = `${action.action} ${action.detail} succeeded${picked}`;
    } catch (err: any) {
      toolResult = `${action.action} ${action.detail} failed: ${err.message}`;
      success = false;
    }
    memory.addToolResult(action.tool_call_id, toolResult);

    actionLog.push({ action: action.action, detail: action.detail, success });

    trace({
      type: 'tool_execution_end',
      toolCallId: action.tool_call_id,
      toolName: action.action,
      result: toolResult,
      isError: !success,
    });

    trace({ type: 'turn_end', turn });
    broadcast();

    await new Promise(r => setTimeout(r, DEFAULTS.LOOP_DELAY_MS));
  }
}

// ─── HTTP server ─────────────────────────────────────────────────────────────

serve({
  port: PORT,
  idleTimeout: 240, // long-lived SSE connections
  async fetch(req) {
    const url = new URL(req.url);

    if (url.pathname === '/health') {
      return new Response('ok');
    }

    // ─── Programmatic test/debug API ────────────────────────────────────────

    // GET /api/state — one-shot full state JSON (same shape as SSE events).
    if (url.pathname === '/api/state') {
      return new Response(JSON.stringify(buildState()), {
        headers: { 'Content-Type': 'application/json' },
      });
    }

    // GET /api/diag — endpoint health + slot status from the upstream llama.cpp.
    if (url.pathname === '/api/diag') {
      const props = await fetch(`${ENDPOINT}/props`).then(r => r.json()).catch(e => ({ error: String(e) }));
      const slots = await fetch(`${ENDPOINT}/slots`).then(r => r.json()).catch(e => ({ error: String(e) }));
      return new Response(JSON.stringify({
        endpoint: ENDPOINT,
        model: remoteModelId,
        serverCtx,
        memory: {
          tokenBudget: memory.tokenBudget,
          triggerCompactTokens: memory.triggerCompactTokens,
          keepRecentTokens: memory.keepRecentTokens,
          messages: memory.messages.length,
          estimateTokens: memory.estimateTokens(),
          compactions: memory.compactions,
          summaryLen: memory.summary?.length ?? 0,
        },
        cacheRetries,
        cacheCooldown,
        upstream: { props, slots },
      }, null, 2), { headers: { 'Content-Type': 'application/json' } });
    }

    // GET /api/memory — full conversation buffer for inspection.
    if (url.pathname === '/api/memory') {
      return new Response(JSON.stringify({
        systemPrompt: memory.systemPrompt,
        summary: memory.summary,
        messages: memory.messages,
      }, null, 2), { headers: { 'Content-Type': 'application/json' } });
    }

    // POST /api/pause — stop hitting the model. Loop idles, world stays intact.
    if (url.pathname === '/api/pause' && req.method === 'POST') {
      paused = true;
      return new Response(JSON.stringify({ paused: true }), {
        headers: { 'Content-Type': 'application/json' },
      });
    }

    // POST /api/resume — resume the loop.
    if (url.pathname === '/api/resume' && req.method === 'POST') {
      paused = false;
      return new Response(JSON.stringify({ paused: false }), {
        headers: { 'Content-Type': 'application/json' },
      });
    }

    // POST /api/reset — wipe world + memory + actionLog and start fresh.
    if (url.pathname === '/api/reset' && req.method === 'POST') {
      Object.assign(game, new SinglePlayerGame());
      memory.messages = [];
      memory.summary = null;
      memory.compactions = 0;
      actionLog.length = 0;
      fallbackCount = 0;
      lastUsage = null;
      lastTotalMs = 0;
      lastThinking = null;
      cacheRetries = 0;
      cacheCooldown = 0;
      trace({ type: 'world_reset' });
      broadcast();
      return new Response(JSON.stringify({ ok: true }), {
        headers: { 'Content-Type': 'application/json' },
      });
    }

    if (url.pathname === '/stream') {
      const stream = new ReadableStream<Uint8Array>({
        start(controller) {
          subscribers.add(controller);
          trace({ type: 'viewer_connect', viewers: subscribers.size });
          // Send immediate snapshot so clients don't wait a full turn.
          controller.enqueue(enc.encode(`data: ${JSON.stringify(buildState())}\n\n`));
        },
        cancel(controller) {
          subscribers.delete(controller);
          trace({ type: 'viewer_disconnect', viewers: subscribers.size });
        },
      });
      return new Response(stream as any, {
        headers: {
          'Content-Type': 'text/event-stream',
          'Cache-Control': 'no-cache',
          Connection: 'keep-alive',
        },
      });
    }

    if (url.pathname === '/' || url.pathname === '/index.html') {
      const data = await readFile(import.meta.dir + '/index.html');
      return new Response(data, { headers: { 'Content-Type': 'text/html' } });
    }

    return new Response('Not found', { status: 404 });
  },
});

trace({ type: 'serve_listening', port: PORT, url: `http://localhost:${PORT}` });

// ─── Boot ────────────────────────────────────────────────────────────────────

await connect();
runAgentLoop().catch((err) => {
  trace({
    type: 'agent_end',
    reason: 'crash',
    errorMessage: String(err?.message ?? err),
  });
  status = 'crashed';
  broadcast();
});