const http = require('http');
const fs = require('fs');
const path = require('path');
const crypto = require('crypto');
const { spawn } = require('child_process');

const PORT = 8010;
const MODEL_ID = 'adolf';
const TIMEOUT_MS = 15 * 60 * 1000;

const WORKSPACE = '/workspace';
const CONV_ROOT = path.join(WORKSPACE, 'conversations');
const STATE_DIR = path.join(WORKSPACE, '.adolf-llm');
const MAP_FILE = path.join(STATE_DIR, 'sessions.json');
const MAX_ENTRIES = 1000; // prune oldest beyond this

fs.mkdirSync(CONV_ROOT, { recursive: true });
fs.mkdirSync(STATE_DIR, { recursive: true });

// ---------------------------------------------------------------------------
// Shared MCP layer (Gate 1). Kimi Code CLI has NO `--mcp-config-file` flag and
// no `kimi mcp` subcommand; it auto-discovers a project-root `.mcp.json` by
// walking up from its cwd to the nearest `.git` (falling back to cwd itself
// when none is found). So we drop a `.mcp.json` into each session's working
// directory before spawning kimi.
//
// Single source of truth: `/shared-mcp.json` (mounted read-only from the repo
// root's `shared-mcp.json`, the same file P6 wires into OpenClaw's own
// `mcp.servers` registry). Adding a server is then a one-file change — no
// server list is hardcoded here anymore.
//
// Gate-1 transport finding (P5, verified by decompiling the installed
// @moonshot-ai/kimi-code package, packages/agent-core/src/config/schema.ts's
// McpServerConfigSchema): Kimi's own field name for remote MCP servers is
// `transport` (literal "stdio" | "http" | "sse"), not `type`. When `transport`
// is omitted, Kimi's config preprocessor infers it from shape: `command` ->
// "stdio", `url` -> "http" (never "sse" — sse requires an explicit
// `transport: "sse"`). It does NOT recognize a `type` key at all; unknown keys
// are silently stripped by the (non-strict) zod schema.
// OpenClaw's own canonical `mcp.servers` schema (docs/gateway/
// configuration-reference.md) uses different literals for the same
// transport: `transport: "streamable-http"` or `"sse"`, with `type: "http"`
// documented as a *CLI-native alias* that `openclaw mcp set` / `openclaw
// doctor --fix` normalize into canonical `transport: "streamable-http"`.
// So the two consumers disagree on the literal value for HTTP streaming
// ("http" vs "streamable-http") under the same field name `transport` --
// writing `transport` explicitly in shared-mcp.json would satisfy at most one
// side. `type: "http"` is the one shape both sides tolerate today: Kimi
// ignores the unrecognized `type` key and correctly infers transport "http"
// from the `url` field alone; OpenClaw recognizes `type` as its documented
// alias and normalizes it on its own terms (P6 concern, not touched here).
// Hence shared-mcp.json intentionally keeps `"type": "http"` for both cognee
// and openclaw-tools rather than switching to `transport`.
let SHARED_MCP_SERVERS = {};
try {
  const raw = fs.readFileSync('/shared-mcp.json', 'utf8');
  SHARED_MCP_SERVERS = JSON.parse(raw).mcpServers || {};
} catch (err) {
  console.error(`shared-mcp.json not loaded (${err.message}); sessions will get no shared MCP servers`);
}

function writeMcpConfig(dir) {
  const cfg = { mcpServers: SHARED_MCP_SERVERS };
  fs.writeFileSync(path.join(dir, '.mcp.json'), JSON.stringify(cfg, null, 2));
}

// ---------------------------------------------------------------------------
// Cognee auto-memory hooks (§3.2). Both are STUBS today — the cognee service is
// P4 and does not exist yet. They are deliberately non-blocking: a turn must
// never wait on (or fail because of) memory. Wire the real cognee-mcp calls in
// P4 and the rest of the turn pipeline stays unchanged.

