import type { ConnectionConfig } from "./connections.js";
import { createServer, connect as netConnect, type Server, type Socket } from "node:net";
import { readFile } from "node:fs/promises";
import { join } from "node:path";
import { homedir, platform } from "node:os";
import Database from "better-sqlite3";
import { sqlSafetyFromEnv } from "./sql-safety.js";
import { isDirectQueryType } from "./diagnostics.js";

export interface TableInfo {
  name: string;
  type: string;
}

export interface ColumnInfo {
  name: string;
  data_type: string;
  is_nullable: boolean;
  column_default: string | null;
  is_primary_key: boolean;
  comment: string | null;
}

export interface QueryResult {
  columns: string[];
  rows: Record<string, unknown>[];
  row_count: number;
}

export interface QueryOptions {
  maxRows?: number;
  timeoutMs?: number;
}

const MAX_ROWS = 100;
const IDLE_TIMEOUT_MS = 5 * 60 * 1000;
const QUERY_TIMEOUT_MS = 30_000;

interface PoolEntry {
  type: "pg" | "mysql";
  pool: unknown;
  timer: ReturnType<typeof setTimeout>;
}

interface RqliteResult {
  columns?: string[];
  values?: unknown[][];
  rows_affected?: number;
  error?: string;
}

interface RqliteResponse {
  results?: RqliteResult[];
}

const pools = new Map<string, PoolEntry>();
const proxyTunnels = new Map<string, { server: Server; port: number; sockets: Set<Socket> }>();

function poolKey(config: ConnectionConfig): string {
  return `${config.id}:${config.database || ""}`;
}

function evictPool(key: string, entry: PoolEntry) {
  pools.delete(key);
  clearTimeout(entry.timer);
  if (entry.type === "pg") {
    (entry.pool as import("pg").Pool).end().catch(() => {});
  } else {
    (entry.pool as import("mysql2/promise").Pool).end().catch(() => {});
  }
}

function resetIdleTimer(key: string, entry: PoolEntry) {
  clearTimeout(entry.timer);
  entry.timer = setTimeout(() => evictPool(key, entry), IDLE_TIMEOUT_MS);
}

export async function closeDatabaseResources(): Promise<void> {
  const poolEntries = [...pools.entries()];
  pools.clear();
  await Promise.all(
    poolEntries.map(async ([, entry]) => {
      clearTimeout(entry.timer);
      if (entry.type === "pg") {
        await (entry.pool as import("pg").Pool).end().catch(() => {});
      } else {
        await (entry.pool as import("mysql2/promise").Pool).end().catch(() => {});
      }
    }),
  );

  const tunnels = [...proxyTunnels.values()];
  proxyTunnels.clear();
  await Promise.all(
    tunnels.map(
      ({ server, sockets }) =>
        new Promise<void>((resolve) => {
          for (const socket of sockets) socket.destroy();
          server.close(() => resolve());
        }),
    ),
  );
}

async function getPgPool(config: ConnectionConfig): Promise<import("pg").Pool> {
  const key = poolKey(config);
  const existing = pools.get(key);
  if (existing?.type === "pg") {
    resetIdleTimer(key, existing);
    return existing.pool as import("pg").Pool;
  }

  const pg = await import("pg");
  const endpoint = await connectionEndpoint(config);
  const pool = new pg.default.Pool({
    connectionString: buildConnectionUrl(config, endpoint),
    max: 3,
    idleTimeoutMillis: 30_000,
    connectionTimeoutMillis: 10_000,
  });
  pool.on("error", () => {});
  const entry: PoolEntry = { type: "pg", pool, timer: setTimeout(() => {}, 0) };
  pools.set(key, entry);
  resetIdleTimer(key, entry);
  return pool;
}

async function getMysqlPool(config: ConnectionConfig): Promise<import("mysql2/promise").Pool> {
  const key = poolKey(config);
  const existing = pools.get(key);
  if (existing?.type === "mysql") {
    resetIdleTimer(key, existing);
    return existing.pool as import("mysql2/promise").Pool;
  }

  const mysql = await import("mysql2/promise");
  const endpoint = await connectionEndpoint(config);
  const pool = mysql.default.createPool({
    uri: buildConnectionUrl(config, endpoint),
    connectionLimit: 3,
    idleTimeout: 30_000,
    connectTimeout: 10_000,
  });
  const entry: PoolEntry = { type: "mysql", pool, timer: setTimeout(() => {}, 0) };
  pools.set(key, entry);
  resetIdleTimer(key, entry);
  return pool;
}

async function connectionEndpoint(config: ConnectionConfig): Promise<{ host: string; port: number }> {
  if (!config.proxy_enabled || !config.proxy_host) return { host: config.host, port: config.port };
  const existing = proxyTunnels.get(config.id);
  if (existing) return { host: "127.0.0.1", port: existing.port };

  const sockets = new Set<Socket>();
  const server = createServer((inbound) => {
    sockets.add(inbound);
    inbound.once("close", () => sockets.delete(inbound));
    connectViaProxy(config)
      .then((outbound) => {
        sockets.add(outbound);
        outbound.once("close", () => sockets.delete(outbound));
        inbound.pipe(outbound);
        outbound.pipe(inbound);
      })
      .catch(() => inbound.destroy());
  });
  const port = await new Promise<number>((resolve, reject) => {
    server.once("error", reject);
    server.listen(0, "127.0.0.1", () => {
      const address = server.address();
      if (address && typeof address === "object") resolve(address.port);
      else reject(new Error("Failed to bind proxy tunnel"));
    });
  });
  proxyTunnels.set(config.id, { server, port, sockets });
  return { host: "127.0.0.1", port };
}

