import { useExternalStoreRuntime, type AppendMessage } from "@assistant-ui/react";
import { useMemo, useSyncExternalStore } from "react";
import type { MessageInput } from "@agent-studio/shared";
import { toThreadMessageLike } from "./message-mapper.js";
import type { HarnessClientState } from "../store/reducer.js";

export type HarnessRuntimeStore = {
  getSnapshot(): HarnessClientState;
  subscribe(subscriber: () => void): () => void;
  sendMessage(input: MessageInput): Promise<unknown>;
  cancel(): Promise<unknown>;
};

export function useHarnessRuntime(store: HarnessRuntimeStore, options: { isSendDisabled?: boolean } = {}) {
  const snapshot = useSyncExternalStore(
    (subscriber) => store.subscribe(subscriber),
    () => store.getSnapshot(),
    () => store.getSnapshot(),
  );
  const messages = useMemo(
    () => snapshot.messageOrder.flatMap((id) => {
      const message = snapshot.messages[id];
      if (message === undefined) return [];
      const tools = Object.values(snapshot.tools).filter((tool) => tool.runId === message.runId);
      return [toThreadMessageLike(message, tools)];
    }),
    [snapshot.messageOrder, snapshot.messages, snapshot.tools],
  );
  const isRunning = Object.values(snapshot.runs).some((run) => run.status === "running");

  return useExternalStoreRuntime({
    messages,
    convertMessage: (message) => message,
    isRunning,
    isSendDisabled: options.isSendDisabled,
    onNew: async (message) => {
      await store.sendMessage(extractMessageInput(message));
    },
    onCancel: async () => {
      await store.cancel();
    },
  });
}

export function extractMessageInput(message: AppendMessage): MessageInput {
  const text = message.content
    .filter((part) => part.type === "text")
    .map((part) => part.text)
    .join("\n")
    .trim();
  const attachmentIds = (message.attachments ?? [])
    .map((attachment) => attachment.id)
    .filter((id) => /^att_[a-zA-Z0-9_-]+$/.test(id));
  return { text, attachmentIds: [...new Set(attachmentIds)] };
}