import { BACKGROUND_CONTEXT as context } from "@earendil-works/chord/context";
import type { ConversationId, Cursor, EntryId, EntryRecord } from "@earendil-works/pi-durable";
import { type Berry, BerryError, Conversations, EventEntry, headOf } from "./berry.ts";

/** One entry of a conversation's branches, as a tree of them shows it. */
export type TreeEntry = {
  id: number;
  /**
   * The entry it follows: the one before it in its branch, or for a branch's first entry, the entry the branch
   * continues from.
   */
  parent: number | null;
  /** The branch it was written in, or for what a fork inherited, the branch it began as. */
  branch: number;
  kind: string;
  /**
   * A user's message in full, to be edited again; another message's text, clipped to a line's worth; an event's
   * extension.
   */
  text: string;
  /** When its message was written; entries without one, like events, have none. */
  time?: number;
  imageCount?: number;
  /** The tools an assistant message called, or the tool a result answers. */
  tools?: string[];
  error?: true;
};

export type Branch = {
  id: number;
  created: number;
  /** Its newest entry, which may be its parent's when it has none of its own. */
  tip: number | null;
};

/** A conversation's branches, and which one it continues on. */
export type ConversationTree = { head: number; branches: Branch[]; entries: TreeEntry[] };

const PREVIEW = 240;

/** `text` on one line, cut to `length` characters. */
export function clip(length: number, text: string): string {
  const flat = text.replace(/\s+/g, " ").trim();
  return flat.length > length ? `${flat.slice(0, length - 1)}…` : flat;
}

export const textOf = (content: string | readonly { type: string; text?: string }[]) =>
  typeof content === "string"
    ? content
    : content.flatMap((block) =>
      block.type === "text" ? [block.text ?? ""] : block.type === "image" ? ["[image]"] : []
    ).join(" ");

/** A user's message as they wrote it, without its images. */
export const userText = (content: string | readonly { type: string; text?: string }[]) =>
  typeof content === "string"
    ? content
    : content.flatMap((block) => (block.type === "text" ? [block.text ?? ""] : [])).join("\n");

function preview(
  entry: EntryRecord,
): Pick<TreeEntry, "text" | "tools" | "error" | "imageCount" | "time"> {
  const message = entry.model?.[0];
  const written = message?.timestamp === undefined ? {} : { time: message.timestamp };
  if (message?.role === "assistant") {
    const tools = message.content.flatMap((
      block,
    ) => (block.type === "toolCall" ? [block.name] : []));
    const text = clip(PREVIEW, textOf(message.content)) ||
      (message.stopReason === "aborted" ? "(aborted)" : clip(PREVIEW, message.errorMessage ?? ""));
    return {
      ...written,
      text,
      ...(tools.length > 0 ? { tools } : {}),
      ...(message.stopReason === "error" ? { error: true } : {}),
    };
  }
  if (message?.role === "toolResult") {
    return {
      ...written,
      text: clip(PREVIEW, textOf(message.content)),
      tools: [message.toolName],
      ...(message.isError ? { error: true } : {}),
    };
  }
  if (message?.role === "user") {
    const imageCount = typeof message.content === "string"
      ? 0
      : message.content.filter((block) => block.type === "image").length;
    return {
      ...written,
      text: entry.kind === "pi.user"
        ? userText(message.content)
        : clip(PREVIEW, textOf(message.content)),
      ...(imageCount > 0 ? { imageCount } : {}),
    };
  }
  if (EventEntry.is(entry)) return { ...written, text: entry.data.ext };
  return { ...written, text: "" };
}

/**
 * Every branch of `tid` as one tree of entries. Entry IDs are session-wide and a branch continues from one entry of
 * another, so the branches' own entries join into one.
 */
export async function conversationTree(berry: Berry, tid: string): Promise<ConversationTree> {
  const info = (await berry.harness.snapshot(Conversations, context))?.[tid];
  if (info === undefined) throw new BerryError(`No conversation ${tid}`);
  const all = [{ id: info.id, at: null, created: info.created }, ...(info.branches ?? [])];

  const branches = await Promise.all(all.map(async ({ id, at, created }) => {
    const conversation = (await berry.harness.conversation(id as ConversationId, context))!;
    // a branch's own entries are those after the entry it continues from; the first has all of its history, what
    // a fork inherited included
    const query = { order: "ascending", minEntryId: ((at ?? 0) + 1) as EntryId } as const;
    const entries: TreeEntry[] = [];
    let previous: number | null = at;
    let cursor: Cursor | undefined;
    do {
      const page = await conversation.entries(query, 256, cursor, context);
      for (const entry of page.items) {
        entries.push({
          id: entry.id,
          parent: previous,
          branch: id,
          kind: entry.kind,
          ...preview(entry),
        });
        previous = entry.id;
      }
      cursor = page.next;
    } while (cursor !== undefined);
    return { branch: { id, created, tip: previous }, entries };
  }));
  return {
    head: headOf(info),
    branches: branches.map(({ branch }) => branch),
    entries: branches.flatMap(({ entries }) => entries).sort((a, b) => a.id - b.id),
  };
}