/* Copyright 2026 Marimo. All rights reserved. */
/* oxlint-disable typescript/no-explicit-any */

import { PlusIcon, RefreshCcw, Trash2Icon } from "lucide-react";
import React from "react";
import {
  type FieldValues,
  FormProvider,
  type Path,
  type UseFormReturn,
  useFieldArray,
} from "react-hook-form";
import { z } from "zod";
import { FieldOptions, randomNumber } from "@/components/forms/options";
import { Combobox, ComboboxItem } from "@/components/ui/combobox";
import { Label } from "@/components/ui/label";
import { NativeSelect } from "@/components/ui/native-select";
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
import {
  Select,
  SelectContent,
  SelectGroup,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { cn } from "@/utils/cn";
import { Events } from "@/utils/events";
import { Strings } from "@/utils/strings";
import { isZodPipe } from "@/utils/zod-utils";
import { Objects } from "../../utils/objects";
import { Button } from "../ui/button";
import { Checkbox } from "../ui/checkbox";
import {
  FormControl,
  FormDescription,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
} from "../ui/form";
import { DebouncedInput, DebouncedNumberInput } from "../ui/input";
import { Textarea } from "../ui/textarea";
import { getDefaults, getUnionLiteral, maybeUnwrap } from "./form-utils";
import {
  ensureStringArray,
  SwitchableMultiSelect,
  TextAreaMultiSelect,
} from "./switchable-multi-select";
export interface FormRenderer<T extends FieldValues = any, S = any> {
  isMatch: (schema: z.ZodType) => schema is z.ZodType<S>;
  Component: React.ComponentType<{
    schema: z.ZodType<S>;
    form: UseFormReturn<T>;
    path: Path<T>;
  }>;
}

interface Props<T extends FieldValues> {
  form: UseFormReturn<T>;
  schema: z.ZodType;
  path?: Path<T>;
  renderers: FormRenderer<T>[] | undefined;
  children?: React.ReactNode;
}

export const ZodForm = <T extends FieldValues>({
  schema,
  form,
  path = "" as Path<T>,
  renderers = [],
  children,
}: Props<T>) => {
  return (
    <FormProvider {...form}>
      {children}
      {renderZodSchema(schema, form, path, renderers)}
    </FormProvider>
  );
};

export function renderZodSchema<T extends FieldValues, S>(
  schema: z.ZodType<S>,
  form: UseFormReturn<T>,
  path: Path<T>,
  renderers: FormRenderer<T>[],
) {
  // Try custom renderers first
  for (const renderer of renderers) {
    const { isMatch, Component } = renderer;
    if (isMatch(schema)) {
      return <Component schema={schema} form={form} path={path} />;
    }
  }

  const {
    label,
    description,
    special,
    direction = "column",
    minLength,
  } = FieldOptions.parse(schema.description || "");

  if (schema instanceof z.ZodDefault) {
    let inner = schema.unwrap() as z.ZodType;
    inner =
      !inner.description && schema.description
        ? inner.describe(schema.description)
        : inner;
    return renderZodSchema(inner, form, path, renderers);
  }

  if (schema instanceof z.ZodOptional) {
    let inner = schema.unwrap() as z.ZodType;
    inner =
      !inner.description && schema.description
        ? inner.describe(schema.description)
        : inner;
    return renderZodSchema(inner, form, path, renderers);
  }

  if (schema instanceof z.ZodObject) {
    return (
      <div
        className={cn(
          "flex",
          direction === "row"
            ? "flex-row gap-6 items-start"
            : direction === "two-columns"
              ? "grid grid-cols-2 gap-y-6"
              : "flex-col gap-6",
        )}
      >
        <FormLabel>{label}</FormLabel>
        {Objects.entries(schema.shape).map(([key, value]) => {
          const isLiteral = value instanceof z.ZodLiteral;
          const childForm = renderZodSchema(
            value as z.ZodType,
            form,
            joinPath(path, key),
            renderers,
          );

          if (isLiteral) {
            return <React.Fragment key={key}>{childForm}</React.Fragment>;
          }

          return (
            <div className="flex flex-row align-start" key={key}>
              {childForm}
            </div>
          );
        })}
      </div>
    );
  }

  if (schema instanceof z.ZodString) {
    if (special === "time") {
      return (
        <FormField
          control={form.control}
          name={path}
          render={({ field }) => (
            <FormItem>
              <FormLabel>{label}</FormLabel>
              <FormDescription>{description}</FormDescription>
              <FormControl>
                <DebouncedInput
                  {...field}
                  onValueChange={field.onChange}
                  type="time"
                  className="my-0"
                />
              </FormControl>
              <FormMessage />
            </FormItem>
          )}
        />
      );
    }

    return <StringFormField schema={schema} form={form} path={path} />;
  }
  if (schema instanceof z.ZodBoolean) {
    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => (
          <FormItem>
            <div className="flex flex-row items-start space-x-2">
              <FormLabel>{label}</FormLabel>
              <FormDescription>{description}</FormDescription>
              <FormControl>
                <Checkbox
                  data-testid="marimo-plugin-data-frames-boolean-checkbox"
                  checked={field.value}
                  onCheckedChange={field.onChange}
                />
              </FormControl>
            </div>
            <FormMessage />
          </FormItem>
        )}
      />
    );
  }
  if (schema instanceof z.ZodNumber) {
    if (special === "random_number_button") {
      return (
        <FormField
          control={form.control}
          name={path}
          render={({ field }) => (
            <Button
              size="xs"
              data-testid="marimo-plugin-data-frames-random-number-button"
              variant="secondary"
              onClick={() => {
                field.onChange(randomNumber());
              }}
            >
              <RefreshCcw className="w-3.5 h-3.5 mr-1" />
              {label}
            </Button>
          )}
        />
      );
    }

    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => (
          <FormItem>
            <FormLabel>{label}</FormLabel>
            <FormDescription>{description}</FormDescription>
            <FormControl>
              <DebouncedNumberInput
                {...field}
                className="my-0"
                onValueChange={field.onChange}
              />
            </FormControl>
            <FormMessage />
          </FormItem>
        )}
      />
    );
  }
  if (schema instanceof z.ZodDate) {
    const inputType =
      special === "datetime"
        ? "datetime-local"
        : special === "time"
          ? "time"
          : "date";

    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => (
          <FormItem>
            <FormLabel>{label}</FormLabel>
            <FormDescription>{description}</FormDescription>
            <FormControl>
              <DebouncedInput
                {...field}
                onValueChange={field.onChange}
                type={inputType}
                className="my-0"
              />
            </FormControl>
            <FormMessage />
          </FormItem>
        )}
      />
    );
  }
  if (schema instanceof z.ZodAny) {
    return <StringFormField schema={schema} form={form} path={path} />;
  }
  if (schema instanceof z.ZodEnum) {
    const values = schema.options.map((option) => option.toString());
    return (
      <SelectFormField
        schema={schema}
        form={form}
        path={path}
        options={values}
      />
    );
  }
  if (schema instanceof z.ZodArray) {
    if (special === "text_area_multiline") {
      return <MultiStringFormField schema={schema} form={form} path={path} />;
    }

    // Inspect child type for a better input
    const childType = schema.element;

    // Show multi-select for enum array
    if (childType instanceof z.ZodEnum) {
      const childOptions: string[] = childType.options.map((option) =>
        option.toString(),
      );
      return (
        <MultiSelectFormField
          schema={schema}
          form={form}
          path={path}
          itemLabel={label}
          options={childOptions}
        />
      );
    }

    // Fallback to generic array
    return (
      <div className="flex flex-col gap-1">
        <Label>{label}</Label>
        <FormArray
          schema={schema.element as z.ZodType}
          form={form}
          path={path}
          key={path}
          minLength={minLength}
          renderers={renderers}
        />
      </div>
    );
  }

  if (schema instanceof z.ZodDiscriminatedUnion) {
    const def = schema.def;
    const options = def.options as z.ZodType[];
    const discriminator = def.discriminator;
    const getSchemaValue = (value: string) => {
      return options.find((option) => {
        return getUnionLiteral(option).value === value;
      });
    };

    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => {
          const value = field.value;
          const types = options.map((option) => {
            return getUnionLiteral(option).value;
          });

          const unionTypeValue: string =
            value && typeof value === "object" && discriminator in value
              ? value[discriminator]
              : types[0];

          const selectedOption = getSchemaValue(unionTypeValue);

          return (
            <div className="flex flex-col">
              <FormLabel>{label}</FormLabel>
              <div className="flex border-b mb-4 -mt-2">
                {types.map((type: string) => (
                  <button
                    key={type}
                    type="button"
                    className={`px-4 py-2 ${
                      unionTypeValue === type
                        ? "border-b-2 border-primary font-medium"
                        : "text-muted-foreground"
                    }`}
                    onClick={() => {
                      const nextSchema = getSchemaValue(type);
                      if (nextSchema) {
                        field.onChange(getDefaults(nextSchema));
                      } else {
                        field.onChange({ [discriminator]: type });
                      }
                    }}
                  >
                    {type}
                  </button>
                ))}
              </div>
              <div className="flex flex-col" key={unionTypeValue}>
                {selectedOption &&
                  renderZodSchema(selectedOption, form, path, renderers)}
              </div>
            </div>
          );
        }}
      />
    );
  }
  if (schema instanceof z.ZodUnion) {
    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => {
          const options = schema.options as z.ZodType[];
          let value: string = field.value;
          const types = options.map((option) => {
            return getUnionLiteral(option).value;
          });

          if (!value) {
            // Default to first
            value = types[0];
          }

          const selectedOption = options.find((option) => {
            return getUnionLiteral(option).value === value;
          });

          return (
            <div className="flex flex-col mb-4 gap-1">
              <FormLabel>{label}</FormLabel>
              <NativeSelect
                data-testid="marimo-plugin-data-frames-union-select"
                {...field}
              >
                {types.map((type: string) => {
                  return (
                    <option key={type} value={type}>
                      {type}
                    </option>
                  );
                })}
              </NativeSelect>
              {selectedOption &&
                renderZodSchema(selectedOption, form, path, renderers)}
            </div>
          );
        }}
      />
    );
  }

  if (schema instanceof z.ZodLiteral) {
    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => (
          <input
            {...field}
            type="hidden"
            value={String([...schema.values][0] ?? "")}
          />
        )}
      />
    );
  }
  if ("unwrap" in schema) {
    // Handle ZodEffects (transforms/refinements)
    return renderZodSchema(maybeUnwrap(schema), form, path, renderers);
  }
  if (isZodPipe(schema)) {
    return renderZodSchema(schema.in, form, path, renderers);
  }

  return (
    <div>
      Unknown schema type{" "}
      {schema == null ? path : JSON.stringify(schema.type ?? schema)}
    </div>
  );
}