function buildConnectionUrl(config: ConnectionConfig, endpoint: { host: string; port: number }): string {
  const db = config.database || "";
  const params = config.url_params || "";
  const suffix = params ? `?${params}` : "";
  if (isMysqlType(config.db_type)) {
    return `mysql://${encodeURIComponent(config.username)}:${encodeURIComponent(config.password)}@${endpoint.host}:${endpoint.port}/${db}${suffix}`;
  }
  return `postgres://${encodeURIComponent(config.username)}:${encodeURIComponent(config.password)}@${endpoint.host}:${endpoint.port}/${db}${suffix}`;
}

function connectViaProxy(config: ConnectionConfig): Promise<Socket> {
  return new Promise((resolve, reject) => {
    const socket = netConnect(config.proxy_port || 1080, config.proxy_host || "127.0.0.1");
    socket.once("error", reject);
    socket.once("connect", () => {
      if ((config.proxy_type || "socks5") === "http") {
        httpConnect(socket, config, resolve, reject);
      } else {
        socks5Connect(socket, config, resolve, reject);
      }
    });
  });
}

function httpConnect(socket: Socket, config: ConnectionConfig, resolve: (socket: Socket) => void, reject: (err: Error) => void) {
  const target = `${config.host}:${config.port}`;
  const lines = [`CONNECT ${target} HTTP/1.1`, `Host: ${target}`];
  if (config.proxy_username || config.proxy_password) {
    const token = Buffer.from(`${config.proxy_username || ""}:${config.proxy_password || ""}`).toString("base64");
    lines.push(`Proxy-Authorization: Basic ${token}`);
  }
  socket.write(`${lines.join("\r\n")}\r\n\r\n`);
  let buffer = Buffer.alloc(0);
  socket.on("data", function onData(chunk: Buffer) {
    buffer = Buffer.concat([buffer, chunk]);
    const end = buffer.indexOf("\r\n\r\n");
    if (end < 0) return;
    socket.off("data", onData);
    const head = buffer.subarray(0, end).toString("utf8");
    if (!/^HTTP\/1\.[01] 200\b/.test(head)) {
      reject(new Error(`HTTP proxy CONNECT failed: ${head.split("\r\n")[0] || "invalid response"}`));
      socket.destroy();
      return;
    }
    const rest = buffer.subarray(end + 4);
    if (rest.length) socket.unshift(rest);
    resolve(socket);
  });
}

function socks5Connect(socket: Socket, config: ConnectionConfig, resolve: (socket: Socket) => void, reject: (err: Error) => void) {
  const wantsAuth = !!(config.proxy_username || config.proxy_password);
  socket.write(Buffer.from(wantsAuth ? [0x05, 0x02, 0x00, 0x02] : [0x05, 0x01, 0x00]));
  socket.once("data", (method) => {
    if (method.length < 2 || method[0] !== 0x05) {
      reject(new Error("Invalid SOCKS greeting"));
      socket.destroy();
      return;
    }
    if (method[1] === 0x02) {
      const user = Buffer.from(config.proxy_username || "");
      const pass = Buffer.from(config.proxy_password || "");
      socket.write(Buffer.concat([Buffer.from([0x01, user.length]), user, Buffer.from([pass.length]), pass]));
      socket.once("data", (auth) => {
        if (auth.length < 2 || auth[1] !== 0x00) {
          reject(new Error("SOCKS proxy authentication failed"));
          socket.destroy();
          return;
        }
        sendSocksConnect(socket, config, resolve, reject);
      });
    } else if (method[1] === 0x00) {
      sendSocksConnect(socket, config, resolve, reject);
    } else {
      reject(new Error("SOCKS proxy rejected authentication methods"));
      socket.destroy();
    }
  });
}

function sendSocksConnect(socket: Socket, config: ConnectionConfig, resolve: (socket: Socket) => void, reject: (err: Error) => void) {
  const host = Buffer.from(config.host);
  socket.write(Buffer.concat([Buffer.from([0x05, 0x01, 0x00, 0x03, host.length]), host, portBytes(config.port)]));
  socket.once("data", (res) => {
    if (res.length < 4 || res[0] !== 0x05 || res[1] !== 0x00) {
      reject(new Error(`SOCKS proxy connect failed with code ${res[1] ?? "unknown"}`));
      socket.destroy();
      return;
    }
    resolve(socket);
  });
}

function portBytes(port: number): Buffer {
  const buf = Buffer.alloc(2);
  buf.writeUInt16BE(port);
  return buf;
}

function isMysqlType(dbType: string): boolean {
  return dbType === "mysql" || dbType === "doris" || dbType === "starrocks";
}

interface BridgeQueryResult {
  columns: string[];
  rows: unknown[][];
  affected_rows: number;
  execution_time_ms: number;
  truncated: boolean;
}

interface BridgeTableInfo {
  name: string;
  table_type: string;
  comment: string | null;
}