// Pre-turn auto-retrieve: returns a string of relevant memories to inject ahead
// of the user prompt, or null for "nothing to inject".
// TODO(P4): call cognee memory_search / cognee.search() scoped by chatId.
async function cogneeSearch(_query, _chatId) {
  return null;
}

// Post-turn auto-ingest: fire-and-forget; never awaited by the turn path.
// TODO(P4): async cognee.add(user + assistant) via cognee-mcp, scoped by chatId.
async function cogneeAdd(_userText, _assistantText, _chatId) {
  return;
}

// ---------------------------------------------------------------------------
// Persistent conversation -> Kimi session map.
// Primary key: `chat:<chat_id>` parsed from OpenClaw's "Conversation info" block
// (Gate 2). Fallback key: `hist:<sha256(prior history)>` when no chat_id is
// present (e.g. webchat surface / future OpenClaw layout change).
// value = { convId, sessionId, dir, ts }
let sessionMap = {};
try {
  sessionMap = JSON.parse(fs.readFileSync(MAP_FILE, 'utf8'));
} catch {
  sessionMap = {};
}

let writeQueue = Promise.resolve();
function persistMap() {
  // prune to the MAX_ENTRIES most-recently-used before writing
  const keys = Object.keys(sessionMap);
  if (keys.length > MAX_ENTRIES) {
    keys.sort((a, b) => (sessionMap[a].ts || 0) - (sessionMap[b].ts || 0));
    for (const k of keys.slice(0, keys.length - MAX_ENTRIES)) delete sessionMap[k];
  }
  const snapshot = JSON.stringify(sessionMap);
  writeQueue = writeQueue.then(
    () => fs.promises.writeFile(MAP_FILE, snapshot),
    () => fs.promises.writeFile(MAP_FILE, snapshot),
  );
  return writeQueue;
}

// ---------------------------------------------------------------------------
// Message helpers.
function textOf(msg) {
  const c = msg.content;
  if (Array.isArray(c)) return c.map(p => (typeof p === 'string' ? p : p.text || '')).join('\n');
  return c == null ? '' : String(c);
}

// only user/assistant turns define conversation identity (system is constant)
function convTurns(messages) {
  return messages.filter(m => m.role === 'user' || m.role === 'assistant');
}

function historyKey(turns) {
  const norm = turns.map(m => ({ role: m.role, text: textOf(m).trim() }));
  return crypto.createHash('sha256').update(JSON.stringify(norm)).digest('hex');
}

function renderTranscript(turns) {
  return turns
    .map(m => `${m.role === 'assistant' ? 'Assistant' : 'User'}: ${textOf(m)}`)
    .join('\n\n');
}

// ---------------------------------------------------------------------------
// Gate 2 — parse the stable chat_id out of OpenClaw's untrusted-metadata block.
// OpenClaw injects, into the *user-role* content, a block that looks like:
//   Conversation info (untrusted metadata):
//   ```json
//   { "chat_id": "matrix:!room:server", "message_id": "...", ... }
//   ```
// We grep on the label string (never a fixed line offset) then pull the first
// balanced JSON object after it and read chat_id. Robust to edited/truncated
// history, which is exactly why it beats a history hash for the common case.
const CONV_INFO_LABEL = 'Conversation info (untrusted metadata):';

function extractBalancedJson(str, from) {
  const start = str.indexOf('{', from);
  if (start === -1) return null;
  let depth = 0;
  let inStr = false;
  let esc = false;
  for (let i = start; i < str.length; i++) {
    const ch = str[i];
    if (inStr) {
      if (esc) esc = false;
      else if (ch === '\\') esc = true;
      else if (ch === '"') inStr = false;
      continue;
    }
    if (ch === '"') inStr = true;
    else if (ch === '{') depth++;
    else if (ch === '}') {
      depth--;
      if (depth === 0) return str.slice(start, i + 1);
    }
  }
  return null;
}