/**
 * Type: T[]
 */
const FormArray = ({
  schema,
  form,
  path,
  minLength,
  renderers,
}: {
  schema: z.ZodType;
  form: UseFormReturn<any>;
  path: Path<any>;
  renderers: FormRenderer[];
  minLength?: number;
}) => {
  const { label, description } = FieldOptions.parse(schema.description || "");

  const control = form.control;
  // prepend, remove, swap, move, insert, replace
  const { fields, append, remove } = useFieldArray({
    control,
    name: path,
  });

  const isBelowMinLength = minLength != null && fields.length < minLength;
  const canRemove = minLength == null || fields.length > minLength;

  return (
    <div className="flex flex-col gap-2 min-w-[220px]">
      <FormLabel>{label}</FormLabel>
      <FormDescription>{description}</FormDescription>
      {fields.map((field, index) => {
        return (
          <div
            className="flex flex-row pl-2 ml-2 border-l-2 border-disabled hover-actions-parent relative pr-5 pt-1 items-center w-fit"
            key={field.id}
            onKeyDown={Events.onEnter((e) => e.preventDefault())}
          >
            {renderZodSchema(schema, form, `${path}[${index}]`, renderers)}
            {canRemove && (
              <Trash2Icon
                className="w-4 h-4 ml-2 my-1 text-muted-foreground hover:text-destructive cursor-pointer absolute right-0 top-5"
                onClick={() => {
                  remove(index);
                }}
              />
            )}
          </div>
        );
      })}
      {isBelowMinLength && (
        <div className="text-destructive text-xs font-semibold">
          <div>At least {minLength} required.</div>
        </div>
      )}
      <div>
        <Button
          size="xs"
          data-testid="marimo-plugin-data-frames-add-array-item"
          variant="text"
          className="hover:text-accent-foreground"
          onClick={() => {
            append(getDefaults(schema));
          }}
        >
          <PlusIcon className="w-3.5 h-3.5 mr-1" />
          Add
        </Button>
      </div>
    </div>
  );
};

