import assert from "node:assert/strict";
import { execFileSync } from "node:child_process";
import test, { after, before } from "node:test";
import pg from "pg";
import { PostgresConnectionRegistry } from "../src/connection-registry.js";
import { rolePassword } from "../src/execution-database.js";
import { PostgresService } from "../src/postgres-service.js";

/**
 * Against a real PostgreSQL, run in a throwaway container, with a catalog user who is not a superuser: one
 * with CREATEDB and CREATEROLE and nothing else, as in a deployment. Skipped where there is no Docker.
 */
function dockerAvailable() {
  try { execFileSync("docker", ["info"], { stdio: "ignore", timeout: 15000 }); return true; } catch { return false; }
}
const enabled = dockerAvailable() && process.env.POSTGRES_MCP_SKIP_DOCKER !== "1";
const SECRET = "0123456789abcdef0123456789abcdef";
let container;
let port;

async function superuser(database = "postgres") {
  const client = new pg.Client({ host: "127.0.0.1", port, user: "postgres", password: "postgres-password", database });
  await client.connect();
  return client;
}

before(async () => {
  if (!enabled) return;
  container = execFileSync("docker", ["run", "-d", "--rm", "-e", "POSTGRES_PASSWORD=postgres-password", "-p", "127.0.0.1::5432", "postgres:17-alpine"]).toString().trim();
  port = Number(execFileSync("docker", ["port", container, "5432/tcp"]).toString().trim().split("\n")[0].split(":").pop());
  for (let attempt = 0; attempt < 60; attempt++) {
    try {
      const admin = await superuser();
      await admin.query("CREATE ROLE ops LOGIN PASSWORD 'ops-password' CREATEDB CREATEROLE");
      await admin.end();
      return;
    } catch {
      await new Promise((resolve) => setTimeout(resolve, 500));
    }
  }
  throw new Error("The PostgreSQL container did not come up");
});

after(() => {
  if (container) execFileSync("docker", ["rm", "-f", container], { stdio: "ignore" });
});

const pinned = (execution) => ({
  "x-postgres-scope": execution,
  "x-postgres-host": "127.0.0.1",
  "x-postgres-port": String(port),
  "x-postgres-database": "postgres",
  "x-postgres-user": "ops",
  "x-postgres-password": "ops-password",
  "x-postgres-sslmode": "disable",
  "x-postgres-role-secret": SECRET
});

function newRegistry() {
  return new PostgresConnectionRegistry({ maxRows: 50 });
}

test("a session gets a database of its own and can build and query what it likes in it", { skip: !enabled }, async () => {
  const registry = newRegistry();
  const context = { sessionId: "s1", pinned: pinned("it-a") };
  assert.deepEqual(await registry.openInternal(context), { database: "exec_it_a" });
  try {
    assert.deepEqual((await registry.callInternal("listTables", {}, context)).tables, []);
    const created = await registry.callInternal("execute", {
      statement: "CREATE TABLE notes (id serial PRIMARY KEY, body text NOT NULL, data jsonb); INSERT INTO notes(body) VALUES ('first')"
    }, context);
    assert.equal(created.results.length, 2);
    const inserted = await registry.callInternal("execute", {
      statement: "INSERT INTO notes(body, data) VALUES ($1, $2) RETURNING id", params: ["it's a quote", { k: [1, 2] }]
    }, context);
    assert.equal(inserted.command, "INSERT");
    assert.equal(inserted.rowCount, 1);
    assert.equal(inserted.rows.length, 1);

    const read = await registry.callInternal("query", { statement: "SELECT body, data FROM notes ORDER BY id" }, context);
    assert.deepEqual(read.rows.map((row) => row.body), ["first", "it's a quote"]);
    assert.deepEqual(read.rows[1].data, { k: [1, 2] });
    assert.deepEqual(read.columns, [{ name: "body" }, { name: "data" }]);
    assert.equal(read.truncated, false);

    const tables = await registry.callInternal("listTables", {}, context);
    assert.deepEqual(tables.tables.map((table) => table.name), ["notes"]);
    const described = await registry.callInternal("describeTable", { table: "notes" }, context);
    assert.deepEqual(described.columns.map((column) => column.name), ["id", "body", "data"]);
    assert.deepEqual(described.primaryKey, ["id"]);
    await assert.rejects(registry.callInternal("describeTable", { table: "nope" }, context), /There is no table public.nope/);
  } finally {
    registry.closeAll();
  }
});

