import { Label, Text } from "../../tremor/Text";
import { Icon } from "../../tremor/Icon";
import { memo, useEffect, useMemo, useRef, useState } from "react";
import dayjs from "dayjs";
import relativeTime from "dayjs/plugin/relativeTime";
import { toast, Toaster } from "sonner";
import { QuestionMessage } from "../../components/QuestionMessage";
import { QuestionHistory } from "../../components/QuestionHistory";
import React from "react";
import { useBackend } from "../Wrapper";
import { useDashboard } from "../Dashboard/useDashboard";
import { ChartLoader } from "../../components/ChartLoader";
import { LogType, Message, Question, Widget } from "@onvo-ai/js";

import { ChevronRightIcon } from "@heroicons/react/16/solid";
import { useTheme } from "../Dashboard/useTheme";
import { PromptInput } from "../../components/PromptInput";
import { useMaxHeight } from "../../lib/useMaxHeight";
import { ChevronDoubleRightIcon, ChevronDownIcon, ClockIcon, PencilSquareIcon, PlusIcon, ViewfinderCircleIcon, XMarkIcon } from "@heroicons/react/24/outline";
import { create } from "zustand";
import { Tooltip } from "../../tremor/Tooltip";
import { DropdownMenu, DropdownMenuContent, DropdownMenuGroup, DropdownMenuItem, DropdownMenuTrigger } from "../../tremor/DropdownMenu";
import { useIsMobile } from "../../lib/utils";

dayjs.extend(relativeTime);


export const useCopilot = create<{
  open: boolean;
  variant?: "default" | "modal";
  setOpen: (
    open: boolean,
  ) => void;
  setVariant: (
    variant: "default" | "modal",
  ) => void;
}>((set) => ({
  open: false,
  variant: "default" as const,
  setVariant: (
    variant: "default" | "modal",
  ) => {
    localStorage.setItem("onvo-copilot-variant", variant);
    set({ variant });
  },
  setOpen: (
    op: boolean,
  ) => {
    let variant = localStorage.getItem("onvo-copilot-variant");
    if (variant) {
      set({ open: op, variant: variant as "default" | "modal" })
    } else {
      set({ open: op, variant: "default" })
    }
  }
}));