/**
 * Type: string
 */
const StringFormField = ({
  schema,
  form,
  path,
}: {
  schema: z.ZodAny | z.ZodString;
  form: UseFormReturn<any>;
  path: Path<any>;
}) => {
  const { label, description, placeholder, disabled, inputType } =
    FieldOptions.parse(schema.description);

  if (inputType === "textarea") {
    return (
      <FormField
        control={form.control}
        name={path}
        render={({ field }) => (
          <FormItem>
            <FormLabel>{label}</FormLabel>
            <FormDescription>{description}</FormDescription>
            <FormControl>
              <Textarea
                {...field}
                value={field.value}
                onChange={field.onChange}
                className="my-0"
                placeholder={placeholder}
                disabled={disabled}
              />
            </FormControl>
            <FormMessage />
          </FormItem>
        )}
      />
    );
  }

  return (
    <FormField
      control={form.control}
      name={path}
      render={({ field }) => (
        <FormItem>
          <FormLabel>{label}</FormLabel>
          <FormDescription>{description}</FormDescription>
          <FormControl>
            <DebouncedInput
              {...field}
              type={inputType}
              value={field.value}
              onValueChange={field.onChange}
              className="my-0"
              placeholder={placeholder}
              disabled={disabled}
            />
          </FormControl>
          <FormMessage />
        </FormItem>
      )}
    />
  );
};

