import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { z } from "zod";
import fs from "node:fs";
import path from "node:path";
import mysql from "mysql2/promise";

const ENV_PATH =
  process.env.BDB_ENV ??
  "C:\\Users\\ahuja\\Westrock\\bdb-access\\bdb-db-readonly.env";

const SQL_DIR =
  process.env.BDB_SQL_DIR ??
  "C:\\Users\\ahuja\\Westrock\\bdb-access-handoff\\sql";

const QUERY_TIMEOUT_MS = 90_000;
const MAX_ROWS = 2000;

function loadEnvFile(filePath: string): Record<string, string> {
  const raw = fs.readFileSync(filePath, "utf-8");
  const out: Record<string, string> = {};
  for (const line of raw.split(/\r?\n/)) {
    const trimmed = line.trim();
    if (!trimmed || trimmed.startsWith("#")) continue;
    const eq = trimmed.indexOf("=");
    if (eq === -1) continue;
    out[trimmed.slice(0, eq).trim()] = trimmed.slice(eq + 1).trim();
  }
  return out;
}

function getConfig() {
  const env = loadEnvFile(ENV_PATH);
  const required = [
    "BDB_DB_HOST",
    "BDB_DB_PORT",
    "BDB_DB_USER",
    "BDB_DB_PASSWORD",
    "BDB_DB_NAME",
  ];
  for (const key of required) {
    if (!env[key]) throw new Error(`Missing ${key} in ${ENV_PATH}`);
  }
  return {
    host: env.BDB_DB_HOST,
    port: Number(env.BDB_DB_PORT),
    user: env.BDB_DB_USER,
    password: env.BDB_DB_PASSWORD,
    database: env.BDB_DB_NAME,
  };
}

/**
 * Only a single SELECT/WITH/SHOW/EXPLAIN/DESCRIBE statement is allowed.
 * A trailing semicolon is tolerated but any statement separator beyond
 * that is rejected, since bdb_ro has no write grants anyway but we also
 * want to fail fast and not rely solely on DB-side permissions.
 */
function assertReadOnlyStatement(sql: string) {
  const body = sql.trim().replace(/;+\s*$/, "");
  if (body.includes(";")) {
    throw new Error("Only a single statement is allowed (no ';' inside the query).");
  }
  const head = body.slice(0, 20).trimStart().toUpperCase();
  const allowed = ["SELECT", "WITH", "SHOW", "EXPLAIN", "DESCRIBE", "DESC "];
  if (!allowed.some((kw) => head.startsWith(kw))) {
    throw new Error(
      `Only SELECT / WITH / SHOW / EXPLAIN / DESCRIBE statements are allowed. Got: ${head}...`
    );
  }
  const forbidden = /\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|REPLACE|GRANT|REVOKE|CALL|LOCK|SET\s+(?!SESSION\s+TRANSACTION))\b/i;
  if (forbidden.test(body)) {
    throw new Error("Query contains a disallowed keyword for a read-only connection.");
  }
}

async function runReadOnlyQuery(sql: string) {
  assertReadOnlyStatement(sql);
  const config = getConfig();
  const conn = await mysql.createConnection({
    host: config.host,
    port: config.port,
    user: config.user,
    password: config.password,
    database: config.database,
    timezone: "Z",
  });
  try {
    await conn.query("SET SESSION TRANSACTION READ ONLY");
    const [rows] = await conn.query({ sql, timeout: QUERY_TIMEOUT_MS });
    const rowArray = Array.isArray(rows) ? rows : [rows];
    const truncated = rowArray.length > MAX_ROWS;
    return {
      rowCount: rowArray.length,
      truncated,
      rows: truncated ? rowArray.slice(0, MAX_ROWS) : rowArray,
    };
  } finally {
    await conn.end();
  }
}

function readSqlFile(name: string): string {
  const full = path.join(SQL_DIR, name);
  return fs.readFileSync(full, "utf-8");
}

export function registerBdbTools(server: McpServer) {
  server.registerTool(
    "bdb_run_query",
    {
      title: "Run a read-only SQL query against BDB",
      description:
        "Run a single SELECT/WITH/SHOW/EXPLAIN/DESCRIBE statement against the BDB production " +
        "database (bdbHaAppDb) using the read-only bdb_ro user. WestRock's companyID is 24227 " +
        "— filter orders on orders.companyID = 24227. 90-second timeout, max 2000 rows returned. " +
        "Any write/DDL statement is rejected before it reaches the database.",
      inputSchema: {
        sql: z.string().describe("A single read-only SQL statement"),
      },
    },
    async ({ sql }) => {
      const result = await runReadOnlyQuery(sql);
      return {
        content: [{ type: "text", text: JSON.stringify(result, null, 2) }],
      };
    }
  );

  server.registerTool(
    "bdb_pending_jobs_report",
    {
      title: "Pending jobs: status, hours in status, due dates",
      description:
        "Runs the saved report listing pending WestRock jobs with current status, time in " +
        "status, due dates, client, address, appraiser and last note.",
      inputSchema: {},
    },
    async () => {
      const sql = readSqlFile("pending_jobs_status_hours_due.sql");
      const result = await runReadOnlyQuery(sql);
      return {
        content: [{ type: "text", text: JSON.stringify(result, null, 2) }],
      };
    }
  );

  server.registerTool(
    "bdb_job_workflow_roles_report",
    {
      title: "People on each pending job (workflow roles)",
      description:
        "Runs the saved report listing who is assigned to each pending WestRock job " +
        "(Team Leader, Appraiser, Appraiser Analyst, Programmer, Inspector).",
      inputSchema: {},
    },
    async () => {
      const sql = readSqlFile("pending_jobs_workflow_roles.sql");
      const result = await runReadOnlyQuery(sql);
      return {
        content: [{ type: "text", text: JSON.stringify(result, null, 2) }],
      };
    }
  );
}