import { Label, Metric, 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 } from "@onvo-ai/js";

import { twMerge } from "tailwind-merge";
import { WidgetLibrary } from "../../components/WidgetLibrary";
import { ChevronRightIcon, XMarkIcon } from "@heroicons/react/16/solid";
import { useTheme } from "../Dashboard/useTheme";
import { PromptInput } from "../../components/PromptInput";
import { useMaxHeight } from "../../lib/useMaxHeight";
import { waitForJob } from "../../lib/waitForJob";
import { Bars3Icon } from "@heroicons/react/24/outline";

dayjs.extend(relativeTime);



const CopilotRaw: React.FC<{
  dashboardId: string;
  variant: "fullscreen" | "copilot";
  trigger?: React.ReactNode;
  coupled?: boolean;
}> = ({ trigger, variant, dashboardId, coupled }): React.ReactNode => {
  const { backend, team, setContainerRef } = useBackend();
  const { dashboard, setId, refreshWidgets, tab: dashboardTab } = useDashboard();
  const theme = useTheme();
  const scroller = useRef<HTMLDivElement>(null);
  const localRef = useRef<HTMLDivElement>(null);
  const { lg, sm } = useMaxHeight();

  const [loading, setLoading] = useState(false);
  const [open, setOpen] = useState(false);
  const [questionLoading, setQuestionLoading] = useState(false);
  const [selectedQuestion, setSelectedQuestion] = useState<Question>();
  const [messages, setMessages] = useState<Message[]>([]);
  const [tab, setTab] = useState<"home" | "history" | "library" | "chat">(
    "home"
  );
  const [sidebar, setSidebar] = useState(false);

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

  useEffect(() => {
    if (dashboard && dashboard.id) {
      setLoading(false);
    }
  }, [dashboard]);

  useEffect(() => {
    if (backend && (!dashboard || dashboard.id !== dashboardId) && !loading) {
      setLoading(true);
      setId(dashboardId, backend);
    }
  }, [dashboard, dashboardId, backend, loading]);

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

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

  useEffect(() => {
    if (!trigger) {
      setOpen(true);
    }
  }, [trigger]);

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

  const getMessages = async (id: string) => {
    let msg = await backend?.question(id).getMessages();
    setMessages(msg || []);
    setTimeout(scrollToBottom, 100);
  }

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


  const addToDashboard = async (message: string | any, library?: boolean) => {
    if (!selectedQuestion || !dashboard || !team) return;

    let textContent =
      typeof message === "string"
        ? message
        : message.map((a: any) => (a.type === "text" ? a.text : "")).join("\n");
    let output = {} as any;
    let code = "";

    if ((textContent || "").search("```") >= 0) {
      if (textContent.split("```python")[1]) {
        code = textContent.split("```python")[1].split("```")[0].trim();
      }
      if (textContent.split("```json")[1]) {
        let out = textContent.split("```json")[1].split("```")[0].trim();
        output = JSON.parse(out);
      }
    }

    const title = output?.options?.plugins?.title?.text || "Chart";

    let newObj: any = {
      title: title,
      layouts: {
        lg: {
          x: 0,
          y: lg,
          w: 4,
          h: output.type === "metric" ? 8 : 20,
        },
        sm: {
          x: 0,
          y: sm,
          w: 3,
          h: output.type === "metric" ? 8 : 20,
        },
      },
      messages: messages.filter((a) => (a.role === "user" || a.role === "assistant") && a.content && (a.content + "").trim() !== ""),
      dashboard: dashboard.id,
      team: team.id || "",
      code: code,
      created_at: new Date().toISOString(),
      settings: {},
      use_in_library: library,
      tab: dashboardTab || 0
    };

    if (!backend) return;

    toast.promise(
      async () => {
        let wid = await backend.widgets.create(newObj, output);
        return wid;
      },
      {
        loading: library
          ? "Adding widget to library..."
          : "Adding widget to dashboard...",
        success: (widget) => {
          refreshWidgets(backend);

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

          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 {
      setQuestionLoading(true);
      setTimeout(() => {
        setTab("chat");
        setMessages([
          {
            role: "user",
            content: query,
            created_at: dayjs().toISOString(),
            id: "new",
          } as any,
        ])
      }, 100);
      let response = await backend?.questions.create({
        dashboard: dashboard?.id,
        query: query,
      });

      if (!response) {
        toast.error("Failed to create question.");
        return;
      }

      setSelectedQuestion(response);
      await waitForJob(response._meta.url);
      getMessages(response.id);
      setQuestionLoading(false);

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

  const handleAskQuestion = async (query: string) => {
    if (!selectedQuestion) return;
    try {
      setQuestionLoading(true);
      const response = await backend?.question(selectedQuestion.id).completion(query);
      getMessages(selectedQuestion.id);

      if (!response) {
        toast.error("Failed to ask question.");
        setQuestionLoading(false);
        return;
      }
      await waitForJob(response._meta.url);
      getMessages(selectedQuestion.id);
      setQuestionLoading(false);

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

  const handleEditMessage = async (messageId: string, content: string) => {
    if (!selectedQuestion) return;
    try {
      setQuestionLoading(true);
      let response = await backend?.question(selectedQuestion.id).updateMessage(messageId, content);
      getMessages(selectedQuestion.id);
      if (!response) {
        toast.error("Failed to edit message.");
        setQuestionLoading(false);
        return;
      }
      console.log("RESPONSE: ", response);
      await waitForJob(response._meta.url);
      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}
        message={message}
        isLast={index === messages.length - 1}
        onReply={(msg) => {
          handleAskQuestion(msg);
        }}
        onEdit={(msg) => {
          handleEditMessage(message.id, msg);
        }}
        onAdd={(library) => {
          addToDashboard(message.content, library);
        }}
      />
    ));
  }, [messages]);

  return (
    <>
      <dialog open={open}>

        <div ref={localRef}
          className={twMerge(
            "onvo-root-style onvo-copilot-modal onvo-@container/questionmodal onvo-foreground-color onvo-animate-dialogOpen onvo-z-[50] onvo-fixed",
            variant === "fullscreen"
              ? "onvo-h-full onvo-w-full onvo-left-0"
              : "onvo-h-[calc(100vh-40px)] onvo-w-[480px] onvo-right-5 onvo-bottom-5 onvo-rounded-2xl onvo-overflow-hidden 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-h-12 onvo-top-0 onvo-w-full onvo-z-10 onvo-flex onvo-flex-row onvo-justify-between onvo-items-center onvo-gap-4 onvo-border-solid onvo-border-b onvo-px-3 onvo-border-black/10 onvo-py-2 dark:onvo-border-white/10"
              }
            >
              {tab === "chat" ? (
                <Icon
                  icon={Bars3Icon}
                  variant="shadow" className="onvo-border-black/10 onvo-background-color dark:onvo-border-white/10"
                  onClick={() => {
                    setSidebar(!sidebar);
                  }}
                />) : <div></div>}
              <div className="onvo-flex onvo-w-full @xl/questionmodal:onvo-w-auto onvo-flex-row onvo-gap-1 onvo-justify-start onvo-items-center">
                <Text className="onvo-hidden @xl/questionmodal:onvo-block">
                  {dashboard?.title}
                </Text>
                <ChevronRightIcon className="onvo-hidden @xl/questionmodal:onvo-block onvo-h-4 onvo-w-4 dark:onvo-fill-slate-500" />
                <Label>{dashboard?.settings?.copilot_title || "Copilot"}</Label>
              </div>
              {trigger ? (
                <Icon
                  icon={XMarkIcon}
                  variant="shadow" className="onvo-border-black/10 onvo-background-color dark:onvo-border-white/10"
                  onClick={() => {
                    setOpen(false);
                  }}
                />) : <div></div>}
            </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 ">
              {["home", "library", "history"].indexOf(tab) >= 0 && (
                <div className="onvo-py-2 onvo-px-2 onvo-w-full onvo-max-w-screen-lg onvo-flex onvo-mx-auto onvo-flex-col onvo-gap-4">
                  {tab === "home" && (
                    <div className="onvo-w-full onvo-flex onvo-flex-col onvo-items-center">
                      <Metric className="onvo-pt-10 onvo-text-center onvo-mb-4">
                        {dashboard?.settings?.copilot_description ||
                          "Create a widget"}
                      </Metric>
                      <PromptInput
                        hideSuggestions={true}
                        url={dashboard?.settings?.help_url}
                        onSubmit={(val) => {
                          createQuestion(val);
                          setTab("chat");
                        }}
                      />
                    </div>
                  )}
                  {(tab === "home" || tab === "library") && (
                    <WidgetLibrary
                      onExpanded={(bool) =>
                        bool ? setTab("library") : setTab("home")
                      }
                    />
                  )}

                  {(tab === "home" || tab === "history") && (
                    <QuestionHistory
                      variant="default"
                      onSelect={(q) => {
                        setSelectedQuestion(q);
                        setMessages([]);
                        setTab("chat");
                      }}
                      onExpanded={(bool) =>
                        bool ? setTab("history") : setTab("home")
                      }
                      onDelete={() => {
                        setSelectedQuestion(undefined);
                      }}
                    />
                  )}
                </div>
              )}

              {tab === "chat" && (
                <div className="onvo-flex onvo-flex-row onvo-w-full onvo-h-full">
                  <div className={"onvo-z-50 onvo-left-0 onvo-absolute @xl/questionmodal:onvo-relative onvo-overflow-x-hidden onvo-flex-shrink-0 onvo-transition-all onvo-ease-in-out " + (sidebar ? "onvo-w-[calc(100%-16px)] @xl/questionmodal:onvo-w-[260px]" : "onvo-w-0 @xl/questionmodal:onvo-w-0")}>
                    <QuestionHistory
                      variant="sidebar"
                      onSelect={(q) => {
                        setSelectedQuestion(q);
                        setMessages([]);
                        setTab("chat");
                      }}
                      selectedId={selectedQuestion?.id}
                      onNew={() => {
                        setSelectedQuestion(undefined);
                        setMessages([]);
                        setTab("home");
                      }}
                      onExpanded={(bool) =>
                        bool ? setTab("history") : setTab("home")
                      }
                      onDelete={() => {
                        setSelectedQuestion(undefined);
                      }}
                    />
                  </div>
                  <div className="onvo-w-full onvo-h-full onvo-flex onvo-flex-col onvo-px-2">
                    <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-4"
                      onSubmit={val => {
                        handleAskQuestion(val);
                      }} />
                  </div>
                </div>
              )}
            </div>
          </div>
        </div>
      </dialog>
      {trigger && React.cloneElement(trigger as React.ReactElement, {
        onClick: (e: React.MouseEvent) => {
          setOpen(true);
        },
      })}
    </>
  );
};

export const Copilot = memo(CopilotRaw);