/**
 * Type: string[]
 */
const MultiStringFormField = ({
  schema,
  form,
  path,
}: {
  schema: z.ZodAny | z.ZodString | z.ZodArray<any>;
  form: UseFormReturn<any>;
  path: Path<any>;
}) => {
  const { label, description, placeholder } = FieldOptions.parse(
    schema.description,
  );

  return (
    <FormField
      control={form.control}
      name={path}
      render={({ field }) => (
        <FormItem>
          <FormLabel>{label}</FormLabel>
          <FormDescription>{description}</FormDescription>
          <FormControl>
            <TextAreaMultiSelect {...field} placeholder={placeholder} />
          </FormControl>
          <FormMessage />
        </FormItem>
      )}
    />
  );
};

/**
 * Type: string
 * Known options
 */
const SelectFormField = ({
  schema,
  form,
  path,
  options,
  textTransform,
}: {
  schema: z.ZodEnum<any> | z.ZodString;
  form: UseFormReturn<any>;
  path: Path<any>;
  options: string[];
  textTransform?: (value: string) => string;
}) => {
  const { label, description, disabled, special } = FieldOptions.parse(
    schema.description,
  );

  if (special === "radio_group") {
    return (
      <FormField
        control={form.control}
        name={path}
        disabled={disabled}
        render={({ field }) => (
          <FormItem>
            <FormLabel>{label}</FormLabel>
            <FormDescription>{description}</FormDescription>
            <FormControl>
              <RadioGroup
                className="flex flex-row gap-2 pt-1 items-center"
                value={field.value}
                onValueChange={field.onChange}
              >
                {options.map((value: string) => {
                  return (
                    <div className="flex items-center gap-1" key={value}>
                      <RadioGroupItem
                        key={value}
                        value={value}
                        id={`${path}-${value}`}
                      />
                      <FormLabel
                        className="whitespace-pre"
                        htmlFor={`${path}-${value}`}
                      >
                        {Strings.startCase(value)}
                      </FormLabel>
                    </div>
                  );
                })}
              </RadioGroup>
            </FormControl>
            <FormMessage />
          </FormItem>
        )}
      />
    );
  }

  return (
    <FormField
      control={form.control}
      name={path}
      disabled={disabled}
      render={({ field }) => (
        <FormItem>
          <FormLabel className="whitespace-pre">{label}</FormLabel>
          <FormDescription>{description}</FormDescription>
          <FormControl>
            <Select
              data-testid="marimo-plugin-data-frames-select"
              value={field.value}
              onValueChange={field.onChange}
            >
              <SelectTrigger className="min-w-[180px]">
                <SelectValue placeholder="--" />
              </SelectTrigger>
              <SelectContent>
                <SelectGroup>
                  {options.map((value: string) => {
                    return (
                      <SelectItem key={value} value={value}>
                        {textTransform?.(value) ?? value}
                      </SelectItem>
                    );
                  })}
                  {options.length === 0 && (
                    <SelectItem disabled={true} value="--">
                      No options
                    </SelectItem>
                  )}
                </SelectGroup>
              </SelectContent>
            </Select>
          </FormControl>
          <FormMessage />
        </FormItem>
      )}
    />
  );
};

