import { useEffect, useRef, useState } from "react";
import { FileUpload } from "./FileUpload";

function Message({ message, role }: Message) {
  return (
    <pre
      className={`p-4 mb-4 rounded-3xl shadow-md whitespace-pre-wrap ${
        role === "user"
          ? "bg-blue-200 text-blue-800 ml-20 self-end"
          : "bg-green-200 text-green-800 mr-20 self-start"
      }`}
    >
      {message}
    </pre>
  );
}

interface Message {
  message: string;
  role: "user" | "assistant";
}

function App() {
  const scrollUpRef = useRef(false);
  const [result, setResult] = useState("");
  const [error, setError] = useState("");
  const [isStreaming, setIsStreaming] = useState(false);
  const [messages, setMessages] = useState<Message[]>([]);
  const [isFileUploaded, setIsFileUploaded] = useState(false);

  useEffect(() => {
    function handleWheel(event: WheelEvent) {
      if (isStreaming && event.deltaY < 0) {
        scrollUpRef.current = true;
      }
    }

    document.body.addEventListener("wheel", handleWheel);

    return () => {
      document.body.removeEventListener("wheel", handleWheel);
    };
  }, [isStreaming]);

  const onKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
    if (e.key === "Enter" && !e.shiftKey) {
      e.preventDefault();
      e.currentTarget.form?.requestSubmit();
    }
  };

  async function onSubmit(e: React.FormEvent<HTMLFormElement>) {
    e.preventDefault();
    const form = e.currentTarget;
    const formData = new FormData(form);
    const msg = formData.get("msg") as string;
    const scrollingElement = document.scrollingElement;

    if (msg.trim() === "") return;

    scrollUpRef.current = false;
    scrollingElement?.scrollTo({
      behavior: "smooth",
      top: scrollingElement?.scrollHeight,
    });

    setIsStreaming(true);
    setMessages((prev) => [
      ...prev,
      {
        message: msg,
        role: "user",
      },
    ]);

    try {
      const session_id = sessionStorage.getItem("session_id");
      const response = await fetch(form.action, {
        method: "POST",
        headers: session_id
          ? {
              "X-Session-ID": session_id,
            }
          : undefined,
        body: msg,
      });

      let result = "";
      const reader = response.body?.getReader();
      const decoder = new TextDecoder();

      form.msg.value = "";

      await reader?.read().then(function processText(msg): Promise<void> {
        const { value, done } = msg;
        const chunk = decoder.decode(value, { stream: true });
        const scrollingElement = document.scrollingElement;

        if (done) {
          setResult("");
          setMessages((prev) => [
            ...prev,
            {
              message: result,
              role: "assistant",
            },
          ]);

          if (scrollingElement) {
            scrollingElement.scrollTop = scrollingElement?.scrollHeight;
          }

          return Promise.resolve();
        }

        result += chunk;
        setResult(result);

        if (!scrollUpRef.current && scrollingElement) {
          scrollingElement.scrollTop = scrollingElement.scrollHeight;
        }

        // Read some more, and call this function again
        return reader.read().then(processText);
      });
    } catch (error) {
      setError("Error submitting message");
      console.error("Error submitting message", error);
    } finally {
      setIsStreaming(false);
    }
  }

  return (
    <div className="max-w-prose flex flex-col items-center justify-end min-h-screen pb-4 mx-auto">
      {!isFileUploaded && (
        <div className="w-full min-h-screen flex flex-col items-center justify-center">
          <h1 className="text-5xl font-bold mb-2">DocTalk</h1>
          <p className="text-xl text-gray-500 mb-10 text-center">
            Talk to your documents.
          </p>
          <FileUpload
            isFileUploaded={isFileUploaded}
            setIsFileUploaded={setIsFileUploaded}
          />
          <div className="text-xs text-center text-gray-500 mt-4">
            Built with{" "}
            <a
              target="_blank"
              href="https://openai.com/"
              rel="noopener noreferrer"
              className="underline text-blue-500"
            >
              OpenAI
            </a>
            ,{" "}
            <a
              target="_blank"
              href="https://fastapi.tiangolo.com/"
              rel="noopener noreferrer"
              className="underline text-blue-500"
            >
              FastAPI
            </a>
            ,{" "}
            <a
              target="_blank"
              href="https://qdrant.tech/"
              rel="noopener noreferrer"
              className="underline text-blue-500"
            >
              Qdrant
            </a>
            ,{" "}
            <a
              target="_blank"
              href="https://react.dev/"
              rel="noopener noreferrer"
              className="underline text-blue-500"
            >
              React
            </a>
            ,{" "}
            <a
              target="_blank"
              href="https://vitejs.dev/"
              rel="noopener noreferrer"
              className="underline text-blue-500"
            >
              Vite
            </a>
            , and{" "}
            <a
              target="_blank"
              href="https://tailwindcss.com/"
              rel="noopener noreferrer"
              className="underline text-blue-500"
            >
              TailwindCSS
            </a>
          </div>
        </div>
      )}

      <div className="w-full flex flex-col justify-end">
        {messages.map(({ message, role }) => (
          <Message message={message} role={role} key={`${message}-${role}`} />
        ))}
        {result && <Message message={result} role="assistant" />}
        {error && <p className="text-red-500 text-center mb-4">{error}</p>}
      </div>

      {isFileUploaded && (
        <form
          method="post"
          action="/api/chat"
          onSubmit={onSubmit}
          className="relative w-full"
        >
          <textarea
            id="msg"
            name="msg"
            rows={3}
            onKeyDown={onKeyDown}
            disabled={isStreaming}
            placeholder="Ask me anything about the document's contents..."
            className="block bg-white w-full p-4 pr-20 m-0 rounded-3xl shadow-md text-xl disabled:opacity-50 disabled:cursor-not-allowed resize-none"
          ></textarea>

          <button
            type="submit"
            disabled={isStreaming}
            aria-label="Send message"
            className="absolute right-0 bottom-0 bg-blue-500 text-white p-2 m-2 rounded-full shadow-md hover:bg-blue-600 focus:bg-blue-600 cursor-pointer transition-colors duration-300 disabled:opacity-50 disabled:cursor-not-allowed"
          >
            <svg
              xmlns="http://www.w3.org/2000/svg"
              viewBox="0 0 24 24"
              fill="currentColor"
              className="w-6 h-6"
            >
              <path d="M3.478 2.404a.75.75 0 0 0-.926.941l2.432 7.905H13.5a.75.75 0 0 1 0 1.5H4.984l-2.432 7.905a.75.75 0 0 0 .926.94 60.519 60.519 0 0 0 18.445-8.986.75.75 0 0 0 0-1.218A60.517 60.517 0 0 0 3.478 2.404Z" />
            </svg>
          </button>
        </form>
      )}
    </div>
  );
}

export default App;
