import { JsonInput, Tabs, Textarea } from "@mantine/core";
import { trpc } from "../../trpc/client";
import { create } from "zustand";
import type { Message as UIMessage } from "ai";
import type { OtherParameters, Store } from "./types.js";

const defaultSystemPrompt = `You are a helpful assistant that answers questions based on the provided context. If you don't know the answer, just say that you don't know, don't try to make up an answer.`;
const defaultParameters = {
  temperature: 0.5,
  max_tokens: 100,
} as OtherParameters;

const useStore = create<Store>()((set) => ({
  messages: [],
  message: "",
  systemPrompt: defaultSystemPrompt,
  parameters: defaultParameters,
  loading: false,
  setMessages: (messages) => set({ messages }),
  setMessage: (message) => set({ message }),
  setSystemPrompt: (systemPrompt) => set({ systemPrompt }),
  setParameters: (parameters) => set({ parameters }),
  setLoading: (loading) => set({ loading }),
}));

export default function ChatPage() {
  const messages = useStore((state) => state.messages);
  const message = useStore((state) => state.message);
  const systemPrompt = useStore((state) => state.systemPrompt);
  const parameters = useStore((state) => state.parameters);
  const loading = useStore((state) => state.loading);
  const setMessages = useStore((state) => state.setMessages);
  const setMessage = useStore((state) => state.setMessage);
  const setSystemPrompt = useStore((state) => state.setSystemPrompt);
  const setParameters = useStore((state) => state.setParameters);
  const setLoading = useStore((state) => state.setLoading);

  return (
    <Tabs defaultValue="message">
      <Tabs.List>
        <Tabs.Tab value="message">Message</Tabs.Tab>
        <Tabs.Tab value="system-prompt">System Prompt</Tabs.Tab>
        <Tabs.Tab value="parameters">Parameters</Tabs.Tab>
      </Tabs.List>
      <Tabs.Panel value="message">
        <Messages messages={messages} />
        <Textarea
          resize="vertical"
          placeholder="Type your message here..."
          value={message}
          disabled={loading}
          onChange={(e) => setMessage(e.target.value)}
          onKeyDown={async (e) => {
            if (e.key === "Enter") {
              e.preventDefault();
              const messagesWithNewUserMessage = [
                ...messages,
                { role: "user" as const, content: message } as UIMessage,
              ];
              setMessages(messagesWithNewUserMessage);
              setLoading(true);
              const response = await trpc.chat.sendMessage.query({
                messages: messagesWithNewUserMessage,
                systemPrompt,
                parameters,
              });
              const messagesWithAssistantMessage = [
                ...messagesWithNewUserMessage,
                { role: "assistant", content: response.text } as UIMessage,
              ];
              setMessages(messagesWithAssistantMessage);
              setMessage("");
              setLoading(false);
            }
          }}
        />
      </Tabs.Panel>
      <Tabs.Panel value="system-prompt">
        <Textarea
          resize="vertical"
          placeholder={defaultSystemPrompt}
          value={systemPrompt}
          onChange={(e) => setSystemPrompt(e.target.value)}
        />
      </Tabs.Panel>
      <Tabs.Panel value="parameters">
        <JsonInput
          resize="vertical"
          formatOnBlur
          placeholder={JSON.stringify(defaultParameters)}
          value={JSON.stringify(parameters)}
          onChange={(value) => setParameters(JSON.parse(value))}
        />
      </Tabs.Panel>
    </Tabs>
  );
}

function Messages({
  messages,
}: {
  messages: Array<UIMessage>;
}) {
  return (
    <div>
      {messages.map((message, index) => (
        // biome-ignore lint/suspicious/noArrayIndexKey: <explanation>
        <div key={index}>{message.content}</div>
      ))}
    </div>
  );
}