/**
 * Type: string[]
 * Known options
 */
const MultiSelectFormField = ({
  schema,
  form,
  path,
  options,
  itemLabel,
  showSwitchable,
}: {
  schema: z.ZodEnum<any> | z.ZodString | z.ZodArray<any>;
  form: UseFormReturn<any>;
  path: Path<any>;
  itemLabel?: string;
  options: string[];
  showSwitchable?: boolean;
}) => {
  const { label, description, placeholder } = FieldOptions.parse(
    schema.description,
  );

  const resolvePlaceholder =
    placeholder ??
    (itemLabel ? `Select ${itemLabel?.toLowerCase()}` : undefined);

  return (
    <FormField
      control={form.control}
      name={path}
      render={({ field }) => {
        const valueAsArray = ensureStringArray(field.value);
        return (
          <FormItem>
            <FormLabel className="whitespace-pre">{label}</FormLabel>
            <FormDescription>{description}</FormDescription>
            <FormControl>
              {showSwitchable ? (
                <SwitchableMultiSelect
                  {...field}
                  value={valueAsArray}
                  options={options}
                  placeholder={resolvePlaceholder}
                />
              ) : (
                <Combobox
                  className="min-w-[180px]"
                  placeholder={resolvePlaceholder}
                  displayValue={(option: string) => Strings.startCase(option)}
                  multiple={true}
                  chips={true}
                  keepPopoverOpenOnSelect={true}
                  value={valueAsArray}
                  onValueChange={field.onChange}
                >
                  {options.map((option) => (
                    <ComboboxItem key={option} value={option}>
                      {option}
                    </ComboboxItem>
                  ))}
                </Combobox>
              )}
            </FormControl>
            <FormMessage />
          </FormItem>
        );
      }}
    />
  );
};

function joinPath<T>(...parts: (string | number)[]): Path<T> {
  return parts.filter((part) => part !== "").join(".") as Path<T>;
}