const express = require("express");
const { spawn, execSync } = require("child_process");
const fs = require("fs");
const path = require("path");
const multer = require("multer");

const app = express();
const PORT = process.env.PORT || 3000;

const INDIGO_DIR = path.resolve(__dirname, "..", "Indigo");
const DATA_DIR = path.join(INDIGO_DIR, "data");
const RUNS_DIR = path.join(INDIGO_DIR, "runs");
const CHAT_DIR = path.join(__dirname, "chats");
const CONFIG_PATH = path.join(__dirname, "config.json");

function findPython() {
  try {
    const cfg = loadConfig();
    if (cfg.pythonPath && fs.existsSync(cfg.pythonPath)) return cfg.pythonPath;
  } catch (e) {}
  try {
    return execSync('python -c "import sys; print(sys.executable)"', { encoding: "utf-8", timeout: 5000 }).trim();
  } catch (e) {
    return "python";
  }
}
const PYTHON = process.env.PYTHON_PATH || findPython();
console.log("PYTHON:", PYTHON);

if (!fs.existsSync(CHAT_DIR)) fs.mkdirSync(CHAT_DIR, { recursive: true });

app.use(express.json({ limit: "10mb" }));
app.use(express.static(path.join(__dirname, "public")));

const upload = multer({ dest: path.join(__dirname, "uploads") });

function loadConfig() {
  const defaults = {
    model: null,
    temperature: 0.8,
    topK: 40,
    topP: 0.95,
    maxNew: 200,
    guard: true,
    device: "auto",
  };
  if (fs.existsSync(CONFIG_PATH)) {
    return { ...defaults, ...JSON.parse(fs.readFileSync(CONFIG_PATH, "utf-8")) };
  }
  return defaults;
}

function saveConfig(cfg) {
  fs.writeFileSync(CONFIG_PATH, JSON.stringify(cfg, null, 2), "utf-8");
}

function listCheckpoints() {
  const results = [];
  if (!fs.existsSync(RUNS_DIR)) return results;
  function scanDir(dir, label) {
    for (const f of fs.readdirSync(dir)) {
      if (f.endsWith(".safetensors") && !f.includes("optimizer")) {
        const metaPath = path.join(dir, f.replace(".safetensors", "_meta.json"));
        let meta = {};
        if (fs.existsSync(metaPath)) {
          meta = JSON.parse(fs.readFileSync(metaPath, "utf-8"));
        }
        results.push({
          name: `${label}/${f}`,
          path: path.join(dir, f),
          step: meta.step || 0,
          val_loss: meta.val_loss || null,
          config: meta.config || {},
          tokenizer: meta.tokenizer || { type: "char" },
          size: fs.statSync(path.join(dir, f)).size,
        });
      }
    }
  }
  for (const run of fs.readdirSync(RUNS_DIR)) {
    const runDir = path.join(RUNS_DIR, run);
    if (!fs.statSync(runDir).isDirectory()) continue;
    const ckptDir = path.join(runDir, "ckpt");
    if (fs.existsSync(ckptDir)) scanDir(ckptDir, run);
    else scanDir(runDir, run);
  }
  return results;
}

const TEXT_EXTS = [".txt", ".md", ".py", ".js", ".ts", ".json", ".csv", ".log", ".html", ".css", ".xml", ".yaml", ".yml", ".toml", ".cfg", ".ini", ".sh", ".bat", ".sql", ".r", ".java", ".c", ".cpp", ".h", ".hpp", ".rs", ".go", ".rb", ".php", ".swift", ".kt", ".scala", ".lua", ".pl", ".tex", ".bib"];

// API: list models
app.get("/api/models", (req, res) => {
  res.json(listCheckpoints());
});

// API: config
app.get("/api/config", (req, res) => {
  res.json(loadConfig());
});

app.post("/api/config", (req, res) => {
  const cfg = { ...loadConfig(), ...req.body };
  saveConfig(cfg);
  res.json(cfg);
});

