// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ========= Copyright 2026 @ Strukto.AI All Rights Reserved. =========

import { CreateBucketCommand, PutObjectCommand, S3Client } from "@aws-sdk/client-s3";
import { MongoClient } from "mongodb";
import {
  buildRuntime,
  CommandSafeguard,
  MongoDBResource,
  MountMode,
  PathSpec,
  RAMResource,
  RedisResource,
  S3Resource,
  ScriptSource,
  VfsRuntime,
  Workspace,
  type RouteScript,
  type Runtime,
  type RunResult,
} from "@struktoai/mirage-node";

const REDIS_URL = process.env.REDIS_URL ?? "redis://localhost:6379/0";
const MONGODB_URI = process.env.MONGODB_URI ?? "mongodb://localhost:27017";
const S3_ENDPOINT = process.env.S3_ENDPOINT ?? "http://localhost:9000";
const S3_KEY = process.env.AWS_ACCESS_KEY_ID ?? "minio";
const S3_SECRET = process.env.AWS_SECRET_ACCESS_KEY ?? "minio123";
const DB = "mirage_integ_runtime";
const BUCKET = "mirage-integ-runtime-ts";

const ENC = new TextEncoder();
const DEC = new TextDecoder();

const SLOW_SCRIPT = "n = 0\nfor i in range(300000000):\n    n = n + 1\n";

const ANALYZE_SCRIPT = `from pathlib import Path
ram = Path('/ram/data.txt').read_text().strip()
s3 = Path('/s3/greeting.txt').read_text().strip()
print('ram:', ram)
print('s3:', s3)
print('argv sum:', int(argv[1]) + int(argv[2]))
`;

const CASES: [string, string][] = [
  [
    "py3_ram_read",
    "python3 -c \"from pathlib import Path; print(Path('/ram/data.txt').read_text().strip().upper())\"",
  ],
  [
    "py3_redis_read",
    "python3 -c \"from pathlib import Path; print(Path('/redis/notes.txt').read_text().strip())\"",
  ],
  [
    "py3_s3_read",
    "python3 -c \"from pathlib import Path; print(Path('/s3/greeting.txt').read_text().strip())\"",
  ],
  [
    "py3_mongodb_meta",
    `python3 -c "import json; from pathlib import Path; meta = json.loads(Path('/mongodb/${DB}/database.json').read_text()); print(meta['database'], len(meta['collections']))"`,
  ],
  [
    "py3_mongodb_schema",
    `python3 -c "import json; from pathlib import Path; s = json.loads(Path('/mongodb/${DB}/collections/books/schema.json').read_text()); print(s['name'], s['kind'], s['document_count'], [f['path'] for f in s['fields']])"`,
  ],
  [
    "py3_mongodb_docs",
    `python3 -c "from pathlib import Path; lines = Path('/mongodb/${DB}/collections/books/documents.jsonl').read_text().splitlines(); print('books:', len(lines))"`,
  ],
  ["py3_script_argv", "python3 /ram/analyze.py 40 2"],
  ["py3_stdin_pipe", "cat /ram/pipe.py | python3"],
  ["py3_pipe_to_grep", "python3 -c \"print('alpha'); print('beta')\" | grep beta"],
  ["js_node_eval", 'node -e "console.log(6 * 7)"'],
  ["js_node_pipe_to_grep", "node -e \"console.log('alpha'); console.log('beta')\" | grep beta"],
  [
    "py3_write_back",
    "python3 -c \"from pathlib import Path; Path('/ram/out.txt').write_text('written-by-python3')\" && cat /ram/out.txt",
  ],
  ["python_alias", "python -c 'print(6 * 7)'"],
  [
    "py3_env",
    "export GREETING=hello-runtime && python3 -c \"import os; print(os.getenv('GREETING'))\"",
  ],
  [
    // str(p) instead of p.name: the monty JS binding returns iterdir
    // entries as plain strings (the Python binding returns Path objects).
    "py3_iterdir",
    "python3 -c \"from pathlib import Path; print(sorted(str(p).split('/')[-1] for p in Path('/ram').iterdir()))\"",
  ],
];

