/* Copyright 2026 Marimo. All rights reserved. */
import { z } from "zod";
import { FieldOptions } from "@/components/forms/options";

function passwordField() {
  return z
    .string()
    .optional()
    .describe(
      FieldOptions.of({
        label: "Password",
        inputType: "password",
        placeholder: "password",
        optionRegex: ".*password.*",
      }),
    );
}

function tokenField(label?: string, required?: boolean) {
  let field: z.ZodString | z.ZodOptional<z.ZodString> = z.string();
  field = required ? field.nonempty() : field.optional();

  field = field.describe(
    FieldOptions.of({
      label: label || "Token",
      inputType: "password",
      placeholder: "token",
      optionRegex: ".*token.*",
    }),
  );
  return field;
}

function warehouseNameField() {
  return z
    .string()
    .optional()
    .describe(
      FieldOptions.of({
        label: "Warehouse Name",
        placeholder: "warehouse",
        optionRegex: ".*warehouse.*",
      }),
    );
}

function uriField(label?: string, required?: boolean) {
  let field: z.ZodString | z.ZodOptional<z.ZodString> = z.string();
  field = required ? field.nonempty() : field.optional();

  return field.describe(
    FieldOptions.of({ label: label || "URI", optionRegex: ".*uri.*" }),
  );
}

function hostField(label?: string) {
  return z
    .string()
    .nonempty()
    .describe(
      FieldOptions.of({
        label: label || "Host",
        placeholder: "localhost",
        optionRegex: ".*host.*",
      }),
    );
}

function databaseField() {
  return z.string().describe(
    FieldOptions.of({
      label: "Database",
      placeholder: "db name",
      optionRegex: ".*database.*",
    }),
  );
}

function schemaField() {
  return z.string().describe(
    FieldOptions.of({
      label: "Schema",
      placeholder: "schema name",
      optionRegex: ".*schema.*",
    }),
  );
}

function usernameField() {
  return z
    .string()
    .nonempty()
    .describe(
      FieldOptions.of({
        label: "Username",
        placeholder: "username",
        optionRegex: ".*username.*",
      }),
    );
}

function portField(defaultPort?: number) {
  const field = z.coerce
    .string()
    .describe(
      FieldOptions.of({
        label: "Port",
        inputType: "number",
        placeholder: defaultPort?.toString(),
      }),
    )
    .transform(Number)
    .refine((n: number) => n >= 0 && n <= 65_535, {
      message: "Port must be between 0 and 65535",
    });

  if (defaultPort !== undefined) {
    return field.default(defaultPort);
  }

  return field;
}

function readOnlyField() {
  return z
    .boolean()
    .default(false)
    .describe(FieldOptions.of({ label: "Read Only" }));
}

