/* Copyright 2026 Marimo. All rights reserved. */

import { describe, expect, it } from "vitest";
import type { SQLTableContext } from "@/core/datasets/data-source-connections";
import { DUCKDB_ENGINE } from "@/core/datasets/engines";
import type {
  Database,
  DatabaseSchema,
  DataTable,
  DataTableColumn,
} from "@/core/kernel/messages";
import {
  schemaSubtreeMatchesSearch,
  shouldExpandDatabaseForSearch,
  sqlCode,
  tableUniqueId,
} from "../utils";

describe("sqlCode", () => {
  const mockTable: DataTable = {
    name: "users" as const,
    columns: [],
    source: "local",
    source_type: "local",
    type: "table",
    engine: null,
    indexes: null,
    num_columns: null,
    num_rows: null,
    variable_name: null,
    primary_keys: null,
  };
  const mockColumn = {
    name: "email" as const,
  } as DataTableColumn;

  describe("basic SQL generation", () => {
    it("should generate basic SQL without sqlTableContext", () => {
      const result = sqlCode({ table: mockTable, columnName: mockColumn.name });
      expect(result).toBe(
        "_df = mo.sql(f'SELECT \"email\" FROM users LIMIT 100')",
      );
    });

    it("should generate SQL with default schema", () => {
      const sqlTableContext: SQLTableContext = {
        engine: DUCKDB_ENGINE,
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "duckdb",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "users" LIMIT 100\n""")',
      );
    });

    it("should generate SQL with non-default schema", () => {
      const sqlTableContext: SQLTableContext = {
        engine: DUCKDB_ENGINE,
        schema: "analytics",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "duckdb",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "analytics"."users" LIMIT 100\n""")',
      );
    });

    it("should generate SQL with non-default engine", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "snowflake",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "snowflake",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT email FROM users LIMIT 100\n""", engine=snowflake)',
      );
    });

    it("should generate SQL with non-default database", () => {
      const sqlTableContext: SQLTableContext = {
        engine: DUCKDB_ENGINE,
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "memory",
        database: "remote",
        dialect: "duckdb",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."users" LIMIT 100\n""")',
      );
    });

    it("should generate SQL for schemaless tables", () => {
      const sqlTableContext: SQLTableContext = {
        engine: DUCKDB_ENGINE,
        schema: "",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "duckdb",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "users" LIMIT 100\n""")',
      );

      const sqlTableContext2: SQLTableContext = {
        engine: DUCKDB_ENGINE,
        schema: "",
        defaultDatabase: "remote",
        database: "another_db",
        dialect: "duckdb",
      };

      const result2 = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext: sqlTableContext2,
      });
      expect(result2).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "another_db"."users" LIMIT 100\n""")',
      );
    });
  });

  describe("BigQuery dialect", () => {
    it("should use backticks for table names", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "bigquery",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "bigquery",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT email FROM `users` LIMIT 100\n""", engine=bigquery)',
      );
    });

    it("should use backticks for database.schema.table", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "bigquery",
        schema: "sales",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "remote",
        dialect: "bigquery",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT email FROM `remote.sales.users` LIMIT 100\n""", engine=bigquery)',
      );
    });

    it("should handle case-insensitive dialect name", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "bigquery",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "BigQuery",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT email FROM `users` LIMIT 100\n""", engine=bigquery)',
      );
    });
  });

  describe("MSSQL dialect", () => {
    it("should use TOP 100 instead of LIMIT", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "mssql",
        schema: "dbo",
        defaultSchema: "dbo",
        defaultDatabase: "master",
        database: "master",
        dialect: "mssql",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT TOP 100 email FROM users\n""", engine=mssql)',
      );
    });

    it("should use TOP 100 with database and schema prefix", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "mssql",
        schema: "sales",
        defaultSchema: "dbo",
        defaultDatabase: "master",
        database: "analytics",
        dialect: "mssql",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT TOP 100 email FROM analytics.sales.users\n""", engine=mssql)',
      );
    });
  });

  describe("TimescaleDB and PostgreSQL dialect", () => {
    it("should wrap table name with double quotes", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "timescaledb",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "timescaledb",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "users" LIMIT 100\n""", engine=timescaledb)',
      );
    });

    it("should wrap database, schema, and table name with double quotes", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "timescaledb",
        schema: "sales",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "remote",
        dialect: "timescaledb",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."sales"."users" LIMIT 100\n""", engine=timescaledb)',
      );
    });

    it("should handle schemaless database with double quotes", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "postgres",
        schema: "",
        defaultDatabase: "mydb",
        database: "remote",
        dialect: "postgres",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."users" LIMIT 100\n""", engine=postgres)',
      );
    });

    it("should preserve dots inside quoted schema names", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "postgres",
        schema: "analytics.events",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "remote",
        dialect: "postgres",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "email" FROM "remote"."analytics.events"."users" LIMIT 100\n""", engine=postgres)',
      );
    });

    it("should not quote * column name", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "postgres",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "postgres",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: "*",
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT * FROM "users" LIMIT 100\n""", engine=postgres)',
      );
    });
  });

  describe("Dremio dialect", () => {
    it("should quote reserved column names and table path parts", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "dremio_conn",
        schema: "operations",
        defaultSchema: "",
        defaultDatabase: "",
        database: "lakehouse",
        dialect: "dremio",
      };

      const result = sqlCode({
        table: { ...mockTable, name: "shipments" as const },
        columnName: "order",
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT "order" FROM "lakehouse"."operations"."shipments" LIMIT 100\n""", engine=dremio_conn)',
      );
    });

    it("should not quote * column name", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "dremio_conn",
        schema: "operations",
        defaultSchema: "",
        defaultDatabase: "",
        database: "lakehouse",
        dialect: "dremio",
      };

      const result = sqlCode({
        table: { ...mockTable, name: "customers" as const },
        columnName: "*",
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT * FROM "lakehouse"."operations"."customers" LIMIT 100\n""", engine=dremio_conn)',
      );
    });

    it("should preserve dots inside quoted schema names", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "dremio_conn",
        schema: "samples.dremio.com",
        defaultSchema: "",
        defaultDatabase: "",
        database: "Samples",
        dialect: "dremio",
      };

      const result = sqlCode({
        table: { ...mockTable, name: "airlines" as const },
        columnName: "*",
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT * FROM "Samples"."samples.dremio.com"."airlines" LIMIT 100\n""", engine=dremio_conn)',
      );
    });
  });

  describe("fallback behavior", () => {
    it("should use default formatter for unknown dialect", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "postgres",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "unknown_dialect",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT email FROM users LIMIT 100\n""", engine=postgres)',
      );
    });

    it("should use default formatter when dialect is not provided", () => {
      const sqlTableContext: SQLTableContext = {
        engine: "postgres",
        schema: "public",
        defaultSchema: "public",
        defaultDatabase: "mydb",
        database: "mydb",
        dialect: "",
      };

      const result = sqlCode({
        table: mockTable,
        columnName: mockColumn.name,
        sqlTableContext,
      });
      expect(result).toBe(
        '_df = mo.sql(f"""\nSELECT email FROM users LIMIT 100\n""", engine=postgres)',
      );
    });
  });
});

