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

const PORT = 8000;
const MODEL_ID = 'kimi-agent';
const TIMEOUT_MS = 15 * 60 * 1000;

const WORKSPACE = '/workspace';
const CONV_ROOT = path.join(WORKSPACE, 'conversations');
const STATE_DIR = path.join(WORKSPACE, '.kimi-agent');
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 });

// --- persistent conversation -> session map ---------------------------------
// key = hash(history-so-far)  ->  { 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 => 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');
}

// --- kimi invocation --------------------------------------------------------
// Returns { text, sessionId }. Parses --output-format stream-json:
//   {"role":"assistant","content":"..."}
//   {"role":"meta","type":"session.resume_hint","session_id":"session_..."}
function runKimi({ prompt, cwd, resumeId }) {
  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 stdout = '';
    let stderr = '';
    child.stdout.on('data', d => { stdout += d; });
    child.stderr.on('data', d => { stderr += d; });

    child.on('error', reject);
    child.on('close', code => {
      const parts = [];
      let sessionId = null;
      for (const line of stdout.split('\n')) {
        const t = line.trim();
        if (!t) continue;
        let obj;
        try { obj = JSON.parse(t); } catch { continue; }
        if (obj.role === 'assistant' && obj.content) parts.push(obj.content);
        if (obj.type === 'session.resume_hint' && obj.session_id) sessionId = obj.session_id;
      }
      const text = parts.join('').trim();
      if (!text && code !== 0) {
        reject(new Error(`kimi exited ${code}: ${stderr.slice(0, 2000)}`));
      } else {
        resolve({ text, sessionId });
      }
    });
  });
}

// Decide session/dir, run kimi, and record the forward mapping.
async function handleTurn(messages) {
  const turns = convTurns(messages);
  // find the last user turn = the new prompt; everything before it is prior history
  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 newPrompt = textOf(turns[lastUserIdx]);
  const prior = turns.slice(0, lastUserIdx);

  let convId;
  let dir;
  let resumeId = null;
  let prompt = newPrompt;

  if (prior.length === 0) {
    // brand-new conversation
    convId = crypto.randomUUID();
    dir = path.join(CONV_ROOT, convId);
  } else {
    const entry = sessionMap[historyKey(prior)];
    if (entry) {
      // known conversation -> resume the same kimi session in its own dir
      convId = entry.convId;
      dir = entry.dir;
      resumeId = entry.sessionId;
    } else {
      // lost mapping (restart / edited history): reseed a fresh session with
      // the full transcript so continuity is preserved
      convId = crypto.randomUUID();
      dir = path.join(CONV_ROOT, convId);
      prompt = renderTranscript(turns.slice(0, lastUserIdx + 1));
    }
  }

  fs.mkdirSync(dir, { recursive: true });
  const { text, sessionId } = await runKimi({ prompt, cwd: dir, resumeId });

  // store forward mapping: next request's prior history == these turns + reply
  const forward = turns.slice(0, lastUserIdx + 1).concat([{ role: 'assistant', content: text }]);
  sessionMap[historyKey(forward)] = {
    convId,
    sessionId: sessionId || resumeId,
    dir,
    ts: Date.now(),
  };
  persistMap();

  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',
    }],
  };
}

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;
      }

      try {
        const text = await handleTurn(parsed.messages || []);

        if (parsed.stream) {
          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(`data: ${JSON.stringify({
            id, object: 'chat.completion.chunk', created, model: MODEL_ID,
            choices: [{ index: 0, delta: { role: 'assistant', content: text }, finish_reason: null }],
          })}\n\n`);
          res.write(`data: ${JSON.stringify({
            id, object: 'chat.completion.chunk', created, model: MODEL_ID,
            choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
          })}\n\n`);
          res.write('data: [DONE]\n\n');
          res.end();
        } else {
          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(`kimi-agent wrapper listening on :${PORT}`));