import {
  CreateFineTuneInputModelEnum,
  FineTuneCohereInputModelEnum,
  FineTuneOpenaiInputModelEnum,
} from "~/browser-client";
import { useFetcher, useNavigate, useParams } from "@remix-run/react";
import { SlidersHorizontal } from "lucide-react";
import { FC, useEffect, useState } from "react";
import { LoadingSpinner } from "~/components/LoadingIcon";
import { Button } from "~/components/ui/button";
import { Dialog, DialogContent, DialogFooter, DialogHeader, DialogTitle, DialogTrigger } from "~/components/ui/dialog";
import { Label } from "~/components/ui/label";
import {
  Select,
  SelectContent,
  SelectGroup,
  SelectItem,
  SelectLabel,
  SelectTrigger,
  SelectValue,
} from "~/components/ui/select";
import { usePaidGuard } from "~/hooks";
import { useAppLayoutData } from "~/hooks/useAppLayoutData";
import { Routes } from "~/routes";
import {
  StartCohereFineTuneBody,
  StartFineTuneAction,
  StartMonteloFineTuneBody,
  StartOpenAIFineTuneBody,
} from "~/types/StartFineTuneAction.types";

type Models = (typeof FineTuneOpenaiInputModelEnum)[keyof typeof FineTuneOpenaiInputModelEnum] &
  (typeof CreateFineTuneInputModelEnum)[keyof typeof CreateFineTuneInputModelEnum] &
  (typeof FineTuneCohereInputModelEnum)[keyof typeof FineTuneCohereInputModelEnum];

export const FineTuneDialog: FC<{ isDialogDisabled: boolean }> = ({ isDialogDisabled }) => {
  const appLayoutData = useAppLayoutData();
  const { isBasicPlan } = usePaidGuard();
  const params = useParams();
  const fetcher = useFetcher<StartFineTuneAction>();
  const navigate = useNavigate();
  const [selectedModel, setSelectedModel] = useState<Models | null>(null);
  const [isDialogOpen, setIsDialogOpen] = useState(false);

  const projectOpenAIKey = appLayoutData?.project.openaiApiKey;
  const projectCohereKey = appLayoutData?.project.cohereApiKey;

  const noOpenAIApiKey = !projectOpenAIKey.length;
  const noCohereApiKey = !projectCohereKey.length;
  const isNotIdle = fetcher.state !== "idle";
  const isStartDisabled = isBasicPlan || isNotIdle || !selectedModel;

  const handleStartFineTune = () => {
    if (!selectedModel || isNotIdle) {
      return;
    }

    const isOpenAI = Object.values(FineTuneOpenaiInputModelEnum).includes(selectedModel);
    const isCohere = Object.values(FineTuneCohereInputModelEnum).includes(selectedModel);

    if (isOpenAI) {
      fetcher.submit(
        {
          model: selectedModel,
        } as StartOpenAIFineTuneBody,
        {
          method: "POST",
          action: Routes.actions.finetuning.createOpenAI({
            projectId: params.projectId!,
            datasetId: params.datasetId!,
          }),
          encType: "application/json",
        },
      );
    } else if (isCohere) {
      fetcher.submit(
        {
          model: selectedModel,
        } as StartCohereFineTuneBody,
        {
          method: "POST",
          action: Routes.actions.finetuning.createCohere({
            projectId: params.projectId!,
            datasetId: params.datasetId!,
          }),
          encType: "application/json",
        },
      );
    } else {
      fetcher.submit(
        {
          model: selectedModel,
        } as StartMonteloFineTuneBody,
        {
          method: "POST",
          action: Routes.actions.finetuning.createMontelo({
            projectId: params.projectId!,
            datasetId: params.datasetId!,
          }),
          encType: "application/json",
        },
      );
    }
  };

  const handleSelect = (value: Models) => {
    setSelectedModel(value);
  };

  useEffect(() => {
    if (fetcher.data?.finetune.id) {
      setIsDialogOpen(false);
      void navigate(
        Routes.app.project.env.finetuning({
          projectId: params.projectId!,
          envId: params.envId!,
        }),
      );
    }
  }, [fetcher.data]);

  return (
    <Dialog open={isDialogOpen} onOpenChange={setIsDialogOpen}>
      <DialogTrigger asChild>
        <Button disabled={isDialogDisabled}>
          <SlidersHorizontal className="mr-2 h-4 w-4" /> Fine-Tune
        </Button>
      </DialogTrigger>
      <DialogContent className={"top-[30%] max-w-[420px] overflow-y-scroll"}>
        <DialogHeader>
          <DialogTitle>Create a Fine-Tuned Model</DialogTitle>
        </DialogHeader>
        <div className={"flex flex-col gap-8"}>
          <div className={"flex flex-col gap-2"}>
            <Label>Base Model</Label>
            <Select onValueChange={handleSelect} disabled={isBasicPlan}>
              <SelectTrigger className="w-full">
                <SelectValue placeholder="Select..." />
              </SelectTrigger>
              <SelectContent>
                <SelectGroup>
                  <SelectLabel>LLama</SelectLabel>
                  {Object.values(CreateFineTuneInputModelEnum).map((model) => (
                    <SelectItem key={model} value={model}>
                      {model}
                    </SelectItem>
                  ))}
                </SelectGroup>
                <SelectGroup>
                  <SelectLabel>OpenAI</SelectLabel>
                  {Object.values(FineTuneOpenaiInputModelEnum).map((model) => (
                    <SelectItem key={model} value={model} disabled={noOpenAIApiKey}>
                      {model} {noOpenAIApiKey ? " - Set OpenAI API Key to enable" : null}
                    </SelectItem>
                  ))}
                </SelectGroup>
                <SelectGroup>
                  <SelectLabel>Cohere</SelectLabel>
                  {Object.values(FineTuneCohereInputModelEnum).map((model) => (
                    <SelectItem key={model} value={model} disabled={noCohereApiKey}>
                      {model} {noCohereApiKey ? " - Set Cohere API Key to enable" : null}
                    </SelectItem>
                  ))}
                </SelectGroup>
              </SelectContent>
            </Select>
            {isBasicPlan ? (
              <h1 className={"text-destructive my-4 text-center"}>Upgrade your plan to fine-tune this dataset!</h1>
            ) : null}
          </div>
        </div>
        <DialogFooter>
          <Button onClick={handleStartFineTune} disabled={isStartDisabled}>
            {fetcher.state === "submitting" || fetcher.state === "loading" ? <LoadingSpinner /> : null}&nbsp;Start
            Fine-Tune
          </Button>
        </DialogFooter>
      </DialogContent>
    </Dialog>
  );
};