function extractChatId(userMsg) {
  const text = textOf(userMsg);
  const at = text.indexOf(CONV_INFO_LABEL);
  if (at === -1) return null;
  const jsonStr = extractBalancedJson(text, at + CONV_INFO_LABEL.length);
  if (!jsonStr) return null;
  try {
    const obj = JSON.parse(jsonStr);
    const id = obj.chat_id;
    return typeof id === 'string' && id ? id : null;
  } catch {
    return null;
  }
}

// ---------------------------------------------------------------------------
// Gate 3 — media. Persist inbound image parts into the session dir and return
// relative path references; the CLI autonomously calls its built-in
// ReadMediaFile tool on referenced paths (no flag/placeholder syntax needed).
const MIME_EXT = {
  'image/png': 'png',
  'image/jpeg': 'jpg',
  'image/jpg': 'jpg',
  'image/webp': 'webp',
  'image/gif': 'gif',
  'image/bmp': 'bmp',
  'image/heic': 'heic',
  'image/heif': 'heif',
};

function extFromMime(mime) {
  return MIME_EXT[(mime || '').toLowerCase()] || 'img';
}

// Persist one image_url part; returns "./img_N.ext" or null if it couldn't.
async function persistImage(url, dir, n) {
  if (typeof url !== 'string' || !url) return null;
  if (url.startsWith('data:')) {
    const m = /^data:([^;,]+)?(;base64)?,(.*)$/s.exec(url);
    if (!m) return null;
    const mime = m[1] || 'application/octet-stream';
    const isB64 = !!m[2];
    const ext = extFromMime(mime);
    const name = `img_${n}.${ext}`;
    const buf = isB64
      ? Buffer.from(m[3], 'base64')
      : Buffer.from(decodeURIComponent(m[3]), 'utf8');
    fs.writeFileSync(path.join(dir, name), buf);
    return `./${name}`;
  }
  // Remote URL: best-effort fetch so the CLI gets a local path to ReadMediaFile.
  // On any failure fall back to handing the raw URL to the model as text.
  if (/^https?:\/\//i.test(url)) {
    try {
      const resp = await fetch(url);
      if (!resp.ok) return url;
      const mime = resp.headers.get('content-type') || '';
      const ext = extFromMime(mime.split(';')[0].trim());
      const name = `img_${n}.${ext}`;
      const buf = Buffer.from(await resp.arrayBuffer());
      fs.writeFileSync(path.join(dir, name), buf);
      return `./${name}`;
    } catch {
      return url;
    }
  }
  return null;
}

// Build the prompt for the current user turn: join text parts, persist any
// image parts, append path references. `injectedMemory` (from cogneeSearch) is
// prepended when present.
async function buildPrompt(userMsg, dir, injectedMemory) {
  const content = userMsg.content;
  const textParts = [];
  const imageRefs = [];

  if (Array.isArray(content)) {
    let n = 0;
    for (const part of content) {
      if (typeof part === 'string') {
        textParts.push(part);
      } else if (part && (part.type === 'text' || typeof part.text === 'string')) {
        textParts.push(part.text || '');
      } else if (part && part.type === 'image_url' && part.image_url && part.image_url.url) {
        n++;
        const ref = await persistImage(part.image_url.url, dir, n);
        if (ref) imageRefs.push(ref);
      } else if (part && part.type === 'input_image' && (part.image_url || part.url)) {
        n++;
        const u = typeof part.image_url === 'string' ? part.image_url : part.url;
        const ref = await persistImage(u, dir, n);
        if (ref) imageRefs.push(ref);
      }
    }
  } else {
    textParts.push(content == null ? '' : String(content));
  }

  let prompt = textParts.join('\n');
  if (imageRefs.length) {
    prompt += '\n\n' + imageRefs.map(r => `See attached image: ${r}`).join('\n');
  }
  if (injectedMemory) {
    prompt = `Relevant memories (retrieved automatically):\n${injectedMemory}\n\n---\n\n${prompt}`;
  }
  return prompt;
}