// API: MCP server check
app.get("/api/mcp-check", (req, res) => {
  const mcpScript = path.join(__dirname, "mcp-server.js");
  if (!fs.existsSync(mcpScript)) return res.json({ ok: false, error: "mcp-server.js not found" });

  const initReq = JSON.stringify({ jsonrpc: "2.0", id: 1, method: "initialize", params: { protocolVersion: "2024-11-05", capabilities: {}, clientInfo: { name: "indigo-chat", version: "1.0.0" } } });
  const listReq = JSON.stringify({ jsonrpc: "2.0", id: 2, method: "tools/list", params: {} });
  const input = initReq + "\n" + listReq + "\n";

  const proc = spawn("node", [mcpScript], { cwd: __dirname, timeout: 10000 });
  let stdout = "";
  let stderr = "";
  proc.stdout.on("data", (d) => { stdout += d.toString(); });
  proc.stderr.on("data", (d) => { stderr += d.toString(); });
  proc.stdin.write(input);
  proc.stdin.end();

  proc.on("close", () => {
    try {
      const lines = stdout.trim().split("\n");
      let tools = [];
      for (const line of lines) {
        try {
          const d = JSON.parse(line);
          if (d.id === 2 && d.result && d.result.tools) tools = d.result.tools;
        } catch (e) {}
      }
      res.json({ ok: true, tools: tools.map((t) => t.name), count: tools.length });
    } catch (e) {
      res.json({ ok: false, error: "Failed to parse response", raw: stdout.substring(0, 300), stderr: stderr.substring(0, 300) });
    }
  });
  proc.on("error", (e) => {
    res.json({ ok: false, error: e.message });
  });
});

// API: chat history
app.get("/api/history", (req, res) => {
  const files = fs.readdirSync(CHAT_DIR)
    .filter((f) => f.endsWith(".json"))
    .map((f) => {
      const data = JSON.parse(fs.readFileSync(path.join(CHAT_DIR, f), "utf-8"));
      return { id: f.replace(".json", ""), title: data.title || "Untitled", updated: data.updated };
    })
    .sort((a, b) => (b.updated || 0) - (a.updated || 0));
  res.json(files);
});

app.get("/api/history/:id", (req, res) => {
  const fp = path.join(CHAT_DIR, `${req.params.id}.json`);
  if (!fs.existsSync(fp)) return res.status(404).json({ error: "not found" });
  const data = JSON.parse(fs.readFileSync(fp, "utf-8"));
  res.json(data);
});

app.post("/api/history", (req, res) => {
  const { id, messages, title } = req.body;
  const chatId = id || `chat_${Date.now()}`;
  const data = {
    title: title || messages?.[0]?.content?.substring(0, 50) || "New Chat",
    messages,
    updated: Date.now(),
  };
  fs.writeFileSync(path.join(CHAT_DIR, `${chatId}.json`), JSON.stringify(data, null, 2), "utf-8");
  res.json({ id: chatId });
});

app.delete("/api/history/:id", (req, res) => {
  const fp = path.join(CHAT_DIR, `${req.params.id}.json`);
  if (fs.existsSync(fp)) fs.unlinkSync(fp);
  res.json({ ok: true });
});

// API: upload file (returns content for text files)
app.post("/api/upload", upload.single("file"), (req, res) => {
  if (!req.file) return res.status(400).json({ error: "no file" });
  const cfg = loadConfig();
  const targetDir = path.join(cfg.projectDir || INDIGO_DIR, "uploads");
  if (!fs.existsSync(targetDir)) fs.mkdirSync(targetDir, { recursive: true });
  const target = path.join(targetDir, req.file.originalname);
  fs.copyFileSync(req.file.path, target);
  fs.unlinkSync(req.file.path);
  let content = "";
  try {
    const ext = path.extname(req.file.originalname).toLowerCase();
    if (TEXT_EXTS.includes(ext)) {
      content = fs.readFileSync(target, "utf-8");
      if (content.length > 8000) content = content.substring(0, 8000) + "\n... (dipotong)";
    }
  } catch (e) {}
  res.json({ path: target, name: req.file.originalname, content });
});