export const PostgresConnectionSchema = z
  .object({
    type: z.literal("postgres"),
    host: hostField(),
    port: portField(5432).optional(),
    database: databaseField(),
    username: usernameField(),
    password: passwordField(),
    ssl: z
      .boolean()
      .default(false)
      .describe(FieldOptions.of({ label: "Use SSL" })),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const MySQLConnectionSchema = z
  .object({
    type: z.literal("mysql"),
    host: hostField(),
    port: portField(3306),
    database: databaseField(),
    username: usernameField(),
    password: passwordField(),
    ssl: z
      .boolean()
      .default(false)
      .describe(FieldOptions.of({ label: "Use SSL" })),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const SQLiteConnectionSchema = z
  .object({
    type: z.literal("sqlite"),
    database: databaseField().describe(
      FieldOptions.of({ label: "Database Path" }),
    ),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const DuckDBConnectionSchema = z
  .object({
    type: z.literal("duckdb"),
    database: databaseField().describe(
      FieldOptions.of({ label: "Database Path" }),
    ),
    read_only: readOnlyField(),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const MotherDuckConnectionSchema = z
  .object({
    type: z.literal("motherduck"),
    database: databaseField()
      .default("my_db")
      .describe(FieldOptions.of({ label: "Database Name" })),
    token: tokenField(),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const SnowflakeConnectionSchema = z
  .object({
    type: z.literal("snowflake"),
    account: z
      .string()
      .nonempty()
      .describe(
        FieldOptions.of({
          label: "Account",
          optionRegex: ".*snowflake.*",
        }),
      ),
    warehouse: z
      .string()
      .optional()
      .describe(
        FieldOptions.of({
          label: "Warehouse",
          optionRegex: ".*snowflake.*",
        }),
      ),
    database: databaseField(),
    schema: z
      .string()
      .optional()
      .describe(
        FieldOptions.of({
          label: "Schema",
          optionRegex: ".*snowflake.*",
        }),
      ),
    role: z
      .string()
      .optional()
      .describe(FieldOptions.of({ label: "Role" })),
    authType: z
      .discriminatedUnion("type", [
        z.object({
          type: z.literal("Password"),
          username: usernameField(),
          password: passwordField(),
          enable_mfa: z
            .boolean()
            .default(false)
            .describe(FieldOptions.of({ label: "Enable MFA (Duo Push)" })),
        }),
        z.object({
          type: z.literal("SSO (Browser)"),
          username: usernameField(),
        }),
        z.object({
          type: z.literal("Key Pair"),
          username: usernameField(),
          private_key_path: z
            .string()
            .nonempty()
            .describe(
              FieldOptions.of({
                label: "Private Key Path",
                placeholder: "/path/to/rsa_key.p8",
              }),
            ),
          private_key_passphrase: z
            .string()
            .optional()
            .describe(
              FieldOptions.of({
                label: "Private Key Passphrase",
                inputType: "password",
                optionRegex: ".*passphrase.*",
              }),
            ),
        }),
        z.object({
          type: z.literal("OAuth / PAT"),
          token: tokenField("Token", true),
        }),
      ])
      .default({
        type: "Password",
        username: "username",
        enable_mfa: false,
      })
      .describe(FieldOptions.of({ special: "tabs" })),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const BigQueryConnectionSchema = z
  .object({
    type: z.literal("bigquery"),
    project: z
      .string()
      .nonempty()
      .describe(
        FieldOptions.of({
          label: "Project ID",
          optionRegex: ".*bigquery.*",
        }),
      ),
    dataset: z
      .string()
      .nonempty()
      .describe(
        FieldOptions.of({
          label: "Dataset",
          optionRegex: ".*bigquery.*",
        }),
      ),
    credentials_json: z
      .string()
      .describe(
        FieldOptions.of({ label: "Credentials JSON", inputType: "textarea" }),
      ),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const ClickhouseConnectionSchema = z
  .object({
    type: z.literal("clickhouse_connect"),
    host: hostField(),
    port: portField(8123).optional(),
    username: usernameField(),
    password: passwordField(),
    secure: z
      .boolean()
      .default(false)
      .describe(FieldOptions.of({ label: "Use HTTPs" })),
    proxy_path: z
      .string()
      .optional()
      .describe(
        FieldOptions.of({
          label: "Proxy Path",
          placeholder: "/clickhouse",
        }),
      ),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const TimeplusConnectionSchema = z
  .object({
    type: z.literal("timeplus"),
    host: hostField().default("localhost"),
    port: portField(8123).optional(),
    username: usernameField().default("default"),
    password: passwordField().default(""),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const ChdbConnectionSchema = z
  .object({
    type: z.literal("chdb"),
    database: databaseField().describe(
      FieldOptions.of({ label: "Database Path" }),
    ),
    read_only: readOnlyField(),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const TrinoConnectionSchema = z
  .object({
    type: z.literal("trino"),
    host: hostField(),
    port: portField(8080),
    database: databaseField(),
    schema: schemaField().optional(),
    username: usernameField(),
    password: passwordField(),
    async_support: z
      .boolean()
      .default(false)
      .describe(FieldOptions.of({ label: "Async Support" })),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const IcebergConnectionSchema = z.object({
  type: z.literal("iceberg"),
  name: z.string().describe(FieldOptions.of({ label: "Catalog Name" })),
  catalog: z
    .discriminatedUnion("type", [
      z.object({
        type: z.literal("REST"),
        warehouse: warehouseNameField(),
        uri: z
          .string()
          .optional()
          .describe(
            FieldOptions.of({
              label: "URI",
              placeholder: "https://",
              optionRegex: ".*uri.*",
            }),
          ),
        token: tokenField(),
      }),
      z.object({
        type: z.literal("SQL"),
        warehouse: warehouseNameField(),
        uri: z
          .string()
          .optional()
          .describe(
            FieldOptions.of({
              label: "URI",
              placeholder: "jdbc:iceberg://host:port/database",
              optionRegex: ".*uri.*",
            }),
          ),
      }),
      z.object({
        type: z.literal("Hive"),
        warehouse: warehouseNameField(),
        uri: uriField(),
      }),
      z.object({
        type: z.literal("Glue"),
        warehouse: warehouseNameField(),
        uri: uriField(),
      }),
      z.object({
        type: z.literal("DynamoDB"),
        "dynamodb.profile-name": z
          .string()
          .optional()
          .describe(FieldOptions.of({ label: "Profile Name" })),
        "dynamodb.region": z
          .string()
          .optional()
          .describe(FieldOptions.of({ label: "Region" })),
        "dynamodb.access-key-id": z
          .string()
          .optional()
          .describe(FieldOptions.of({ label: "Access Key ID" })),
        "dynamodb.secret-access-key": z
          .string()
          .optional()
          .describe(
            FieldOptions.of({
              label: "Secret Access Key",
              inputType: "password",
            }),
          ),
        "dynamodb.session-token": z
          .string()
          .optional()
          .describe(
            FieldOptions.of({
              label: "Session Token",
              inputType: "password",
            }),
          ),
      }),
    ])
    .default({
      type: "REST",
      token: undefined,
    })
    .describe(FieldOptions.of({ special: "tabs" })),
});

export const DataFusionConnectionSchema = z.object({
  type: z.literal("datafusion"),
  sessionContext: z
    .boolean()
    .optional()
    .describe(
      FieldOptions.of({
        label: "Use Session Context",
      }),
    ),
});

// Ideally, we can conditionally render the username, host, and port fields.
export const PySparkConnectionSchema = z.object({
  type: z.literal("pyspark"),
  host: hostField().optional(),
  port: portField().optional(),
});

// Ref: https://github.com/aws/amazon-redshift-python-driver/blob/master/tutorials/001%20-%20Connecting%20to%20Amazon%20Redshift.ipynb
export const RedshiftConnectionSchema = z
  .object({
    type: z.literal("redshift"),
    host: hostField(),
    port: portField(5439),
    connectionType: z
      .discriminatedUnion("type", [
        z.object({
          type: z.literal("IAM credentials"),
          region: z.string().describe(FieldOptions.of({ label: "Region" })),
          aws_access_key_id: z
            .string()
            .nonempty()
            .describe(
              FieldOptions.of({
                label: "AWS Access Key ID",
                inputType: "password",
                optionRegex: ".*aws_access_key_id.*",
              }),
            ),
          aws_secret_access_key: z
            .string()
            .nonempty()
            .describe(
              FieldOptions.of({
                label: "AWS Secret Access Key",
                inputType: "password",
                optionRegex: ".*aws_secret_access_key.*",
              }),
            ),
          aws_session_token: z
            .string()
            .optional()
            .describe(
              FieldOptions.of({
                label: "AWS Session Token",
                inputType: "password",
                optionRegex: ".*aws_session_token.*",
              }),
            ),
        }),
        z.object({
          type: z.literal("DB credentials"),
          user: usernameField(),
          password: passwordField(),
        }),
      ])
      .default({
        type: "IAM credentials",
        aws_access_key_id: "",
        aws_secret_access_key: "",
        region: "",
      }),
    database: databaseField(),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const DatabricksConnectionSchema = z
  .object({
    type: z.literal("databricks"),
    access_token: tokenField("Access Token", true),
    server_hostname: hostField("Server Hostname"),
    http_path: uriField("HTTP Path", true),
    catalog: z
      .string()
      .optional()
      .describe(FieldOptions.of({ label: "Catalog" })),
    schema: z
      .string()
      .optional()
      .describe(FieldOptions.of({ label: "Schema" })),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const SupabaseConnectionSchema = z
  .object({
    type: z.literal("supabase"),
    host: hostField(),
    port: portField(5432).optional(),
    database: databaseField(),
    username: usernameField(),
    password: passwordField(),
    disable_client_pooling: z
      .boolean()
      .default(false)
      .describe(
        FieldOptions.of({
          label: "Disable Client-Side Pooling",
        }),
      ),
  })
  .describe(FieldOptions.of({ direction: "two-columns" }));

export const DatabaseConnectionSchema = z.discriminatedUnion("type", [
  PostgresConnectionSchema,
  MySQLConnectionSchema,
  SQLiteConnectionSchema,
  DuckDBConnectionSchema,
  MotherDuckConnectionSchema,
  SnowflakeConnectionSchema,
  BigQueryConnectionSchema,
  ClickhouseConnectionSchema,
  TimeplusConnectionSchema,
  ChdbConnectionSchema,
  TrinoConnectionSchema,
  IcebergConnectionSchema,
  DataFusionConnectionSchema,
  PySparkConnectionSchema,
  RedshiftConnectionSchema,
  DatabricksConnectionSchema,
  SupabaseConnectionSchema,
]);

export type DatabaseConnection = z.infer<typeof DatabaseConnectionSchema>;