// ---------------------------------------------------------------------------
// Kimi invocation with REAL streaming. Parses `--output-format stream-json`
// incrementally: each complete stdout line is one JSON object.
//   {"role":"assistant","content":"..."}                 -> emit as a delta
//   {"type":"session.resume_hint","session_id":"..."}    -> capture session id
// onDelta(chunk) is called per assistant content fragment as it arrives.
// Resolves { text, sessionId } once the process closes.
function runKimi({ prompt, cwd, resumeId, onDelta }) {
  return new Promise((resolve, reject) => {
    const args = [];
    if (resumeId) args.push('-r', resumeId);
    args.push('-p', prompt, '--output-format', 'stream-json');

    const child = spawn('kimi', args, { cwd, timeout: TIMEOUT_MS });

    let buf = '';
    let stderr = '';
    const parts = [];
    let sessionId = null;

    function handleLine(line) {
      const t = line.trim();
      if (!t) return;
      let obj;
      try { obj = JSON.parse(t); } catch { return; }
      if (obj.role === 'assistant' && typeof obj.content === 'string' && obj.content) {
        parts.push(obj.content);
        if (onDelta) onDelta(obj.content);
      }
      if (obj.type === 'session.resume_hint' && obj.session_id) sessionId = obj.session_id;
    }

    child.stdout.on('data', d => {
      buf += d;
      let nl;
      while ((nl = buf.indexOf('\n')) !== -1) {
        const line = buf.slice(0, nl);
        buf = buf.slice(nl + 1);
        handleLine(line);
      }
    });
    child.stderr.on('data', d => { stderr += d; });

    child.on('error', reject);
    child.on('close', code => {
      if (buf) handleLine(buf); // flush any trailing partial line
      const text = parts.join('').trim();
      if (!text && code !== 0) {
        reject(new Error(`kimi exited ${code}: ${stderr.slice(0, 2000)}`));
      } else {
        resolve({ text, sessionId });
      }
    });
  });
}

// ---------------------------------------------------------------------------
// One turn: resolve session (chat_id primary, history-hash fallback), persist
// media + .mcp.json, run kimi (streaming through onDelta), record the mapping,
// and fire the async cognee ingest. Returns { text }.
async function handleTurn(messages, onDelta) {
  const turns = convTurns(messages);
  let lastUserIdx = -1;
  for (let i = turns.length - 1; i >= 0; i--) {
    if (turns[i].role === 'user') { lastUserIdx = i; break; }
  }
  if (lastUserIdx === -1) throw new Error('no user message found');

  const userMsg = turns[lastUserIdx];
  const prior = turns.slice(0, lastUserIdx);
  const chatId = extractChatId(userMsg);

  // Resolve the session key + working dir + resume id.
  let key;
  let convId;
  let dir;
  let resumeId = null;
  let reseed = false; // when true, prompt with the full transcript to rebuild continuity

  if (chatId) {
    key = `chat:${chatId}`;
    const entry = sessionMap[key];
    if (entry) {
      convId = entry.convId;
      dir = entry.dir;
      resumeId = entry.sessionId;
    } else {
      convId = crypto.randomUUID();
      dir = path.join(CONV_ROOT, convId);
      // First time we see this chat_id but history exists (server restart / lost
      // map): reseed the fresh session with the transcript so context survives.
      reseed = prior.length > 0;
    }
  } else {
    // Fallback: no chat_id -> forward history-hash mapping (kimi-agent style).
    if (prior.length === 0) {
      convId = crypto.randomUUID();
      dir = path.join(CONV_ROOT, convId);
    } else {
      const entry = sessionMap[`hist:${historyKey(prior)}`];
      if (entry) {
        convId = entry.convId;
        dir = entry.dir;
        resumeId = entry.sessionId;
      } else {
        convId = crypto.randomUUID();
        dir = path.join(CONV_ROOT, convId);
        reseed = true;
      }
    }
  }

  fs.mkdirSync(dir, { recursive: true });
  writeMcpConfig(dir); // Gate 1: shared MCP via project-root .mcp.json

  // Pre-turn auto-retrieve (STUB no-op today; never blocks meaningfully).
  const injectedMemory = await cogneeSearch(textOf(userMsg), chatId);

  let prompt;
  if (reseed) {
    // Rebuild the whole conversation for a fresh session, plus current media/memory.
    const base = renderTranscript(turns.slice(0, lastUserIdx + 1));
    const media = await buildPrompt(userMsg, dir, injectedMemory);
    // buildPrompt already includes the current user text; for reseed we want the
    // transcript to carry it, so only append image refs / injected memory extras.
    prompt = base;
    const extra = media.replace(textOf(userMsg), '').trim();
    if (extra) prompt += `\n\n${extra}`;
  } else {
    prompt = await buildPrompt(userMsg, dir, injectedMemory);
  }

  const { text, sessionId } = await runKimi({ prompt, cwd: dir, resumeId, onDelta });

  // Record the forward mapping.
  const entry = { convId, sessionId: sessionId || resumeId, dir, ts: Date.now() };
  if (chatId) {
    sessionMap[`chat:${chatId}`] = entry;
  } else {
    const forward = turns.slice(0, lastUserIdx + 1).concat([{ role: 'assistant', content: text }]);
    sessionMap[`hist:${historyKey(forward)}`] = entry;
  }
  persistMap();

  // Post-turn auto-ingest (STUB no-op today; fire-and-forget, never awaited).
  cogneeAdd(textOf(userMsg), text, chatId).catch(() => {});

  return { text };
}

