import type { UIMessage } from "ai";
import type { generateText } from "ai";
import type { Conversation, Fact, FactTrigger } from "./database/common";
import type { chat } from "./server/trpc/chat";

export type OtherParameters = Omit<
  Parameters<typeof generateText>[0],
  "model" | "messages" | "abortSignal" | "prompt"
>;

export type ConversationUI = Conversation & {};

/** Helper to extract the inner type of an async iterable. */
type AsyncIterated<T> = T extends AsyncIterable<infer U> ? U : never;
export type SendMessageStatus = AsyncIterated<
  Awaited<ReturnType<typeof chat.sendMessage>>
>;
export type SendMessageStatusUI =
  | SendMessageStatus
  | {
      readonly status: "error";
      readonly message: string;
      readonly result?: undefined;
    }
  | null;

export type Store = {
  /** This is a string because Milvus sends it as a string, and the value
   * overflows the JS integer anyway. */
  selectedConversationId: string;
  conversations: Array<ConversationUI>;
  messages: Array<DraftMessage | CommittedMessage>;
  message: string;
  systemPrompt: string;
  parameters: OtherParameters;
  facts: Array<Fact>;
  factTriggers: Array<FactTrigger>;
  loading: boolean;
  sendMessageStatus: SendMessageStatusUI;
  isSendingMessage: boolean;
  setConversationId: (conversationId: string) => void;
  setConversationTitle: (conversationTitle: string) => void;
  setConversations: (conversations: Array<ConversationUI>) => void;
  addConversation: (conversation: ConversationUI) => void;
  removeConversation: (conversationId: string) => void;
  setMessages: (messages: Array<DraftMessage | CommittedMessage>) => void;
  setMessage: (message: string) => void;
  setSystemPrompt: (systemPrompt: string) => void;
  setParameters: (parameters: OtherParameters) => void;
  setFacts: (facts: Array<Fact>) => void;
  setFactTriggers: (factTriggers: Array<FactTrigger>) => void;
  removeFact: (factId: string) => void;
  removeFactTrigger: (factTriggerId: string) => void;
  setLoading: (loading: boolean) => void;
  setSendMessageStatus: (status: SendMessageStatusUI) => void;
  setIsSendingMessage: (isSending: boolean) => void;
};

/** The message while it's being typed in the input box. */
export type DraftMessage = Omit<UIMessage, "id" | "createdAt">;
/** The message after it's been saved to the database. */
export type CommittedMessage = DraftMessage & {
  id: string;
  conversationId: string;
  index: number;
  runningSummary?: string;
  createdAt: string;
};