const ERROR_CASES: [string, string][] = [
  [
    "py3_missing_file",
    "python3 -c \"from pathlib import Path; print(Path('/ram/nope.txt').read_text())\"",
  ],
  [
    "py3_host_fs_invisible",
    "python3 -c \"from pathlib import Path; print(Path('/etc/passwd').read_text())\"",
  ],
];

async function seedMongo(): Promise<void> {
  const client = new MongoClient(MONGODB_URI);
  try {
    await client.db(DB).dropDatabase();
    await client
      .db(DB)
      .collection("books")
      .insertMany([
        { _id: 1 as never, title: "alpha" },
        { _id: 2 as never, title: "beta" },
      ]);
    await client
      .db(DB)
      .collection("authors")
      .insertMany([{ _id: 1 as never, name: "ada" }]);
  } finally {
    await client.close();
  }
}

async function putS3(key: string, body: string): Promise<void> {
  const client = new S3Client({
    region: "us-east-1",
    endpoint: S3_ENDPOINT,
    forcePathStyle: true,
    credentials: { accessKeyId: S3_KEY, secretAccessKey: S3_SECRET },
  });
  await client.send(new PutObjectCommand({ Bucket: BUCKET, Key: key, Body: body }));
  client.destroy();
}

async function seedS3(): Promise<void> {
  const client = new S3Client({
    region: "us-east-1",
    endpoint: S3_ENDPOINT,
    forcePathStyle: true,
    credentials: { accessKeyId: S3_KEY, secretAccessKey: S3_SECRET },
  });
  try {
    await client.send(new CreateBucketCommand({ Bucket: BUCKET }));
  } catch {
    // bucket already exists from a prior run
  }
  client.destroy();
  await putS3("greeting.txt", "hello from s3\n");
  await putS3("cache_inval.txt", "seeded-cold\n");
  await putS3("warm.txt", "cold-bytes\n");
}

function buildWorkspace(runId: string): Workspace {
  const ram = new RAMResource();
  ram.store.files.set("/data.txt", ENC.encode("hello from ram\n"));
  ram.store.files.set("/analyze.py", ENC.encode(ANALYZE_SCRIPT));
  ram.store.files.set("/pipe.py", ENC.encode("print('came-through-pipe')\n"));
  const redis = new RedisResource({
    url: REDIS_URL,
    keyPrefix: `mirage-integ-runtime-ts-${runId}/`,
  });
  const s3 = new S3Resource({
    bucket: BUCKET,
    region: "us-east-1",
    endpoint: S3_ENDPOINT,
    forcePathStyle: true,
    accessKeyId: S3_KEY,
    secretAccessKey: S3_SECRET,
  });
  const mongodb = new MongoDBResource({ uri: MONGODB_URI, databases: [DB] });
  return new Workspace(
    { "/ram": ram, "/redis": redis, "/s3": s3, "/mongodb": mongodb },
    { mode: MountMode.EXEC, runtimes: ["monty", "quickjs", "vfs"] },
  );
}

class EchoBox implements Runtime {
  readonly name = "sandbox";
  captures: readonly string[] = ["nvidia-smi"];
  readonly runsLines = true;
  script?: RouteScript;
  attach(): void {}
  run(): Promise<RunResult> {
    return Promise.reject(new Error("whole-line runtimes take lines"));
  }
  runLine(
    line: string,
    _stdin: Uint8Array | null,
    _env: Record<string, string>,
    _cwd: string,
  ): Promise<RunResult> {
    return Promise.resolve({ stdout: ENC.encode(`box:${line}\n`), stderr: null, exitCode: 0 });
  }
  close(): Promise<void> {
    return Promise.resolve();
  }
}

async function run(ws: Workspace, name: string, cmd: string): Promise<void> {
  const io = await ws.execute(cmd);
  const out = DEC.decode(io.stdout);
  console.log(`=== ${name} ===`);
  if (out) process.stdout.write(out.endsWith("\n") ? out : out + "\n");
}

async function runError(ws: Workspace, name: string, cmd: string): Promise<void> {
  const io = await ws.execute(cmd);
  console.log(`=== ${name} ===`);
  console.log(`exit_code=${io.exitCode}`);
}