const CopilotRaw: React.FC<{
  coupled?: boolean;
}> = ({ coupled }): React.ReactNode => {
  const { backend, setContainerRef, } = useBackend();
  const { dashboard, refreshWidgets, tab: dashboardTab } = useDashboard();
  const theme = useTheme();
  const { open, setOpen, setVariant, variant } = useCopilot();
  const scroller = useRef<HTMLDivElement>(null);
  const localRef = useRef<HTMLDivElement>(null);
  const { lg, sm } = useMaxHeight();
  const isMobile = useIsMobile();

  const [questionLoading, setQuestionLoading] = useState(false);
  const [selectedQuestion, setSelectedQuestion] = useState<Question>();
  const [messages, setMessages] = useState<Message[]>([]);
  const [userQuery, setUserQuery] = useState("");

  useEffect(() => {
    if (isMobile) {
      setVariant("modal");
    }
  }, [isMobile]);

  useEffect(() => {
    if (localRef?.current && !coupled) {
      setContainerRef(localRef);
    }
  }, [localRef, coupled])

  useEffect(() => {
    if (!(window as any).Onvo) {
      (window as any).Onvo = {};
    }
    (window as any).Onvo.setCopilotOpen = (val: boolean) => setOpen(val);
  }, []);

  useEffect(() => {
    if (!open) {
      setMessages([]);
      setSelectedQuestion(undefined);
      setUserQuery("");
    }
  }, [open]);

  useEffect(() => {
    if (selectedQuestion) {
      getMessages(selectedQuestion.id);
    } else {
      setMessages([]);
    }
  }, [selectedQuestion, userQuery]);

  const getMessages = async (id: string) => {
    let msg = await backend?.question(id).getMessages();
    if (!msg || msg.length === 0) {
      if (userQuery && userQuery.trim() !== "") {
        msg = [{
          role: "user",
          content: userQuery,
          created_at: dayjs().toISOString(),
          id: "new",
        } as any];
      }
    }
    setMessages(msg && msg.length > 0 ? msg : []);
    setTimeout(scrollToBottom, 100);
  }

  const scrollToBottom = () => {
    if (scroller.current) {
      scroller.current.scrollTop = scroller.current.scrollHeight;
    }
  };


  const addToDashboard = async (widget: Widget, library?: boolean) => {
    if (!backend) return;

    toast.promise(
      async () => {
        let wid = await backend.widgets.create({
          ...widget,
          use_in_library: library || false,
          use_in_chat: false,
          tab: dashboardTab || 0,
          layouts: {
            lg: {
              x: 0,
              y: lg,
              w: 4,
              h: widget.type === "metric" ? 8 : 20,
            },
            sm: {
              x: 0,
              y: sm,
              w: 3,
              h: widget.type === "metric" ? 8 : 20,
            },
          }
        });

        await backend.widget(wid.id).updateCache();
        return wid;
      },
      {
        loading: library
          ? "Adding widget to library..."
          : "Adding widget to dashboard...",
        success: (widget) => {
          refreshWidgets(backend);

          setSelectedQuestion(undefined);
          setMessages([]);
          setOpen(false);
          setUserQuery("");

          if (backend) {
            backend.logs.create({
              type: LogType.EditWidget,
              dashboard: widget.dashboard,
              widget: widget.id,
            })
          }
          if (library) {
            return "Widget added to library";
          } else {
            return "Widget added to dashboard";
          }
        },
        error: (error) =>
          library
            ? "Error adding widget to library: " + error.message
            : "Error adding widget to dashboard: " + error.message,
      }
    );
  };


  const createQuestion = async (query: string) => {
    if (!dashboard) return;

    try {
      let response = await backend?.questions.create({
        dashboard: dashboard?.id,
        query: query,
      });

      if (!response) {
        throw new Error("Failed to create question.");
      }

      setUserQuery(query);
      setSelectedQuestion(response);
      return response.id;

    } catch (e: any) {
      toast.error("Failed to create question: ", e.message);
    }
  };

  const handleAskQuestion = async (query: string) => {
    let qId = selectedQuestion?.id;
    if (!qId) {
      qId = await createQuestion(query);
    }

    if (!qId) return;
    try {
      setQuestionLoading(true);
      await backend?.question(qId).completion(query, log => {
        getMessages(qId);
        setUserQuery("");
      });
      getMessages(qId);
      setQuestionLoading(false);

    } catch (error: any) {
      console.error("Error asking question:", error);
      toast.error(error.message, {
        duration: Infinity,
        closeButton: true
      });
      setQuestionLoading(false);
    }
  };

  const handleEditMessage = async (messageId: string, content: string) => {
    if (!selectedQuestion) return;
    try {
      setQuestionLoading(true);

      const editedMessage = messages.find((msg) => msg.id === messageId);

      setMessages((prev) => {
        return prev.filter((msg) => {
          if (msg.created_at && editedMessage?.created_at && msg.created_at > editedMessage.created_at) {
            return false
          }
          return true;
        });
      });

      // waitForJob
      let response = await backend?.question(selectedQuestion.id).updateMessage(messageId, content);
      if (!response) {
        toast.error("Failed to edit message.");
        setQuestionLoading(false);
        return;
      }
      getMessages(selectedQuestion.id);
      setQuestionLoading(false);

    } catch (error) {
      console.error("Error asking question:", error);
      setQuestionLoading(false);
    }
  };

  const questionMessageList = useMemo(() => {
    return messages.map((message, index) => (
      <QuestionMessage
        key={message.id}
        questionLoading={questionLoading}
        message={message}
        isLast={index === messages.length - 1}
        onReply={(msg) => {
          handleAskQuestion(msg);
        }}
        onEdit={(msg) => {
          handleEditMessage(message.id, msg);
        }}
        onAdd={(widget, library) => {
          addToDashboard(widget, library);
        }}
      />
    ));
  }, [messages, questionLoading]);


  return (
    <>

      <div ref={localRef}
        className={
          "onvo-root-style onvo-copilot-modal onvo-@container/questionmodal onvo-foreground-color onvo-animate-dialogOpen onvo-z-40 onvo-overflow-hidden " + (variant === "default" ? "onvo-h-full onvo-relative onvo-w-full" : "onvo-fixed onvo-h-[640px] onvo-w-[420px] onvo-right-5 onvo-bottom-5 onvo-rounded-3xl onvo-border-solid onvo-border onvo-border-slate-200 dark:onvo-border-slate-800 onvo-shadow-xl")
        }
      >
        {!coupled && <Toaster position="bottom-right" richColors />}
        <div className="onvo-question-modal-question-list onvo-flex onvo-flex-col onvo-foreground-color onvo-absolute onvo-w-full onvo-right-0 onvo-top-0 onvo-z-20 onvo-h-full">
          <div
            className={
              "onvo-foreground-color onvo-relative onvo-h-auto onvo-top-0 onvo-w-full onvo-z-10 onvo-flex onvo-flex-row onvo-justify-between onvo-items-center onvo-gap-2 onvo-border-solid onvo-border-b onvo-px-3 onvo-border-black/10 onvo-py-2 dark:onvo-border-white/10"
            }
          >
            <div className="onvo-flex onvo-flex-col onvo-w-full onvo-gap-0 onvo-justify-center onvo-items-start">

              <Label>{dashboard?.settings?.copilot_title || "Copilot"}</Label>
              {(dashboard?.settings?.copilot_description || "").trim() !== "" && (
                <Text className="onvo-text-xs">{dashboard?.settings?.copilot_description || ""}</Text>
              )}
            </div>
            <Tooltip content="New chat" side="bottom">
              <Icon
                icon={PencilSquareIcon}
                className="onvo-relative onvo-border-black/10 onvo-background-color dark:onvo-border-white/10"
                onClick={() => {
                  setSelectedQuestion(undefined);
                  setMessages([]);
                }}
              />
            </Tooltip>

            <DropdownMenu>
              <Tooltip triggerAsChild content="History" side="bottom">
                <DropdownMenuTrigger asChild>
                  <Icon
                    icon={ClockIcon}
                    className="onvo-relative onvo-border-black/10 onvo-background-color dark:onvo-border-white/10"
                  />
                </DropdownMenuTrigger>
              </Tooltip>
              <DropdownMenuContent side="bottom" className="onvo-backdrop-blur-lg onvo-min-w-56 onvo-z-50 onvo-overflow-y-auto onvo-max-h-[300px]">

                <QuestionHistory
                  onSelect={(q) => {
                    setSelectedQuestion(q);
                    setMessages([]);
                  }}
                  variant="dropdown"
                />
              </DropdownMenuContent>
            </DropdownMenu>
            {coupled && !isMobile && (
              <Tooltip content={variant === "modal" ? "Show as sidebar" : "Show as modal"} side="bottom">
                <Icon
                  icon={ViewfinderCircleIcon}
                  className="onvo-relative onvo-border-black/10 onvo-background-color dark:onvo-border-white/10"
                  onClick={() => {
                    setVariant(variant === "default" ? "modal" : "default");
                  }}
                />
              </Tooltip>)}
            <Tooltip content="Close copilot" side="bottom">
              <Icon
                icon={variant === "default" ? ChevronDoubleRightIcon : XMarkIcon}
                className="relative onvo-border-black/10 onvo-background-color dark:onvo-border-white/10"
                onClick={() => {
                  setOpen(false);
                }}
              />
            </Tooltip>

          </div>
          <div className="onvo-flex onvo-flex-grow onvo-w-full onvo-h-[calc(100%-52px)] onvo-overflow-y-auto onvo-scrollbar-thin onvo-flex-col ">

            <div className="onvo-w-full onvo-h-full onvo-flex onvo-flex-col onvo-px-2">
              {!selectedQuestion ? (
                <>
                  <div className="onvo-flex-grow onvo-w-full"></div>
                  <div className={"onvo-z-50 onvo-left-0 onvo-overflow-x-hidden onvo-flex-shrink-0 onvo-transition-all onvo-ease-in-out onvo-mb-2"}>
                    <QuestionHistory
                      onSelect={(q) => {
                        setSelectedQuestion(q);
                        setMessages([]);
                      }}
                      limit={4}
                      variant="list"
                    />
                  </div></>
              ) : (
                <div
                  className="onvo-flex onvo-w-full onvo-flex-grow onvo-flex-col onvo-gap-4 onvo-overflow-y-auto onvo-scrollbar-thin onvo-py-2"
                  ref={scroller}
                >
                  <div className="onvo-flex onvo-flex-col onvo-relative onvo-mx-auto onvo-w-full onvo-max-w-screen-md">
                    {questionMessageList}
                    {questionLoading && (
                      <ChartLoader id={selectedQuestion?.id} />
                    )}
                  </div>
                </div>
              )}
              <PromptInput
                url={dashboard?.settings?.help_url}
                className="onvo-mb-2"
                onSubmit={val => {
                  handleAskQuestion(val);
                }} />
            </div>
          </div>
        </div>
      </div>
    </>
  );
};

export const Copilot: React.FC<{
  coupled?: boolean;
}> = ({ coupled }): React.ReactNode => {
  const { variant, open, setVariant } = useCopilot();

  useEffect(() => {

    let storedVariant = localStorage.getItem("onvo-copilot-variant");
    if (!coupled) {
      setVariant("modal");
      return;
    }
    if (storedVariant) {
      setVariant(storedVariant as "default" | "modal");
    } else {
      setVariant("default");
    }
  }, []);

  if (variant === "default") {
    return (
      <CopilotRaw coupled={coupled} />
    )
  }
  return (<>

    <dialog open={open}>
      <CopilotRaw coupled={coupled} />
    </dialog>
  </>)
}