import { DatapointDtoSourceEnum, TraceWithLogsDto } from "@montelo/browser-client";
import { useParams } from "@remix-run/react";
import { FC, useState } from "react";
import { Button } from "~/components/ui/button";
import { Dialog, DialogContent, DialogHeader, DialogTitle, DialogTrigger } from "~/components/ui/dialog";
import { Label } from "~/components/ui/label";
import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "~/components/ui/select";
import { useFetcherWithReset } from "~/hooks";
import { AddToDatasetContent } from "~/pages/traces/components/AddToDatasetContent";
import { Routes } from "~/routes";
import { DatasetsLoaderTypes } from "~/types/DatasetsLoader.types";

export const AddToDatasetDialog: FC<{ trace: TraceWithLogsDto }> = ({ trace }) => {
  const params = useParams();
  const fetcher = useFetcherWithReset<DatasetsLoaderTypes>();
  const [isDialogOpen, setIsDialogOpen] = useState(false);
  const [source, setSource] = useState<DatapointDtoSourceEnum>(DatapointDtoSourceEnum.Trace);
  const [sourceId, setSourceId] = useState<string>("");
  const [selectedDatasetId, setSelectedDatasetId] = useState<string | null>(null);
  const [inputSchema, setInputSchema] = useState<string>("");
  const [outputSchema, setOutputSchema] = useState<string>("");
  const [logSchema, setLogSchema] = useState<string>("");

  const selectedDataset = fetcher.data?.datasets.find((dataset) => dataset.id === selectedDatasetId);
  const canEdit = !!selectedDataset;

  const fetchDatasets = () => {
    if (fetcher.data) {
      return;
    }

    fetcher.load(
      Routes.app.project.env.datasets({
        projectId: params.projectId!,
        envId: params.envId!,
      }),
    );

    const currentUrl = window.location.href;
    const urlParams = new URLSearchParams(new URL(currentUrl).search);
    const selectedLogId = urlParams.get("logId");
    const selectedLog = trace.logs.find((log) => log.id === selectedLogId);
    if (!selectedLogId || !selectedLog) {
      setSource(DatapointDtoSourceEnum.Trace);
      setSourceId(trace.id);
      setLogSchema(JSON.stringify({ ...trace.logs[0].input, ...trace.logs.at(-1)?.output }, null, 2));
      return;
    }

    setSource(DatapointDtoSourceEnum.Log);
    setSourceId(selectedLogId);
    setLogSchema(JSON.stringify({ ...selectedLog.input, ...selectedLog.output }, null, 2));
  };

  const handleDialogChange = (open: boolean) => {
    if (!open) {
      setSelectedDatasetId(null);
      setInputSchema("");
      setLogSchema("");
      fetcher.reset();
    }
    setIsDialogOpen(open);
  };

  const handleSelect = (id: string) => {
    setSelectedDatasetId(id);
    const dataset = fetcher.data?.datasets.find((dataset) => dataset.id === id);

    if (dataset?.isFineTuning) {
      const inp = {
        messages: [
          {
            role: "",
            content: "",
          },
        ],
      };

      const out = {
        messages: [
          {
            role: "assistant",
            content: "",
          },
        ],
      };
      setInputSchema(JSON.stringify(inp, null, 2));
      setOutputSchema(JSON.stringify(out, null, 2));
      return;
    }

    setInputSchema(JSON.stringify(dataset?.inputSchema, null, 2));
    setOutputSchema(JSON.stringify(dataset?.outputSchema, null, 2));
  };

  return (
    <Dialog open={isDialogOpen} onOpenChange={handleDialogChange}>
      <DialogTrigger asChild>
        <Button onMouseEnter={fetchDatasets}>Add to Dataset</Button>
      </DialogTrigger>
      <DialogContent className="min-h-[700px] max-w-[72rem] overflow-y-scroll">
        <DialogHeader>
          <DialogTitle>Add to Dataset</DialogTitle>
        </DialogHeader>
        <div>
          <Label>Dataset</Label>
          <Select onValueChange={handleSelect}>
            <SelectTrigger className="w-[180px]">
              <SelectValue placeholder="Select a dataset" />
            </SelectTrigger>
            <SelectContent>
              {fetcher.data?.datasets.map((dataset) => (
                <SelectItem key={dataset.id} value={dataset.id}>
                  {dataset.name}
                </SelectItem>
              ))}
            </SelectContent>
          </Select>
        </div>
        <AddToDatasetContent
          source={source}
          sourceId={sourceId}
          logSchema={logSchema}
          inputSchema={inputSchema}
          setInputSchema={setInputSchema}
          outputSchema={outputSchema}
          setOutputSchema={setOutputSchema}
          canEdit={canEdit}
          selectedDataset={selectedDataset}
          handleDialogChange={handleDialogChange}
        />
      </DialogContent>
    </Dialog>
  );
};