async function main(): Promise<void> {
  await seedS3();
  await seedMongo();
  const runId = Math.random().toString(36).slice(2, 10);
  const ws = buildWorkspace(runId);
  await ws.execute("echo redis says hi > /redis/notes.txt");
  for (const [name, cmd] of CASES) await run(ws, name, cmd);
  for (const [name, cmd] of ERROR_CASES) await runError(ws, name, cmd);

  // Sandbox I/O shares the shell file cache (s3 caches reads).
  // Invalidate: prime the cache with cat (applyIo populates it after
  // the command), then a sandbox write must drop the entry so the
  // follow-up cat re-fetches the new bytes.
  await run(ws, "s3_inval_prime", "cat /s3/cache_inval.txt");
  await run(
    ws,
    "py3_s3_cache_invalidate",
    "python3 -c \"from pathlib import Path; Path('/s3/cache_inval.txt').write_text('updated-by-sandbox')\" && cat /s3/cache_inval.txt",
  );
  // Warm read: prime, mutate the backend out of band, and the sandbox
  // read (LAZY, never revalidated) must keep serving the cached bytes
  // instead of hitting the backend.
  await run(ws, "s3_warm_prime", "cat /s3/warm.txt");
  await putS3("warm.txt", "hot-bytes\n");
  await run(
    ws,
    "py3_s3_warm_read",
    "python3 -c \"from pathlib import Path; print(Path('/s3/warm.txt').read_text().strip())\"",
  );

  // Op safeguards bind to the executing (post-symlink-follow) mount: a
  // sandbox read through a link on an unsafeguarded mount still gets
  // the target mount's byte cap.
  const wsLink = new Workspace(
    {
      "/data": [
        new RAMResource(),
        MountMode.EXEC,
        { read: new CommandSafeguard({ maxBytes: 8 }) },
      ],
      "/r": new RAMResource(),
    },
    { mode: MountMode.EXEC, runtimes: ["monty", "vfs"] },
  );
  await wsLink.execute("echo 0123456789abcdef > /data/big.txt");
  await wsLink.execute("ln -s /data/big.txt /r/link");
  await run(
    wsLink,
    "sbx_link_guard",
    "python3 -c \"from pathlib import Path; print(Path('/r/link').read_text())\"",
  );
  await wsLink.close();

  // Cross-mount rename through dispatch: both languages address the dst
  // against the source mount, so the file lands on the source backend
  // under the dst's virtual path (EXDEV follow-up).
  const wsMv = new Workspace(
    { "/a": new RAMResource(), "/b": new RAMResource() },
    { mode: MountMode.EXEC, runtimes: ["monty", "vfs"] },
  );
  await wsMv.execute("echo moved-bytes > /a/x.txt");
  await wsMv.dispatch("rename", "/a/x.txt", [PathSpec.fromStrPath("/b/y.txt")]);
  await run(wsMv, "xmount_rename", "cat /a/b/y.txt");
  await runError(wsMv, "xmount_rename_src", "cat /a/x.txt");
  await runError(wsMv, "xmount_rename_dst_mount", "cat /b/y.txt");
  await wsMv.close();

  const slowRam = new RAMResource();
  slowRam.store.files.set("/slow.py", ENC.encode(SLOW_SCRIPT));
  const wsSg = new Workspace(
    { "/ram": slowRam },
    {
      mode: MountMode.EXEC,
      runtimes: ["monty", "vfs"],
      commandSafeguards: {
        "/ram": { python3: new CommandSafeguard({ timeoutSeconds: 1 }) },
      },
    },
  );
  await runError(wsSg, "py3_safeguard_timeout", "python3 /ram/slow.py");
  await wsSg.close();

  // Routing ladder: monty's entry script refuses lines carrying a
  // skip-monty marker; with no other capturer those lines are
  // admission failures (126) while vfs commands keep running. A
  // scripted vfs entry locks down the lines it refuses.
  const routedMonty = buildRuntime("monty");
  routedMonty.script = new ScriptSource("'skip-monty' not in ctx['line']");
  const wsRoute = new Workspace(
    { "/ram": new RAMResource() },
    { mode: MountMode.EXEC, runtimes: [routedMonty, "vfs"] },
  );
  await run(wsRoute, "route_script_monty", 'python3 -c "print(6 * 7)"');
  const denied = await wsRoute.execute('python3 -c "x" && echo skip-monty');
  console.log("=== route_refused ===");
  console.log(`exit_code=${denied.exitCode}`);
  process.stdout.write(DEC.decode(denied.stderr));
  await run(wsRoute, "route_vfs_open", "echo vfs-open");
  await wsRoute.close();

  const wsLock = new Workspace(
    { "/ram": new RAMResource() },
    {
      mode: MountMode.EXEC,
      runtimes: [new VfsRuntime(new ScriptSource("'secret' not in ctx['line']"))],
    },
  );
  await run(wsLock, "lockdown_allow", "echo fine");
  await runError(wsLock, "lockdown_deny", "cat /ram/secret.txt");
  await wsLock.close();

  // Per-runtime context: monty's script sees its own stage, so a
  // pipeline led by cat still routes python3 onto monty.
  const ctxMonty = buildRuntime("monty");
  ctxMonty.script = new ScriptSource("ctx['command'] == 'python3'");
  const wsCtx = new Workspace(
    { "/ram": new RAMResource() },
    { mode: MountMode.EXEC, runtimes: [ctxMonty, "vfs"] },
  );
  await wsCtx.execute("printf \"print('came-through-pipe')\" > /ram/pipe.py");
  await run(wsCtx, "ctx_command_pipeline", "cat /ram/pipe.py | python3");
  await wsCtx.close();

  // Overriding the vfs runtime: explicit captures restrict the
  // workspace to those commands; interpreter bindings untouched.
  const wsCap = new Workspace(
    { "/ram": new RAMResource() },
    {
      mode: MountMode.EXEC,
      runtimes: ["quickjs", new VfsRuntime({ captures: ["echo", "cat"] })],
    },
  );
  await run(wsCap, "vfs_captures_allow", "echo vfs-captured");
  await runError(wsCap, "vfs_captures_deny", "ls /ram");
  await run(wsCap, "vfs_captures_node", 'node -e "console.log(6 * 7)"');
  await wsCap.close();

  // addRuntime: the runtime argument can only name a workspace entry,
  // so it fails loud until the entry is added at runtime.
  const wsAdd = new Workspace(
    { "/ram": new RAMResource() },
    { mode: MountMode.EXEC, runtimes: ["quickjs", "vfs"] },
  );
  try {
    await wsAdd.execute('python3 -c "x"', { runtime: "monty" });
  } catch {
    console.log("=== add_runtime_arg_before ===");
    console.log("unknown-runtime-rejected");
  }
  wsAdd.addRuntime("monty");
  const added = await wsAdd.execute('python3 -c "print(40 + 2)"', { runtime: "monty" });
  console.log("=== add_runtime_arg_after ===");
  process.stdout.write(DEC.decode(added.stdout));
  await wsAdd.close();

  // A runtime that runs whole lines: the raw line lands there
  // wholesale when it captures one of the line's commands or "*";
  // a refused line falls back to the workspace shell.
  const wsLine = new Workspace(
    { "/ram": new RAMResource() },
    { mode: MountMode.EXEC, runtimes: [new EchoBox(), "vfs"] },
  );
  await run(wsLine, "whole_line_capture", "nvidia-smi -L | grep GPU");
  await run(wsLine, "whole_line_uncaptured", "echo still-the-workspace");
  await wsLine.close();

  const star = new EchoBox();
  star.captures = ["*"];
  star.script = new ScriptSource("'keep-out' not in ctx['line']");
  const wsStar = new Workspace(
    { "/ram": new RAMResource() },
    { mode: MountMode.EXEC, runtimes: [star, "vfs"] },
  );
  await run(wsStar, "whole_line_star", "ls /ram && echo done");
  await run(wsStar, "whole_line_refused", "echo keep-out");
  await wsStar.close();
  await ws.close();
}

main()
  .then(() => process.exit(0))
  .catch((err: unknown) => {
    console.error(err);
    process.exit(1);
  });