describe("tableUniqueId", () => {
  const ctx = (
    overrides: Partial<SQLTableContext> & { database: string; schema: string },
  ): SQLTableContext => ({
    engine: "e",
    dialect: "duckdb",
    ...overrides,
  });

  it("returns just the table name without a context", () => {
    expect(tableUniqueId(undefined, "t")).toBe("t");
  });

  it("uses database + schema for flat engines", () => {
    expect(tableUniqueId(ctx({ database: "db", schema: "public" }), "t")).toBe(
      "db.public.t",
    );
  });

  it("does not duplicate the leaf schema for nested namespaces", () => {
    // Regression: previously produced "top.nested.nested.t".
    expect(
      tableUniqueId(
        ctx({ database: "top", schema: "nested", schemaPath: ["nested"] }),
        "t",
      ),
    ).toBe("top.nested.t");
  });

  it("includes the full schema path for deeply nested namespaces", () => {
    expect(
      tableUniqueId(
        ctx({
          database: "top",
          schema: "deep",
          schemaPath: ["nested", "deep"],
        }),
        "t",
      ),
    ).toBe("top.nested.deep.t");
  });

  it("falls back to the flat schema when schemaPath is empty", () => {
    expect(
      tableUniqueId(
        ctx({ database: "db", schema: "public", schemaPath: [] }),
        "t",
      ),
    ).toBe("db.public.t");
  });

  it("does not emit a double dot for schemaless tables", () => {
    expect(
      tableUniqueId(ctx({ database: "db", schema: "", schemaPath: [] }), "t"),
    ).toBe("db.t");
    expect(tableUniqueId(ctx({ database: "db", schema: "" }), "t")).toBe(
      "db.t",
    );
  });
});