test("a read cannot change anything, and says so", { skip: !enabled }, async () => {
  const registry = newRegistry();
  const context = { sessionId: "s2", pinned: pinned("it-readonly") };
  await registry.openInternal(context);
  try {
    await registry.callInternal("execute", { statement: "CREATE TABLE t (id int)" }, context);
    await assert.rejects(registry.callInternal("query", { statement: "INSERT INTO t VALUES (1)" }, context), /read-only transaction/);
    await assert.rejects(registry.callInternal("query", { statement: "DELETE FROM t" }, context), /read-only transaction/);
    assert.equal((await registry.callInternal("query", { statement: "SELECT count(*)::int AS n FROM t" }, context)).rows[0].n, 0);
  } finally {
    registry.closeAll();
  }
});

test("a long result is cut at the row limit and says it was", { skip: !enabled }, async () => {
  const registry = newRegistry();
  const context = { sessionId: "s3", pinned: pinned("it-rows") };
  await registry.openInternal(context);
  try {
    const all = await registry.callInternal("query", { statement: "SELECT g FROM generate_series(1, 500) g", maxRows: 50 }, context);
    assert.equal(all.rows.length, 50);
    assert.equal(all.truncated, true);
    const few = await registry.callInternal("query", { statement: "SELECT g FROM generate_series(1, 3) g" }, context);
    assert.equal(few.rows.length, 3);
    assert.equal(few.truncated, false);
    await assert.rejects(registry.callInternal("query", { statement: "SELECT 1", maxRows: 5000 }, context), /maxRows must be an integer between 1 and 50/);
  } finally {
    registry.closeAll();
  }
});

test("the user of one execution cannot reach another's database, nor become anyone more powerful", { skip: !enabled }, async () => {
  const registry = newRegistry();
  const a = { sessionId: "sa", pinned: pinned("it-iso-a") };
  const b = { sessionId: "sb", pinned: pinned("it-iso-b") };
  await registry.openInternal(a);
  await registry.openInternal(b);
  try {
    await registry.callInternal("execute", { statement: "CREATE TABLE secret (v text); INSERT INTO secret VALUES ('only for a')" }, a);
    // The two databases are separate: b has no such table, and no statement of its own can name a's.
    await assert.rejects(registry.callInternal("query", { statement: "SELECT v FROM secret" }, b), /secret/);

    // Nor can a's user connect to b's database, with the password of its own or any other.
    const stranger = new pg.Client({ host: "127.0.0.1", port, database: "exec_it_iso_b", user: "exec_it_iso_a", password: "x" });
    await assert.rejects(stranger.connect());

    // Whatever a's user tries, it stays that user: it cannot switch to a role with more rights...
    await assert.rejects(registry.callInternal("execute", { statement: "SET ROLE ops" }, a), /permission denied/);
    await assert.rejects(registry.callInternal("execute", { statement: "SET ROLE postgres" }, a), /permission denied/);
    // ...make users or databases, or read the server's files.
    await assert.rejects(registry.callInternal("execute", { statement: "CREATE ROLE intruder LOGIN" }, a), /permission denied/);
    await assert.rejects(registry.callInternal("execute", { statement: "CREATE DATABASE another" }, a), /permission denied/);
    await assert.rejects(registry.callInternal("query", { statement: "SELECT pg_read_file('/etc/passwd')" }, a), /permission denied/);
    await assert.rejects(registry.callInternal("execute", { statement: "COPY secret TO PROGRAM 'id'" }, a), /permission denied|must be superuser|pg_execute_server_program/);
    const who = await registry.callInternal("query", { statement: "SELECT current_user AS who, current_database() AS db" }, a);
    assert.deepEqual(who.rows[0], { who: "exec_it_iso_a", db: "exec_it_iso_a" });
  } finally {
    registry.closeAll();
  }
});

