import { css } from "@emotion/react";
import { useCallback, useEffect, useMemo, useState } from "react";
import type { Key } from "react-aria-components";
import { Input, TextArea } from "react-aria-components";
import { Controller, useForm } from "react-hook-form";
import { graphql, useMutation } from "react-relay";

import {
  Alert,
  Button,
  FieldError,
  Flex,
  Label,
  TextField,
  View,
} from "@phoenix/components";
import { ComboBox, ComboBoxItem } from "@phoenix/components/core/combobox";
import { useNotifySuccess } from "@phoenix/contexts";

import type { DocumentAnnotationFormCreateMutation } from "./__generated__/DocumentAnnotationFormCreateMutation.graphql";
import type { DocumentAnnotationFormPatchMutation } from "./__generated__/DocumentAnnotationFormPatchMutation.graphql";

export type DocumentAnnotationFormData = {
  id: string;
  name: string;
  label: string | null;
  score: number | null;
  explanation: string | null;
};

type FormValues = {
  name: string;
  label: string;
  score: string;
  explanation: string;
};

const LABEL_SCORE_MAP: Record<string, number> = {
  relevant: 1,
  irrelevant: 0,
};

const PRESET_LABELS = [
  { id: "relevant", name: "relevant" },
  { id: "irrelevant", name: "irrelevant" },
];