// ---------------------------------------------------------------------------
// OpenAI-compatible HTTP surface.
function completionBody(text) {
  return {
    id: `chatcmpl-${Date.now()}`,
    object: 'chat.completion',
    created: Math.floor(Date.now() / 1000),
    model: MODEL_ID,
    choices: [{
      index: 0,
      message: { role: 'assistant', content: text },
      finish_reason: 'stop',
    }],
  };
}

function sseChunk(id, created, delta, finishReason) {
  return `data: ${JSON.stringify({
    id, object: 'chat.completion.chunk', created, model: MODEL_ID,
    choices: [{ index: 0, delta, finish_reason: finishReason }],
  })}\n\n`;
}

const server = http.createServer((req, res) => {
  if (req.method === 'GET' && req.url === '/v1/models') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      object: 'list',
      data: [{ id: MODEL_ID, object: 'model', owned_by: 'moonshot' }],
    }));
    return;
  }

  if (req.method === 'POST' && req.url === '/v1/chat/completions') {
    let body = '';
    req.on('data', d => { body += d; });
    req.on('end', async () => {
      let parsed;
      try {
        parsed = JSON.parse(body);
      } catch {
        res.writeHead(400, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: 'invalid JSON body' }));
        return;
      }

      const messages = parsed.messages || [];

      if (parsed.stream) {
        // Real streaming: open SSE, emit role chunk, then forward kimi deltas.
        res.writeHead(200, {
          'Content-Type': 'text/event-stream',
          'Cache-Control': 'no-cache',
          Connection: 'keep-alive',
        });
        const id = `chatcmpl-${Date.now()}`;
        const created = Math.floor(Date.now() / 1000);
        res.write(sseChunk(id, created, { role: 'assistant' }, null));
        try {
          await handleTurn(messages, delta => {
            res.write(sseChunk(id, created, { content: delta }, null));
          });
          res.write(sseChunk(id, created, {}, 'stop'));
          res.write('data: [DONE]\n\n');
          res.end();
        } catch (err) {
          // Headers already sent — surface the error inside the stream.
          res.write(sseChunk(id, created, { content: `\n[error: ${String(err.message || err)}]` }, 'stop'));
          res.write('data: [DONE]\n\n');
          res.end();
        }
        return;
      }

      try {
        const { text } = await handleTurn(messages, null);
        res.writeHead(200, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify(completionBody(text)));
      } catch (err) {
        res.writeHead(500, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: String(err.message || err) }));
      }
    });
    return;
  }

  res.writeHead(404, { 'Content-Type': 'application/json' });
  res.end(JSON.stringify({ error: 'not found' }));
});

server.listen(PORT, () => console.log(`adolf-llm wrapper listening on :${PORT}`));