import { css } from "@emotion/react";
import { Suspense, useMemo, useState } from "react";
import { graphql, useLazyLoadQuery, useMutation } from "react-relay";

import {
  Button,
  Checkbox,
  Dialog,
  DialogTrigger,
  Icon,
  Icons,
  Loading,
  Popover,
  PopoverArrow,
  View,
} from "@phoenix/components";
import type { TagPromptVersionButtonTagsQuery } from "@phoenix/components/prompt/__generated__/TagPromptVersionButtonTagsQuery.graphql";
import { DEFAULT_PROMPT_VERSION_TAGS } from "@phoenix/constants";
import { useNotifySuccess } from "@phoenix/contexts";
import { NewPromptVersionDialog } from "@phoenix/pages/prompt/NewPromptVersionTagDialog";

type TagPromptVersionButtonProps = {
  promptId: string;
  versionId: string;
};

export function TagPromptVersionButton({
  promptId,
  versionId,
}: TagPromptVersionButtonProps) {
  const [showTagList, setShowTagList] = useState<boolean>(false);
  const [showNewTagDialog, setShowNewTagDialog] = useState<boolean>(false);
  const [fetchKey, setFetchKey] = useState<number>(0);
  const notifySuccess = useNotifySuccess();

  return (
    <div>
      <DialogTrigger isOpen={showTagList} onOpenChange={setShowTagList}>
        <Button
          size="S"
          leadingVisual={<Icon svg={<Icons.PriceTags />} />}
          css={css`
            text-wrap: nowrap;
          `}
        >
          Tag Version
        </Button>
        <Popover placement="bottom end">
          <PopoverArrow />
          <Dialog>
            <Suspense
              fallback={
                <View padding="size-100">
                  <Loading size="S" />
                </View>
              }
            >
              <TagList
                promptId={promptId}
                versionId={versionId}
                fetchKey={fetchKey}
                onTagSet={(tagName) => {
                  setFetchKey((prev) => prev + 1);
                  notifySuccess({
                    title: "Tag Set",
                    message: `The tag ${tagName} has been set on the version`,
                  });
                }}
              />
            </Suspense>
            <View padding="size-100" width="250px">
              <Button
                leadingVisual={<Icon svg={<Icons.Plus />} />}
                size="S"
                css={css`
                  width: 100%;
                `}
                onPress={() => {
                  setShowTagList(false);
                  setShowNewTagDialog(true);
                }}
              >
                New Tag
              </Button>
            </View>
          </Dialog>
        </Popover>
      </DialogTrigger>
      <DialogTrigger
        isOpen={showNewTagDialog}
        onOpenChange={(isOpen) => {
          setShowNewTagDialog(isOpen);
        }}
      >
        <NewPromptVersionDialog
          promptVersionId={versionId}
          onDismiss={() => setShowNewTagDialog(false)}
          onNewTagCreated={(newTag) => {
            setFetchKey((prev) => prev + 1);
            notifySuccess({
              title: "Tag Created",
              message: `The tag ${newTag.name} has been created and set on the prompt version`,
            });
          }}
        />
      </DialogTrigger>
    </div>
  );
}

function TagList({
  promptId,
  versionId,
  fetchKey,
  onTagSet,
}: {
  promptId: string;
  versionId: string;
  fetchKey: number;
  onTagSet: (tagName: string) => void;
}) {
  const data = useLazyLoadQuery<TagPromptVersionButtonTagsQuery>(
    graphql`
      query TagPromptVersionButtonTagsQuery($promptId: ID!, $versionId: ID!) {
        prompt: node(id: $promptId) {
          ... on Prompt {
            versionTags {
              id
              name
            }
          }
        }
        promptVersion: node(id: $versionId) {
          ... on PromptVersion {
            tags {
              id
              name
            }
          }
        }
      }
    `,
    { promptId, versionId },
    { fetchKey, fetchPolicy: "store-and-network" }
  );

  const allVersionTags = useMemo(() => {
    return data?.prompt?.versionTags?.map((tag) => tag.name) || [];
  }, [data?.prompt?.versionTags]);

  const tagsSetOnVersion = useMemo(() => {
    return data?.promptVersion?.tags?.map((tag) => tag.name) || [];
  }, [data?.promptVersion?.tags]);

  const allVersionTagsWithDefaults = useMemo(() => {
    return Array.from(
      new Set([
        ...DEFAULT_PROMPT_VERSION_TAGS.map(
          (tagDefinition) => tagDefinition.name
        ),
        ...allVersionTags,
      ])
    );
  }, [allVersionTags]);

  const [commitSetTag, isCommitting] = useMutation(graphql`
    mutation TagPromptVersionButtonSetTagMutation(
      $input: SetPromptVersionTagInput!
      $promptId: ID!
    ) {
      setPromptVersionTag(input: $input) {
        query {
          prompt: node(id: $promptId) {
            ... on Prompt {
              ...PromptVersionsList__main
            }
          }
        }
      }
    }
  `);
  return (
    <ul>
      {allVersionTagsWithDefaults.map((tagName: string) => {
        const isTagSet = tagsSetOnVersion.includes(tagName);
        return (
          <li key={tagName}>
            <View
              paddingY="size-100"
              paddingX="size-200"
              borderBottomColor="default"
              borderBottomWidth="thin"
            >
              <Checkbox
                name={tagName}
                isSelected={isTagSet}
                isDisabled={isTagSet || isCommitting}
                onChange={(isSelected) => {
                  if (isSelected) {
                    const isCreate = !tagsSetOnVersion.includes(tagName);
                    let description = "";
                    if (isCreate) {
                      const tagDefinition = DEFAULT_PROMPT_VERSION_TAGS.find(
                        (tagDefinition) => tagDefinition.name === tagName
                      );
                      description = tagDefinition?.description || "";
                    }

                    commitSetTag({
                      variables: {
                        input: {
                          name: tagName,
                          promptVersionId: versionId,
                          description,
                        },
                        promptId: promptId,
                      },
                      onCompleted: () => {
                        onTagSet(tagName);
                      },
                    });
                  }
                }}
              >
                {tagName}
              </Checkbox>
            </View>
          </li>
        );
      })}
    </ul>
  );
}