export function DocumentAnnotationForm({
  spanNodeId,
  documentPosition,
  existingAnnotation,
  existingAnnotationNames = [],
  onSaved,
  onCancel,
}: {
  spanNodeId: string;
  documentPosition: number;
  existingAnnotation?: DocumentAnnotationFormData | null;
  existingAnnotationNames?: string[];
  onSaved?: () => void;
  onCancel: () => void;
}) {
  const [error, setError] = useState<string | null>(null);
  const notifySuccess = useNotifySuccess();

  const {
    control,
    setValue,
    trigger,
    handleSubmit,
    formState: { isValid },
  } = useForm<FormValues>({
    defaultValues: {
      name: existingAnnotation?.name ?? "relevance",
      label: existingAnnotation?.label ?? "",
      score:
        existingAnnotation?.score != null
          ? String(existingAnnotation.score)
          : "",
      explanation: existingAnnotation?.explanation ?? "",
    },
    mode: "onChange",
  });

  useEffect(() => {
    trigger("name");
  }, [trigger]);

  const takenNamesSet = useMemo(
    () => new Set(existingAnnotationNames),
    [existingAnnotationNames]
  );

  const validateNameUnique = useCallback(
    (value: string) =>
      !takenNamesSet.has(value.trim()) ||
      "An annotation with this name already exists",
    [takenNamesSet]
  );

  const [commitCreate, isCreating] =
    useMutation<DocumentAnnotationFormCreateMutation>(graphql`
      mutation DocumentAnnotationFormCreateMutation(
        $input: [CreateDocumentAnnotationInput!]!
        $spanId: ID!
      ) {
        createDocumentAnnotations(input: $input) {
          query {
            node(id: $spanId) {
              ... on Span {
                documentEvaluations {
                  id
                  annotatorKind
                  documentPosition
                  name
                  label
                  score
                  explanation
                  createdAt
                  updatedAt
                  user {
                    username
                    profilePictureUrl
                  }
                }
              }
            }
          }
        }
      }
    `);

  const [commitPatch, isPatching] =
    useMutation<DocumentAnnotationFormPatchMutation>(graphql`
      mutation DocumentAnnotationFormPatchMutation(
        $input: [PatchAnnotationInput!]!
        $spanId: ID!
      ) {
        patchDocumentAnnotations(input: $input) {
          query {
            node(id: $spanId) {
              ... on Span {
                documentEvaluations {
                  id
                  annotatorKind
                  documentPosition
                  name
                  label
                  score
                  explanation
                  createdAt
                  updatedAt
                  user {
                    username
                    profilePictureUrl
                  }
                }
              }
            }
          }
        }
      }
    `);

  const isBusy = isCreating || isPatching;

  const onSubmit = ({ name, label, score, explanation }: FormValues) => {
    const trimmedName = name.trim();
    const parsedScore = score !== "" ? Number(score) : null;
    const finalLabel = label.trim() || null;
    const finalExplanation = explanation.trim() || null;

    if (existingAnnotation) {
      commitPatch({
        variables: {
          input: [
            {
              annotationId: existingAnnotation.id,
              name: trimmedName,
              label: finalLabel,
              score: parsedScore,
              explanation: finalExplanation,
              annotatorKind: "HUMAN",
              source: "APP",
            },
          ],
          spanId: spanNodeId,
        },
        onCompleted: () => {
          notifySuccess({
            title: "Annotation updated",
            message: "Document annotation saved successfully.",
          });
          onSaved?.();
        },
        onError: (error) => {
          setError(error.message);
        },
      });
    } else {
      commitCreate({
        variables: {
          input: [
            {
              spanId: spanNodeId,
              documentPosition,
              name: trimmedName,
              annotatorKind: "HUMAN",
              label: finalLabel,
              score: parsedScore,
              explanation: finalExplanation,
              metadata: {},
              source: "APP",
            },
          ],
          spanId: spanNodeId,
        },
        onCompleted: () => {
          notifySuccess({
            title: "Annotation created",
            message: "Document annotation saved successfully.",
          });
          onSaved?.();
        },
        onError: (error) => {
          setError(error.message);
        },
      });
    }
  };

  return (
    <View padding="size-200">
      <Flex direction="column" gap="size-100">
        <Controller
          name="name"
          control={control}
          rules={{
            required: "Name is required",
            validate: { unique: validateNameUnique },
          }}
          render={({ field, fieldState: { invalid, error } }) => (
            <TextField
              size="S"
              value={field.value}
              onChange={field.onChange}
              isInvalid={invalid}
            >
              <Label>Name</Label>
              <Input placeholder="e.g. relevance" />
              {error?.message && <FieldError>{error.message}</FieldError>}
            </TextField>
          )}
        />
        <Flex direction="row" gap="size-100" alignItems="end">
          <Controller
            name="label"
            control={control}
            render={({ field }) => (
              <ComboBox
                label="Label"
                placeholder="Select or type a label"
                allowsCustomValue
                selectedKey={field.value as Key}
                inputValue={field.value}
                onInputChange={(value) => {
                  field.onChange(value);
                  const mapped = LABEL_SCORE_MAP[value];
                  if (mapped != null) {
                    setValue("score", String(mapped));
                  }
                }}
                onSelectionChange={(key) => {
                  if (key != null) {
                    const value = String(key);
                    field.onChange(value);
                    const mapped = LABEL_SCORE_MAP[value];
                    if (mapped != null) {
                      setValue("score", String(mapped));
                    }
                  }
                }}
                size="M"
                css={css`
                  flex: 1 1 0%;
                  min-width: 0;
                  .combobox__container {
                    min-width: 0;
                  }
                `}
              >
                {PRESET_LABELS.map((item) => (
                  <ComboBoxItem
                    key={item.id}
                    id={item.id}
                    textValue={item.name}
                  >
                    {item.name}
                  </ComboBoxItem>
                ))}
              </ComboBox>
            )}
          />
          <Controller
            name="score"
            control={control}
            render={({ field }) => (
              <TextField
                size="S"
                value={field.value}
                onChange={field.onChange}
                css={css`
                  flex: 1 1 0%;
                  min-width: 0;
                `}
              >
                <Label>Score</Label>
                <Input type="number" placeholder="e.g. 1" />
              </TextField>
            )}
          />
        </Flex>
        <Controller
          name="explanation"
          control={control}
          render={({ field }) => (
            <TextField
              size="S"
              value={field.value}
              onChange={field.onChange}
              css={css`
                width: 100%;
                & .react-aria-TextArea {
                  resize: vertical;
                  transition: none;
                }
              `}
            >
              <Label>Explanation</Label>
              <TextArea rows={2} placeholder="Optional explanation" />
            </TextField>
          )}
        />
        {error && <Alert variant="danger">{error}</Alert>}
        <Flex direction="row" gap="size-100" justifyContent="end">
          <Button size="S" onPress={onCancel}>
            Cancel
          </Button>
          <Button
            variant="primary"
            size="S"
            onPress={() => handleSubmit(onSubmit)()}
            isDisabled={isBusy || !isValid}
          >
            Save
          </Button>
        </Flex>
      </Flex>
    </View>
  );
}