interface BridgeColumnInfo {
  name: string;
  data_type: string;
  is_nullable: boolean;
  column_default: string | null;
  is_primary_key: boolean;
  comment: string | null;
}

interface MongoDocumentResult {
  documents: unknown[];
  total: number;
}

function bridgeAppDataDir(): string {
  const home = homedir();
  switch (platform()) {
    case "darwin":
      return join(home, "Library", "Application Support", "com.dbx.app");
    case "win32":
      return join(process.env.APPDATA || join(home, "AppData", "Roaming"), "com.dbx.app");
    default:
      return join(home, ".config", "com.dbx.app");
  }
}

async function bridgeDataRequest<T>(path: string, body: Record<string, unknown>): Promise<T> {
  let bridgeUrl: string;
  try {
    const portFile = join(bridgeAppDataDir(), "mcp-bridge-port");
    const port = (await readFile(portFile, "utf-8")).trim();
    bridgeUrl = `http://127.0.0.1:${port}`;
  } catch {
    throw new Error("DBX desktop app is not running. This database type requires DBX to be running for query execution.");
  }
  const res = await fetch(`${bridgeUrl}${path}`, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });
  if (!res.ok) {
    const errBody = await res.text().catch(() => "");
    let errorMsg: string;
    try {
      const parsed = JSON.parse(errBody);
      errorMsg = parsed.error || errBody;
    } catch {
      errorMsg = errBody;
    }
    throw new Error(errorMsg || `Bridge request failed: ${res.status}`);
  }
  return res.json() as Promise<T>;
}

function resolveMaxRows(options?: QueryOptions): number {
  return options?.maxRows ?? MAX_ROWS;
}

function resolveTimeoutMs(options?: QueryOptions): number {
  return options?.timeoutMs ?? QUERY_TIMEOUT_MS;
}

function convertBridgeQueryResult(result: BridgeQueryResult, options?: QueryOptions): QueryResult {
  const rows = result.rows.slice(0, resolveMaxRows(options)).map((row) => {
    const obj: Record<string, unknown> = {};
    result.columns.forEach((col, i) => { obj[col] = row[i]; });
    return obj;
  });
  return { columns: result.columns, rows, row_count: rows.length };
}

function withTimeout<T>(promise: Promise<T>, ms: number): Promise<T> {
  return new Promise((resolve, reject) => {
    const timer = setTimeout(() => reject(new Error(`Query timed out after ${ms}ms`)), ms);
    promise.then(resolve, reject).finally(() => clearTimeout(timer));
  });
}

async function queryWithRetry(config: ConnectionConfig, fn: () => Promise<QueryResult>, options?: QueryOptions): Promise<QueryResult> {
  const timeoutMs = resolveTimeoutMs(options);
  try {
    return await withTimeout(fn(), timeoutMs);
  } catch (e: unknown) {
    const msg = e instanceof Error ? e.message : String(e);
    const retriable = /terminating connection|Connection lost|ECONNRESET|EPIPE|connection refused/i.test(msg);
    if (retriable) {
      const key = poolKey(config);
      const entry = pools.get(key);
      if (entry) evictPool(key, entry);
      return withTimeout(fn(), timeoutMs);
    }
    throw e;
  }
}

async function pgQuery(config: ConnectionConfig, sql: string, params?: unknown[], options?: QueryOptions): Promise<QueryResult> {
  return queryWithRetry(config, async () => {
    const pool = await getPgPool(config);
    const result = await pool.query(sql, params);
    const rows = (result.rows || []).slice(0, resolveMaxRows(options));
    return { columns: result.fields?.map((f) => f.name) ?? [], rows, row_count: rows.length };
  }, options);
}

async function mysqlQuery(config: ConnectionConfig, sql: string, params?: unknown[], options?: QueryOptions): Promise<QueryResult> {
  return queryWithRetry(config, async () => {
    const pool = await getMysqlPool(config);
    const [results, fields] = await pool.query(sql, params);
    const rows = (Array.isArray(results) ? results : []).slice(0, resolveMaxRows(options)) as Record<string, unknown>[];
    return { columns: (fields as Array<{ name: string }>)?.map((f) => f.name) ?? [], rows, row_count: rows.length };
  }, options);
}

async function query(config: ConnectionConfig, sql: string, params?: unknown[], options?: QueryOptions): Promise<QueryResult> {
  if (config.db_type === "sqlite") return sqliteQuery(config, sql, options);
  if (config.db_type === "rqlite") return rqliteQuery(config, sql, options);
  if (isMysqlType(config.db_type)) return mysqlQuery(config, sql, params, options);
  return pgQuery(config, sql, params, options);
}

function sqlitePath(config: ConnectionConfig): string {
  return expandTilde(config.host || config.database || "");
}

function expandTilde(path: string): string {
  if (path === "~") return homedir();
  if (path.startsWith("~/")) return join(homedir(), path.slice(2));
  return path;
}

function quoteSqliteIdentifier(identifier: string): string {
  return `"${identifier.replace(/"/g, '""')}"`;
}

function sqliteQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): QueryResult {
  const db = new Database(sqlitePath(config), { readonly: !sqlSafetyFromEnv().allowWrites });
  try {
    const stmt = db.prepare(sql);
    if (stmt.reader) {
      const rows = stmt.all().slice(0, resolveMaxRows(options)) as Record<string, unknown>[];
      return { columns: stmt.columns().map((column) => column.name), rows, row_count: rows.length };
    }
    const result = stmt.run();
    return { columns: [], rows: [], row_count: result.changes };
  } finally {
    db.close();
  }
}

async function rqliteQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): Promise<QueryResult> {
  const isReader = /^\s*(?:--[^\n]*\n|\s|\/\*[\s\S]*?\*\/)*(select|pragma|explain|with)\b/i.test(sql);
  const endpoint = isReader ? "/db/query" : "/db/execute";
  const result = await rqliteRequest(config, endpoint, sql);
  if (isReader) {
    const columns = result.columns ?? [];
    const rows = (result.values ?? []).slice(0, resolveMaxRows(options)).map((row) => {
      const record: Record<string, unknown> = {};
      columns.forEach((column, index) => {
        record[column] = row[index];
      });
      return record;
    });
    return { columns, rows, row_count: rows.length };
  }
  return { columns: [], rows: [], row_count: result.rows_affected ?? 0 };
}

async function rqliteRequest(config: ConnectionConfig, endpoint: "/db/query" | "/db/execute", sql: string): Promise<RqliteResult> {
  const { host, port } = await connectionEndpoint(config);
  const scheme = config.ssl ? "https" : "http";
  const params = (config.url_params || "").trim().replace(/^\?/, "");
  const url = `${scheme}://${host}:${port}${endpoint}${params ? `?${params}` : ""}`;
  const headers: Record<string, string> = { "content-type": "application/json" };
  if (config.username) {
    headers.authorization = `Basic ${Buffer.from(`${config.username}:${config.password || ""}`).toString("base64")}`;
  }
  const response = await fetch(url, {
    method: "POST",
    headers,
    body: JSON.stringify([sql]),
  });
  const text = await response.text();
  if (!response.ok) throw new Error(`rqlite error (${response.status}): ${text}`);
  const payload = JSON.parse(text) as RqliteResponse;
  const result = payload.results?.[0];
  if (!result) throw new Error("rqlite returned no result");
  if (result.error) throw new Error(`rqlite error: ${result.error}`);
  return result;
}

export async function executeQuery(config: ConnectionConfig, sql: string, options?: QueryOptions): Promise<QueryResult> {
  if (config.db_type === "mongodb") {
    const find = parseMongoFindCommand(sql);
    if (find) {
      const result = await withTimeout(mongoFindDocuments(config, find.collection, find.skip, find.limit, find.filter, find.sort), resolveTimeoutMs(options));
      return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total);
    }
    const count = parseMongoCountDocumentsCommand(sql);
    if (count) {
      const result = await withTimeout(mongoFindDocuments(config, count.collection, 0, 1, count.filter), resolveTimeoutMs(options));
      return { columns: ["count"], rows: [{ count: result.total }], row_count: 1 };
    }
    const aggregate = parseMongoAggregateCommand(sql);
    if (aggregate) {
      const safety = evaluateMongoAggregateSafety(aggregate, sqlSafetyFromEnv());
      if (!safety.allowed) throw new Error(safety.reason);
      const result = await withTimeout(
        mongoAggregateDocuments(config, aggregate.collection, aggregate.pipeline, resolveMaxRows(options)),
        resolveTimeoutMs(options),
      );
      return mongoDocumentsToQueryResult(result.documents.slice(0, resolveMaxRows(options)), result.total);
    }
    const write = parseMongoWriteCommand(sql);
    if (write) {
      const safety = evaluateMongoWriteSafety(write, sqlSafetyFromEnv());
      if (!safety.allowed) throw new Error(safety.reason);
      const affected = await withTimeout(executeMongoWrite(config, write), resolveTimeoutMs(options));
      return { columns: [], rows: [], row_count: affected };
    }
    throw new Error(
      "Use MongoDB shell-style commands, for example: db.projects.find({}).limit(100), db.projects.countDocuments({}), db.projects.insertOne({...}), db.projects.updateOne({...}, {$set: {...}}), or db.projects.deleteOne({...})",
    );
  }
  if (isDirectQueryType(config.db_type)) {
    return query(config, sql, undefined, options);
  }
  const result = await withTimeout(bridgeDataRequest<BridgeQueryResult>("/data/execute-query", {
    connection_name: config.name,
    database: config.database || "",
    sql,
  }), resolveTimeoutMs(options));
  return convertBridgeQueryResult(result, options);
}