// MCP agent: server calls MCP tools, injects context into prompt
function callMCPServer(jsonl) {
  return new Promise((resolve) => {
    const mcpScript = path.join(__dirname, "mcp-server.js");
    if (!fs.existsSync(mcpScript)) return resolve(null);
    const proc = spawn("node", [mcpScript], { cwd: __dirname, timeout: 10000 });
    let stdout = "";
    proc.stdout.on("data", (d) => { stdout += d.toString(); });
    proc.on("close", () => {
      const results = {};
      for (const line of stdout.trim().split("\n")) {
        try {
          const d = JSON.parse(line);
          if (d.id !== undefined && d.result) results[d.id] = d.result;
        } catch (e) {}
      }
      resolve(results);
    });
    proc.on("error", () => resolve(null));
    proc.stdin.write(jsonl);
    proc.stdin.end();
  });
}

const FILE_KEYWORDS = ["file", "folder", "directory", "direktori", "berkas", "list", "search", "cari", "baca", "read", "tampilkan", "show", "kode", "code", "isi", "nama file", "file apa", "isi folder", "file mana", "cari file"];

function needsFileContext(text) {
  const t = text.toLowerCase();
  return FILE_KEYWORDS.some((k) => t.includes(k));
}

function mcpInit() {
  return JSON.stringify({ jsonrpc: "2.0", id: 1, method: "initialize", params: { protocolVersion: "2024-11-05", capabilities: {}, clientInfo: { name: "indigo-chat", version: "1.0.0" } } }) + "\n";
}

async function agentGatherContext(prompt) {
  if (!needsFileContext(prompt)) return null;
  const initLine = mcpInit();
  const p = prompt.toLowerCase();
  const listDir = /(list|isi|sebutkan|tampilkan|show).*(folder|direktori|directory|root|berkas|semua)/.test(p);

  const calls = [];
  if (listDir) {
    calls.push({ id: 2, name: "list_directory", args: {} });
  } else {
    // extract potential filenames (e.g. "train.py", "kamus_id.txt")
    const fileRefs = p.match(/[a-z0-9_\\-]+\.[a-z0-9]{1,5}/g) || [];
    if (fileRefs.length > 0) {
      for (let i = 0; i < Math.min(fileRefs.length, 2); i++) {
        calls.push({ id: 2 + i, name: "search_files", args: { query: fileRefs[i].replace(/^[^a-z0-9]/g, "") } });
      }
    } else {
      // fallback: search by a keyword from the prompt
      const kw = p.split(/[^a-z0-9]+/).filter((w) => w.length >= 3).pop();
      if (kw) calls.push({ id: 2, name: "search_files", args: { query: kw } });
    }
  }

  if (calls.length === 0) return null;
  const jsonl = initLine + calls.map((c) => JSON.stringify({ jsonrpc: "2.0", id: c.id, method: "tools/call", params: { name: c.name, arguments: c.args } })).join("\n") + "\n";
  const resp = await callMCPServer(jsonl);
  if (!resp) return null;

  const sections = [];
  const readCalls = [];
  let nextId = calls.length + 2;

  for (const c of calls) {
    const result = resp[c.id];
    if (!result) continue;
    let text = result.content?.[0]?.text || "";
    if (c.name === "list_directory") {
      sections.push(`[Konteks: konten folder]\n${text.substring(0, 3000)}`);
      continue;
    }
    let items = [];
    try { items = JSON.parse(text); } catch (e) { continue; }
    if (!Array.isArray(items)) continue;
    sections.push(`[Hasil pencarian "${c.args.query}"]\n` + items.map((f) => f.path).join("\n").substring(0, 1000));
    // antrekan pembacaan file (batch dalam satu spawn MCP)
    for (const item of items.filter((f) => !f.isDir).slice(0, 2)) {
      readCalls.push({ id: nextId, path: item.path, name: item.name });
      nextId++;
    }
  }

  // Baca semua file dalam satu sesi MCP (satu spawn, bukan per-file)
  if (readCalls.length > 0) {
    const readJsonl = initLine + readCalls.map((r) => JSON.stringify({ jsonrpc: "2.0", id: r.id, method: "tools/call", params: { name: "read_file", arguments: { path: r.path, limit: 2000 } } })).join("\n") + "\n";
    const readResp = await callMCPServer(readJsonl);
    for (const r of readCalls) {
      const rc = readResp?.[r.id];
      if (rc?.content?.[0]) {
        sections.push(`[Isi file: ${r.name}]\n${rc.content[0].text.substring(0, 2000)}`);
      }
    }
  }

  if (sections.length === 0) return null;
  return { context: sections.join("\n\n") };
}