function makeTable(name: string): DataTable {
  return {
    name,
    columns: [],
    source: "memory",
    source_type: "local",
    type: "table",
    engine: null,
    indexes: null,
    num_columns: null,
    num_rows: null,
    variable_name: null,
    primary_keys: null,
  };
}

function makeSchema(opts: {
  name: string;
  tables: DataTable[];
  tables_resolved?: boolean;
  child_schemas?: DatabaseSchema[];
  child_schemas_resolved?: boolean;
}): DatabaseSchema {
  return {
    name: opts.name,
    tables: opts.tables,
    tables_resolved: opts.tables_resolved ?? true,
    child_schemas: opts.child_schemas ?? [],
    child_schemas_resolved: opts.child_schemas_resolved ?? true,
  };
}

function makeDatabase(
  name: string,
  schemas: DatabaseSchema[],
  schemas_resolved = true,
): Database {
  return { name, dialect: "duckdb", schemas, schemas_resolved, engine: null };
}

describe("schemaSubtreeMatchesSearch", () => {
  it("returns false for an empty query", () => {
    const schema = makeSchema({ name: "main", tables: [makeTable("users")] });
    expect(schemaSubtreeMatchesSearch(schema, "")).toBe(false);
    expect(schemaSubtreeMatchesSearch(schema, "   ")).toBe(false);
    expect(schemaSubtreeMatchesSearch(schema, undefined)).toBe(false);
  });

  it("matches a table name case-insensitively", () => {
    const schema = makeSchema({ name: "main", tables: [makeTable("Users")] });
    expect(schemaSubtreeMatchesSearch(schema, "user")).toBe(true);
    expect(schemaSubtreeMatchesSearch(schema, "orders")).toBe(false);
  });

  it("matches a table in a resolved child schema", () => {
    const schema = makeSchema({
      name: "parent",
      tables: [],
      child_schemas: [
        makeSchema({ name: "child", tables: [makeTable("orders")] }),
      ],
    });
    expect(schemaSubtreeMatchesSearch(schema, "orders")).toBe(true);
  });

  it("ignores deferred tables so search never triggers a fetch", () => {
    const schema = makeSchema({
      name: "main",
      tables: [makeTable("users")],
      tables_resolved: false,
    });
    expect(schemaSubtreeMatchesSearch(schema, "users")).toBe(false);
  });

  it("ignores deferred child schemas", () => {
    const schema = makeSchema({
      name: "parent",
      tables: [],
      child_schemas: [
        makeSchema({ name: "child", tables: [makeTable("orders")] }),
      ],
      child_schemas_resolved: false,
    });
    expect(schemaSubtreeMatchesSearch(schema, "orders")).toBe(false);
  });
});

describe("shouldExpandDatabaseForSearch", () => {
  it("returns false for an empty query", () => {
    const db = makeDatabase("memory", [
      makeSchema({ name: "main", tables: [makeTable("users")] }),
    ]);
    expect(shouldExpandDatabaseForSearch(db, "")).toBe(false);
    expect(shouldExpandDatabaseForSearch(db, undefined)).toBe(false);
  });

  it("expands when a loaded schema contains a matching table", () => {
    const db = makeDatabase("memory", [
      makeSchema({ name: "main", tables: [makeTable("users")] }),
      makeSchema({ name: "other", tables: [makeTable("orders")] }),
    ]);
    expect(shouldExpandDatabaseForSearch(db, "user")).toBe(true);
    expect(shouldExpandDatabaseForSearch(db, "nomatch")).toBe(false);
  });

  it("does not expand when the schema list itself is deferred", () => {
    const db = makeDatabase(
      "memory",
      [makeSchema({ name: "main", tables: [makeTable("users")] })],
      /* schemas_resolved */ false,
    );
    expect(shouldExpandDatabaseForSearch(db, "users")).toBe(false);
  });
});