export async function listTables(config: ConnectionConfig, schema?: string): Promise<TableInfo[]> {
  if (config.db_type === "mongodb") {
    const collections = await bridgeDataRequest<string[]>("/data/mongo/list-collections", {
      connection_name: config.name,
      database: config.database || "",
      schema: schema || "",
    });
    return collections.map((name) => ({ name, type: "COLLECTION" }));
  }
  if (config.db_type === "sqlite" || config.db_type === "rqlite") {
    const result = await query(
      config,
      `SELECT name, type FROM sqlite_master WHERE type IN ('table', 'view') AND name NOT LIKE 'sqlite_%' ORDER BY name`,
    );
    return result.rows.map((r) => ({ name: String(r.name || ""), type: String(r.type || "table") }));
  }
  if (!isDirectQueryType(config.db_type)) {
    const tables = await bridgeDataRequest<BridgeTableInfo[]>("/data/list-tables", {
      connection_name: config.name,
      database: config.database || "",
      schema: schema || "",
    });
    return tables.map((t) => ({ name: t.name, type: t.table_type || "TABLE" }));
  }
  let result: QueryResult;
  if (isMysqlType(config.db_type)) {
    result = await query(config, `SELECT TABLE_NAME AS name, TABLE_TYPE AS type FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() ORDER BY TABLE_NAME`);
  } else {
    result = await query(
      config,
      `SELECT table_name AS name, table_type AS type FROM information_schema.tables WHERE table_schema = $1 ORDER BY table_name`,
      [schema || "public"],
    );
  }
  return result.rows.map((r) => ({ name: String(r.name || r.NAME), type: String(r.type || r.TYPE || "TABLE") }));
}

export async function describeTable(config: ConnectionConfig, table: string, schema?: string): Promise<ColumnInfo[]> {
  if (config.db_type === "mongodb") {
    const result = await mongoFindDocuments(config, table, 0, 20, "{}");
    return inferMongoColumns(result.documents);
  }
  if (config.db_type === "sqlite" || config.db_type === "rqlite") {
    const result = await query(config, `PRAGMA table_info(${quoteSqliteIdentifier(table)})`);
    return result.rows.map((r) => ({
      name: String(r.name || ""),
      data_type: String(r.type || ""),
      is_nullable: Number(r.notnull || 0) === 0,
      column_default: r.dflt_value != null ? String(r.dflt_value) : null,
      is_primary_key: Number(r.pk || 0) > 0,
      comment: null,
    }));
  }
  if (!isDirectQueryType(config.db_type)) {
    const columns = await bridgeDataRequest<BridgeColumnInfo[]>("/data/describe-table", {
      connection_name: config.name,
      database: config.database || "",
      schema: schema || "",
      table,
    });
    return columns.map((c) => ({
      name: c.name,
      data_type: c.data_type,
      is_nullable: c.is_nullable,
      column_default: c.column_default,
      is_primary_key: c.is_primary_key,
      comment: c.comment,
    }));
  }
  let result: QueryResult;
  if (isMysqlType(config.db_type)) {
    result = await query(
      config,
      `SELECT c.COLUMN_NAME AS name, c.DATA_TYPE AS data_type, c.IS_NULLABLE = 'YES' AS is_nullable, c.COLUMN_DEFAULT AS column_default, c.COLUMN_KEY = 'PRI' AS is_primary_key, c.COLUMN_COMMENT AS comment FROM information_schema.COLUMNS c WHERE c.TABLE_SCHEMA = DATABASE() AND c.TABLE_NAME = ? ORDER BY c.ORDINAL_POSITION`,
      [table],
    );
  } else {
    result = await query(
      config,
      `SELECT c.column_name AS name, c.data_type, c.is_nullable = 'YES' AS is_nullable, c.column_default, CASE WHEN tc.constraint_type = 'PRIMARY KEY' THEN true ELSE false END AS is_primary_key, col_description(cls.oid, c.ordinal_position) AS comment FROM information_schema.columns c LEFT JOIN information_schema.key_column_usage kcu ON kcu.table_schema = c.table_schema AND kcu.table_name = c.table_name AND kcu.column_name = c.column_name LEFT JOIN information_schema.table_constraints tc ON tc.constraint_name = kcu.constraint_name AND tc.table_schema = kcu.table_schema AND tc.constraint_type = 'PRIMARY KEY' LEFT JOIN pg_class cls ON cls.relname = c.table_name AND cls.relnamespace = (SELECT oid FROM pg_namespace WHERE nspname = c.table_schema) WHERE c.table_schema = $1 AND c.table_name = $2 ORDER BY c.ordinal_position`,
      [schema || "public", table],
    );
  }
  return result.rows.map((r) => ({
    name: String(r.name || ""),
    data_type: String(r.data_type || ""),
    is_nullable: Boolean(r.is_nullable),
    column_default: r.column_default != null ? String(r.column_default) : null,
    is_primary_key: Boolean(r.is_primary_key),
    comment: r.comment != null ? String(r.comment) : null,
  }));
}

async function mongoFindDocuments(
  config: ConnectionConfig,
  collection: string,
  skip: number,
  limit: number,
  filter: string,
  sort?: string,
): Promise<MongoDocumentResult> {
  return bridgeDataRequest<MongoDocumentResult>("/data/mongo/find-documents", {
    connection_name: config.name,
    database: config.database || "",
    collection,
    skip,
    limit,
    filter,
    sort,
  });
}

async function executeMongoWrite(config: ConnectionConfig, command: MongoWriteCommand): Promise<number> {
  if (command.kind === "insert") {
    const result = await bridgeDataRequest<{ affected_rows: number }>("/data/mongo/insert-documents", {
      connection_name: config.name,
      database: config.database || "",
      collection: command.collection,
      docs_json: command.docsJson,
    });
    return result.affected_rows;
  }
  if (command.kind === "update") {
    const result = await bridgeDataRequest<{ affected_rows: number }>("/data/mongo/update-documents", {
      connection_name: config.name,
      database: config.database || "",
      collection: command.collection,
      filter_json: command.filter,
      update_json: command.update,
      many: command.many,
    });
    return result.affected_rows;
  }
  const result = await bridgeDataRequest<{ affected_rows: number }>("/data/mongo/delete-documents", {
    connection_name: config.name,
    database: config.database || "",
    collection: command.collection,
    filter_json: command.filter,
    many: command.many,
  });
  return result.affected_rows;
}

