import "dotenv/config";
import { readFile, writeFile } from "node:fs/promises";
import mysql, { type RowDataPacket } from "mysql2/promise";

// drizzle-kit 0.31 introspects MariaDB as if it were MySQL, which breaks three things:
// - MariaDB returns COLUMN_DEFAULT quoted ('abc', and 'NULL' for no default), so defaults come
//   out as .default('NULL') or broken literals like '''Z'''.
// - drizzle-kit reads the primary key rows as row["TABLE_NAME"], but MariaDB returns lowercase
//   keys, so every primary key is dropped.
// - The import line misses some column builders (bigint, tinyint).
// This script rewrites db/amigo/schema.ts from information_schema after each pull.
const SCHEMA_PATH = "db/amigo/schema.ts";

// `.default(...)` at the end of a column line, optionally followed by `.notNull()`/`.primaryKey()`.
const DEFAULT_CALL =
  /\.default\(.*?\)(?=(?:\.notNull\(\))?(?:\.primaryKey\(\))?,$)/;

type ColumnRow = RowDataPacket & {
  tableName: string;
  columnName: string;
  columnDefault: string | null;
  isPrimary: number;
};

function fixDefault(line: string, rawDefault: string | null): string {
  if (rawDefault === null || rawDefault === "NULL") {
    return line.replace(DEFAULT_CALL, "");
  }
  if (rawDefault.startsWith("'") && rawDefault.endsWith("'")) {
    const value = rawDefault.slice(1, -1).replaceAll("''", "'");
    return line.replace(DEFAULT_CALL, `.default(${JSON.stringify(value)})`);
  }
  // Numbers and expressions like current_timestamp() are already generated correctly.
  return line;
}

function addPrimaryKey(line: string): string {
  if (line.includes(".primaryKey()")) return line;
  return line.replace(/,$/, ".primaryKey(),");
}

// Rebuilds the import lines from the builders the schema actually calls.
function rebuildImports(schema: string): string {
  const used = new Set(["mysqlTable"]);
  for (const match of schema.matchAll(/^\t\w+: (\w+)\(/gm)) used.add(match[1]);
  for (const match of schema.matchAll(/^\t(index|uniqueIndex|primaryKey)\(/gm))
    used.add(match[1]);

  const body = schema.replace(/^import .*\n/gm, "");
  const imports = [
    `import { ${[...used].join(", ")} } from "drizzle-orm/mysql-core";`,
  ];
  if (body.includes("sql`")) imports.push('import { sql } from "drizzle-orm";');
  return `${imports.join("\n")}\n${body}`;
}

async function main() {
  const url = process.env.AMIGO_DATABASE_URL;
  if (!url) throw new Error("AMIGO_DATABASE_URL is not set in .env");

  const connection = await mysql.createConnection(url);
  const [rows] = await connection.query<ColumnRow[]>(
    `SELECT c.table_name AS tableName,
            c.column_name AS columnName,
            c.column_default AS columnDefault,
            c.column_key = 'PRI' AS isPrimary
     FROM information_schema.columns c
     WHERE c.table_schema = DATABASE()
     ORDER BY c.table_name, c.ordinal_position`,
  );
  await connection.end();

  const tables = new Map<string, ColumnRow[]>();
  for (const row of rows) {
    const columns = tables.get(row.tableName) ?? [];
    columns.push(row);
    tables.set(row.tableName, columns);
  }

  let schema = await readFile(SCHEMA_PATH, "utf8");

  for (const [tableName, columns] of tables) {
    const start = schema.indexOf(`mysqlTable("${tableName}", {`);
    if (start === -1)
      throw new Error(`Table ${tableName} not found in ${SCHEMA_PATH}`);
    const next = schema.indexOf("\nexport const ", start);
    const end = next === -1 ? schema.length : next;
    let block = schema.slice(start, end);

    const primaryColumns = columns
      .filter((c) => c.isPrimary)
      .map((c) => c.columnName);

    for (const column of columns) {
      const linePattern = new RegExp(`^\\t${column.columnName}: .*,$`, "m");
      const line = block.match(linePattern)?.[0];
      if (!line)
        throw new Error(`Column ${tableName}.${column.columnName} not found`);

      let fixed = fixDefault(line, column.columnDefault);
      if (primaryColumns.length === 1 && column.isPrimary)
        fixed = addPrimaryKey(fixed);
      block = block.replace(line, () => fixed);
    }

    if (primaryColumns.length > 1 && !block.includes("primaryKey({")) {
      if (!block.includes("(table) => [\n")) {
        throw new Error(
          `Composite key on ${tableName} needs a (table) => [...] block`,
        );
      }
      const columnsList = primaryColumns
        .map((name) => `table.${name}`)
        .join(", ");
      block = block.replace(
        "(table) => [\n",
        `(table) => [\n\tprimaryKey({ columns: [${columnsList}] }),\n`,
      );
    }

    schema = schema.slice(0, start) + block + schema.slice(end);
  }

  await writeFile(SCHEMA_PATH, rebuildImports(schema));
  console.log(
    `Fixed defaults and primary keys for ${tables.size} tables in ${SCHEMA_PATH}`,
  );
}

main().catch((error) => {
  console.error(error);
  process.exit(1);
});