test("two sessions of one execution share its database, even when they open it at once", { skip: !enabled }, async () => {
  const registry = newRegistry();
  const contexts = [1, 2, 3].map((n) => ({ sessionId: `race-${n}`, pinned: pinned("it-race") }));
  try {
    const opened = await Promise.all(contexts.map((context) => registry.openInternal(context)));
    assert.deepEqual(opened, [{ database: "exec_it_race" }, { database: "exec_it_race" }, { database: "exec_it_race" }]);
    await registry.callInternal("execute", { statement: "CREATE TABLE shared (v int); INSERT INTO shared VALUES (7)" }, contexts[0]);
    const read = await registry.callInternal("query", { statement: "SELECT v FROM shared" }, contexts[2]);
    assert.deepEqual(read.rows, [{ v: 7 }]);
  } finally {
    registry.closeAll();
  }
});

test("the database carries the age the manager's janitor reads, and only the derived password gets in", { skip: !enabled }, async () => {
  const registry = newRegistry();
  const context = { sessionId: "s4", pinned: pinned("it-age") };
  await registry.openInternal(context);
  try {
    const admin = await superuser();
    const { rows } = await admin.query("SELECT shobj_description(oid, 'pg_database') AS comment FROM pg_database WHERE datname = 'exec_it_age'");
    await admin.end();
    assert.match(rows[0].comment, /^humainflow:created=\d{10}$/);

    const wrong = new pg.Client({ host: "127.0.0.1", port, database: "exec_it_age", user: "exec_it_age", password: "not-the-derived-one" });
    await assert.rejects(wrong.connect(), /password authentication failed/);
    // The user and the database are the same name, and the password is derived from the name and the secret.
    const right = new pg.Client({ host: "127.0.0.1", port, database: "exec_it_age", user: "exec_it_age", password: rolePassword(SECRET, "exec_it_age") });
    await right.connect();
    await right.end();
  } finally {
    registry.closeAll();
  }
});

test("an external session works against the node's host, and a bad password is refused without echoing it", { skip: !enabled }, async () => {
  // The container is on loopback, which only a listed host may be: unlisted, it would be refused as not public.
  const registry = new PostgresConnectionRegistry({ allowedHosts: ["127.0.0.1"] });
  const context = { sessionId: "s6", pinned: { "x-postgres-host": "127.0.0.1", "x-postgres-port": String(port), "x-postgres-database": "postgres" } };
  const opened = await registry.open({ user: "ops", password: "ops-password", sslmode: "disable" }, context);
  try {
    assert.equal(opened.user, "ops");
    assert.equal((await registry.call("query", { postgresSessionId: opened.postgresSessionId, statement: "SELECT 41 + 1 AS answer" }, context)).rows[0].answer, 42);
    await assert.rejects(
      registry.open({ user: "ops", password: "wrong-password-9", sslmode: "disable" }, context),
      (error) => /PostgreSQL/.test(error.message) && !/wrong-password-9/.test(error.message));
  } finally {
    registry.closeAll();
  }
});

test("without an allowlist, a database on a private address is not reached", { skip: !enabled }, async () => {
  const registry = new PostgresConnectionRegistry();
  const context = { sessionId: "s7", pinned: { "x-postgres-host": "127.0.0.1", "x-postgres-port": String(port), "x-postgres-database": "postgres" } };
  await assert.rejects(registry.open({ user: "ops", password: "ops-password", sslmode: "disable" }, context), /not a public address/);
  const named = { sessionId: "s8", pinned: { "x-postgres-host": "localhost", "x-postgres-port": String(port), "x-postgres-database": "postgres" } };
  await assert.rejects(registry.open({ user: "ops", password: "ops-password", sslmode: "disable" }, named), /not a public one/);
});

test("a statement that runs too long is stopped", { skip: !enabled }, async () => {
  const service = new PostgresService({ host: "127.0.0.1", port, database: "postgres", user: "ops", password: "ops-password", sslmode: "disable", statementTimeoutMs: 500 });
  await service.verify();
  try {
    await assert.rejects(service.query({ statement: "SELECT pg_sleep(5)" }), /canceling statement due to statement timeout/);
    assert.equal((await service.query({ statement: "SELECT 1 AS ok" })).rows[0].ok, 1, "and the session is still usable");
  } finally {
    service.destroy();
  }
});