async function mongoAggregateDocuments(
  config: ConnectionConfig,
  collection: string,
  pipelineJson: string,
  maxRows: number,
): Promise<MongoDocumentResult> {
  return bridgeDataRequest<MongoDocumentResult>("/data/mongo/aggregate-documents", {
    connection_name: config.name,
    database: config.database || "",
    collection,
    pipeline_json: pipelineJson,
    max_rows: maxRows,
  });
}

export function mongoDocumentsToQueryResult(documents: unknown[], total: number): QueryResult {
  const columns: string[] = [];
  for (const doc of documents) {
    if (isRecord(doc)) {
      for (const key of Object.keys(doc)) {
        if (!columns.includes(key)) columns.push(key);
      }
    } else if (!columns.includes("value")) {
      columns.push("value");
    }
  }
  const rows = documents.map((doc) => {
    const row: Record<string, unknown> = {};
    for (const column of columns) {
      row[column] = isRecord(doc) ? toCellValue(doc[column]) : column === "value" ? toCellValue(doc) : null;
    }
    return row;
  });
  return { columns, rows, row_count: rows.length };
}

export function inferMongoColumns(documents: unknown[]): ColumnInfo[] {
  const columns = new Map<string, { types: Set<string>; nullable: boolean }>();
  for (const doc of documents) {
    if (!isRecord(doc)) {
      const entry = columns.get("value") ?? { types: new Set<string>(), nullable: false };
      entry.types.add(mongoTypeName(doc));
      columns.set("value", entry);
      continue;
    }
    for (const [name, value] of Object.entries(doc)) {
      const entry = columns.get(name) ?? { types: new Set<string>(), nullable: false };
      entry.types.add(mongoTypeName(value));
      if (value === null || value === undefined) entry.nullable = true;
      columns.set(name, entry);
    }
  }
  return Array.from(columns.entries()).map(([name, entry]) => ({
    name,
    data_type: Array.from(entry.types).sort().join(" | ") || "unknown",
    is_nullable: entry.nullable,
    column_default: null,
    is_primary_key: name === "_id",
    comment: null,
  }));
}

interface MongoFindCommand {
  collection: string;
  filter: string;
  skip: number;
  limit: number;
  sort?: string;
}

interface MongoCountDocumentsCommand {
  collection: string;
  filter: string;
}

interface MongoAggregateCommand {
  collection: string;
  pipeline: string;
}

export type MongoWriteCommand =
  | { kind: "insert"; collection: string; docsJson: string }
  | { kind: "update"; collection: string; filter: string; update: string; many: boolean }
  | { kind: "delete"; collection: string; filter: string; many: boolean };

export function parseMongoFindCommand(input: string): MongoFindCommand | null {
  const source = input.trim().replace(/;$/, "").trim();
  const target = parseCollectionMethodTarget(source, "find");
  if (!target) return null;
  const findOpenIndex = source.indexOf("(", target.methodCallIndex);
  const findCloseIndex = findMatchingParen(source, findOpenIndex);
  if (findCloseIndex < 0) return null;
  const findArgs = splitTopLevel(source.slice(findOpenIndex + 1, findCloseIndex));
  const filter = normalizeJsonArgument(findArgs[0] || "{}");
  if (!filter) return null;
  const chain = source.slice(findCloseIndex + 1).trim();
  if (chain && !chain.startsWith(".")) return null;
  const sortArg = readChainedCallArgument(chain, "sort");
  let sort: string | undefined;
  if (sortArg !== undefined) {
    const parsedSort = normalizeJsonArgument(sortArg);
    if (!parsedSort) return null;
    sort = parsedSort;
  }
  const skip = readChainedIntegerArgument(chain, "skip", 0);
  const limit = readChainedIntegerArgument(chain, "limit", MAX_ROWS);
  if (skip === null || limit === null) return null;
  return { collection: target.collection, filter, skip, limit, sort };
}

export function parseMongoCountDocumentsCommand(input: string): MongoCountDocumentsCommand | null {
  const source = input.trim().replace(/;$/, "").trim();
  const target = parseCollectionMethodTarget(source, "countDocuments");
  if (!target) return null;
  const openIndex = source.indexOf("(", target.methodCallIndex);
  const closeIndex = findMatchingParen(source, openIndex);
  if (closeIndex < 0 || source.slice(closeIndex + 1).trim()) return null;
  const args = splitTopLevel(source.slice(openIndex + 1, closeIndex));
  if (args.length > 1 && args.slice(1).some((arg) => arg.trim())) return null;
  const filter = normalizeJsonArgument(args[0] || "{}");
  return filter ? { collection: target.collection, filter } : null;
}

export function parseMongoAggregateCommand(input: string): MongoAggregateCommand | null {
  const source = input.trim().replace(/;$/, "").trim();
  const target = parseCollectionMethodTarget(source, "aggregate");
  if (!target) return null;
  const args = parseMethodArgs(source, target.methodCallIndex);
  if (!args || args.length !== 1) return null;
  const pipeline = normalizeJsonArgument(args[0]);
  if (!pipeline) return null;
  return Array.isArray(JSON.parse(pipeline)) ? { collection: target.collection, pipeline } : null;
}