// API: generate (streaming)
app.post("/api/generate", async (req, res) => {
  const { prompt, modelPath, history } = req.body;
  const cfg = loadConfig();
  const ckpt = modelPath ? path.resolve(INDIGO_DIR, modelPath) : cfg.model;

  if (!ckpt) return res.status(400).json({ error: "no model selected" });

  res.setHeader("Content-Type", "text/event-stream");
  res.setHeader("Cache-Control", "no-cache");
  res.setHeader("Connection", "keep-alive");

  const genScript = path.join(__dirname, "generate.py");

  // agent loop: gather file context via MCP if needed
  let agentCtx = null;
  if (cfg.agent !== false) {
    try {
      agentCtx = await agentGatherContext(prompt);
      if (agentCtx) res.write(`data: ${JSON.stringify({ type: "agent", context: agentCtx.context })}\n\n`);
    } catch (e) {
      console.error("[agent error]", e.message);
      res.write(`data: ${JSON.stringify({ type: "agent", error: e.message })}\n\n`);
    }
  }

  const finalPrompt = agentCtx ? `Konteks dari MCP:\n${agentCtx.context}\n\nPertanyaan: ${prompt}` : prompt;

  const env = {
    ...process.env,
    INDIGO_DIR: INDIGO_DIR,
    CKPT: ckpt,
    PROMPT: finalPrompt,
    HISTORY: JSON.stringify(history || []),
    TEMPERATURE: String(cfg.temperature ?? 0.8),
    TOP_K: String(cfg.topK ?? 40),
    TOP_P: String(cfg.topP ?? 0.95),
    MAX_NEW: String(cfg.maxNew ?? 200),
  };
  const proc = spawn(PYTHON, ["-u", genScript], { cwd: __dirname, env });

  let output = "";
  let stderrBuf = "";
  proc.on("error", (e) => { console.error("[generate spawn error]", e.message); });
  proc.stdout.on("data", (d) => { output += d.toString(); });
  proc.stderr.on("data", (d) => { stderrBuf += d.toString(); });
  proc.on("close", () => {
    try {
      const lines = output.trim().split("\n");
      const lastLine = lines[lines.length - 1];
      const result = JSON.parse(lastLine);
      res.write(`data: ${JSON.stringify(result)}\n\n`);
    } catch (e) {
      res.write(`data: ${JSON.stringify({ error: "parse error", raw: output.substring(0, 500), stderr: stderrBuf.substring(0, 300) })}\n\n`);
    }
    res.write("data: [DONE]\n\n");
    res.end();
  });

  res.on("close", () => {
    if (!proc.killed) proc.kill();
  });
});

// API: guard check
app.post("/api/guard", (req, res) => {
  const { text } = req.body;
  const kamusPath = path.join(DATA_DIR, "kamus_id.txt");
  if (!fs.existsSync(kamusPath)) return res.json({ ratio: 1.0, words: 0, known: 0 });
  const words = new Set(fs.readFileSync(kamusPath, "utf-8").split("\n").map((w) => w.trim().toLowerCase()).filter(Boolean));
  const tokens = text.toLowerCase().split(/\s+/).filter(Boolean);
  const known = tokens.filter((t) => words.has(t)).length;
  res.json({ ratio: tokens.length ? known / tokens.length : 1.0, words: tokens.length, known });
});

app.listen(PORT, () => {
  console.log(`indigo-chat running at http://localhost:${PORT}`);
});