export function mongoAggregateWriteStage(pipelineJson: string): "$out" | "$merge" | null {
  try {
    const pipeline = JSON.parse(pipelineJson);
    if (!Array.isArray(pipeline)) return null;
    for (const stage of pipeline) {
      if (!isRecord(stage)) continue;
      if (Object.prototype.hasOwnProperty.call(stage, "$out")) return "$out";
      if (Object.prototype.hasOwnProperty.call(stage, "$merge")) return "$merge";
    }
    return null;
  } catch {
    return null;
  }
}

export function parseMongoWriteCommand(input: string): MongoWriteCommand | null {
  const source = input.trim().replace(/;$/, "").trim();
  const insertOne = parseCollectionMethodTarget(source, "insertOne");
  if (insertOne) {
    const args = parseMethodArgs(source, insertOne.methodCallIndex);
    if (!args || args.length !== 1) return null;
    const doc = normalizeJsonArgument(args[0]);
    return doc ? { kind: "insert", collection: insertOne.collection, docsJson: doc } : null;
  }

  const insertMany = parseCollectionMethodTarget(source, "insertMany");
  if (insertMany) {
    const args = parseMethodArgs(source, insertMany.methodCallIndex);
    if (!args || args.length !== 1) return null;
    const docs = normalizeJsonArgument(args[0]);
    if (!docs) return null;
    return Array.isArray(JSON.parse(docs)) ? { kind: "insert", collection: insertMany.collection, docsJson: docs } : null;
  }

  for (const method of ["updateOne", "updateMany"] as const) {
    const target = parseCollectionMethodTarget(source, method);
    if (!target) continue;
    const args = parseMethodArgs(source, target.methodCallIndex);
    if (!args || args.length !== 2) return null;
    const filter = normalizeJsonArgument(args[0]);
    const update = normalizeJsonArgument(args[1]);
    if (!filter || !update) return null;
    return { kind: "update", collection: target.collection, filter, update, many: method === "updateMany" };
  }

  for (const method of ["deleteOne", "deleteMany"] as const) {
    const target = parseCollectionMethodTarget(source, method);
    if (!target) continue;
    const args = parseMethodArgs(source, target.methodCallIndex);
    if (!args || args.length !== 1) return null;
    const filter = normalizeJsonArgument(args[0]);
    if (!filter) return null;
    return { kind: "delete", collection: target.collection, filter, many: method === "deleteMany" };
  }

  return null;
}

export function evaluateMongoWriteSafety(
  command: MongoWriteCommand,
  options: { allowWrites?: boolean; allowDangerous?: boolean },
): { allowed: boolean; reason?: string } {
  if (!options.allowWrites) {
    return {
      allowed: false,
      reason: "MCP MongoDB execution is read-only by default. Set DBX_MCP_ALLOW_WRITES=1 to allow write commands.",
    };
  }
  if (!options.allowDangerous && (command.kind === "update" || command.kind === "delete") && isEmptyJsonObject(command.filter)) {
    return {
      allowed: false,
      reason: "MongoDB update/delete commands must include a non-empty filter unless DBX_MCP_ALLOW_DANGEROUS_SQL=1 is set.",
    };
  }
  return { allowed: true };
}

export function evaluateMongoAggregateSafety(
  command: MongoAggregateCommand,
  options: { allowWrites?: boolean; allowDangerous?: boolean },
): { allowed: boolean; reason?: string } {
  const writeStage = mongoAggregateWriteStage(command.pipeline);
  if (!writeStage) return { allowed: true };
  if (!options.allowWrites) {
    return {
      allowed: false,
      reason: `MongoDB aggregate stage "${writeStage}" writes data. Set DBX_MCP_ALLOW_WRITES=1 to allow write commands.`,
    };
  }
  if (!options.allowDangerous) {
    return {
      allowed: false,
      reason: `MongoDB aggregate stage "${writeStage}" is dangerous. Set DBX_MCP_ALLOW_DANGEROUS_SQL=1 to allow it.`,
    };
  }
  return { allowed: true };
}

function parseCollectionMethodTarget(source: string, method: string): { collection: string; methodCallIndex: number } | null {
  const direct = new RegExp(`^db\\.([A-Za-z_$][\\w$]*)\\.${method}\\s*\\(`).exec(source);
  if (direct) return { collection: direct[1], methodCallIndex: source.indexOf(`.${method}`) };
  const quoted = new RegExp(`^db\\.getCollection\\(\\s*(['"])([^'"]+)\\1\\s*\\)\\.${method}\\s*\\(`).exec(source);
  if (quoted) return { collection: quoted[2], methodCallIndex: source.indexOf(`.${method}`) };
  return null;
}

function parseMethodArgs(source: string, methodCallIndex: number): string[] | null {
  const openIndex = source.indexOf("(", methodCallIndex);
  const closeIndex = findMatchingParen(source, openIndex);
  if (closeIndex < 0 || source.slice(closeIndex + 1).trim()) return null;
  return splitTopLevel(source.slice(openIndex + 1, closeIndex));
}

function readChainedCallArgument(chain: string, method: string): string | undefined {
  const pattern = new RegExp(`\\.${method}\\s*\\(`, "g");
  const match = pattern.exec(chain);
  if (!match) return undefined;
  const openIndex = match.index + match[0].lastIndexOf("(");
  const closeIndex = findMatchingParen(chain, openIndex);
  return closeIndex < 0 ? undefined : chain.slice(openIndex + 1, closeIndex);
}

function readChainedIntegerArgument(chain: string, method: string, fallback: number): number | null {
  const arg = readChainedCallArgument(chain, method);
  if (arg === undefined) return fallback;
  if (!/^\d+$/.test(arg.trim())) return null;
  return Number(arg.trim());
}

function normalizeJsonArgument(arg: string): string | null {
  const value = quoteUnquotedObjectKeys(
    convertSingleQuotedStrings((arg.trim() || "{}").replace(/ObjectId\s*\(\s*["']([^"']+)["']\s*\)/g, '{"$oid":"$1"}')),
  );
  try {
    JSON.parse(value);
    return value;
  } catch {
    return null;
  }
}

function convertSingleQuotedStrings(source: string): string {
  let result = "";
  let copiedUntil = 0;
  let quote: string | null = null;
  let start = 0;
  let value = "";
  let escaped = false;

  for (let i = 0; i < source.length; i += 1) {
    const char = source[i];
    if (!quote) {
      if (char === "'") {
        quote = char;
        start = i;
        value = "";
        escaped = false;
      } else if (char === '"') {
        quote = char;
      }
      continue;
    }

    if (quote === '"') {
      if (escaped) escaped = false;
      else if (char === "\\") escaped = true;
      else if (char === '"') quote = null;
      continue;
    }

    if (escaped) {
      value += char;
      escaped = false;
    } else if (char === "\\") {
      escaped = true;
    } else if (char === "'") {
      result += source.slice(copiedUntil, start) + JSON.stringify(value);
      copiedUntil = i + 1;
      quote = null;
    } else {
      value += char;
    }
  }

  return quote === "'" ? source : result + source.slice(copiedUntil);
}

function quoteUnquotedObjectKeys(source: string): string {
  let result = "";
  let quote: string | null = null;
  let escaped = false;

  for (let i = 0; i < source.length; i += 1) {
    const char = source[i];
    if (quote) {
      result += char;
      if (escaped) escaped = false;
      else if (char === "\\") escaped = true;
      else if (char === quote) quote = null;
      continue;
    }

    if (char === '"' || char === "'") {
      quote = char;
      result += char;
      continue;
    }

    if (/[A-Za-z_$]/.test(char) && shouldQuoteObjectKey(source, i)) {
      let end = i + 1;
      while (/[\w$]/.test(source[end] || "")) end += 1;
      result += `"${source.slice(i, end)}"`;
      i = end - 1;
      continue;
    }

    result += char;
  }

  return result;
}

function shouldQuoteObjectKey(source: string, index: number): boolean {
  let before = index - 1;
  while (/\s/.test(source[before] || "")) before -= 1;
  if (source[before] !== "{" && source[before] !== ",") return false;

  let after = index + 1;
  while (/[\w$]/.test(source[after] || "")) after += 1;
  while (/\s/.test(source[after] || "")) after += 1;
  return source[after] === ":";
}

function isEmptyJsonObject(json: string): boolean {
  try {
    const parsed = JSON.parse(json);
    return isRecord(parsed) && Object.keys(parsed).length === 0;
  } catch {
    return false;
  }
}

function splitTopLevel(source: string): string[] {
  const parts: string[] = [];
  let depth = 0;
  let start = 0;
  let quote: string | null = null;
  for (let i = 0; i < source.length; i += 1) {
    const ch = source[i];
    if (quote) {
      if (ch === "\\" && i + 1 < source.length) i += 1;
      else if (ch === quote) quote = null;
      continue;
    }
    if (ch === "'" || ch === '"') quote = ch;
    else if (ch === "{" || ch === "[" || ch === "(") depth += 1;
    else if (ch === "}" || ch === "]" || ch === ")") depth -= 1;
    else if (ch === "," && depth === 0) {
      parts.push(source.slice(start, i));
      start = i + 1;
    }
  }
  parts.push(source.slice(start));
  return parts;
}

function findMatchingParen(source: string, openIndex: number): number {
  if (openIndex < 0 || source[openIndex] !== "(") return -1;
  let depth = 0;
  let quote: string | null = null;
  for (let i = openIndex; i < source.length; i += 1) {
    const ch = source[i];
    if (quote) {
      if (ch === "\\" && i + 1 < source.length) i += 1;
      else if (ch === quote) quote = null;
      continue;
    }
    if (ch === "'" || ch === '"') quote = ch;
    else if (ch === "(") depth += 1;
    else if (ch === ")") {
      depth -= 1;
      if (depth === 0) return i;
    }
  }
  return -1;
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

function mongoTypeName(value: unknown): string {
  if (value === null || value === undefined) return "null";
  if (Array.isArray(value)) return "array";
  if (isRecord(value)) return "object";
  return typeof value;
}

function toCellValue(value: unknown): unknown {
  return typeof value === "object" && value !== null ? JSON.stringify(value) : value;
}