import { Signal } from "@char/aftercare";
import { api } from "./client.ts";
import type { CommandInfo } from "./conversation.tsx";
import { Scope, swap, when } from "./ui.ts";

export type TreeEntry = {
  id: number;
  parent: number | null;
  branch: number;
  kind: string;
  text: string;
  time?: number;
  imageCount?: number;
  tools?: string[];
  error?: true;
};
export type Branch = {
  id: number;
  created: number;
  tip: number | null;
};
export type ConversationTree = { head: number; branches: Branch[]; entries: TreeEntry[] };

/**
 * A conversation, or a draft that continues it from an entry, or with `fork` a new conversation from one; and what its
 * composer starts with.
 */
export type Place = { tid: string; at?: number; fork?: boolean; text?: string };

type Filter = "messages" | "user" | "all";
const FILTERS: Record<Filter, { label: string; shows: (entry: TreeEntry) => boolean }> = {
  messages: {
    label: "messages",
    shows: (entry) =>
      entry.kind === "pi.user" || entry.kind === "pi.compaction" || entry.kind === "pi.reset" ||
      (entry.kind === "pi.assistant" && (entry.text !== "" || entry.error === true)),
  },
  user: { label: "yours", shows: (entry) => entry.kind === "pi.user" },
  all: { label: "everything", shows: () => true },
};

type Row = { entry: TreeEntry; lead: string; branches: Branch[]; onPath: boolean };

/**
 * The rows of a tree as tree(1) draws one, except that a run without branch points keeps its indent, as pi's does.
 * An entry `shows` leaves out hands its children to its nearest shown ancestor; a branch's end always shows.
 */
function layout(
  tree: ConversationTree,
  shows: (entry: TreeEntry) => boolean,
  path: Set<number>,
): Row[] {
  const byId = new Map(tree.entries.map((entry) => [entry.id, entry]));
  const ends = Map.groupBy(tree.branches, (branch) => branch.tip);
  const shown = (id: number) => {
    const entry = byId.get(id);
    return entry !== undefined && (shows(entry) || ends.has(id));
  };
  const children = new Map<number | null, TreeEntry[]>();
  for (const entry of tree.entries) {
    if (!shown(entry.id)) continue;
    let parent = entry.parent;
    while (parent !== null && !shown(parent)) parent = byId.get(parent)?.parent ?? null;
    children.set(parent, [...children.get(parent) ?? [], entry]);
  }

  const rows: Row[] = [];
  const stack: { entry: TreeEntry; lead: string; indent: string }[] = [];
  const descend = (list: readonly TreeEntry[], indent: string) => {
    // the way to where we are comes first
    const kids = [...list.filter((e) => path.has(e.id)), ...list.filter((e) => !path.has(e.id))];
    if (kids.length === 1) return void stack.push({ entry: kids[0], lead: indent, indent });
    for (let i = kids.length - 1; i >= 0; i--) {
      const last = i === kids.length - 1;
      stack.push({
        entry: kids[i],
        lead: indent + (last ? "└─ " : "├─ "),
        indent: indent + (last ? "   " : "│  "),
      });
    }
  };
  descend(children.get(null) ?? [], "");
  while (stack.length > 0) {
    const { entry, lead, indent } = stack.pop()!;
    rows.push({ entry, lead, branches: ends.get(entry.id) ?? [], onPath: path.has(entry.id) });
    descend(children.get(entry.id) ?? [], indent);
  }
  return rows;
}

/** A row's label, by kind, and its text. */
export function describe(entry: TreeEntry): [label: string, text: string] {
  switch (entry.kind) {
    case "pi.user":
      return [
        "user",
        [
          entry.text,
          entry.imageCount ? `[${entry.imageCount} image${entry.imageCount === 1 ? "" : "s"}]` : "",
        ].filter(Boolean).join(" "),
      ];
    case "pi.assistant":
      return ["assistant", entry.text || `→ ${entry.tools?.join(", ") ?? ""}`];
    case "pi.tool-result":
      return [entry.tools?.[0] ?? "result", entry.text];
    case "pi.system":
      return ["prompt", ""];
    case "pi.reset":
      return ["context reset", entry.text];
    case "pi.compaction":
      return ["compacted", entry.text];
    case "berry.event":
      return ["event", entry.text];
    case "berry.start":
      return ["start", ""];
    case "berry.subagent-report":
      return ["report", entry.text];
    default:
      return [entry.kind, entry.text];
  }
}

/** A move along a conversation's own tree, which its view shows once its transcript has caught up. */
export type Move = {
  how: "continue" | "edit" | "regenerate";
  /** The entry chosen: one to continue after, a message of the user's to edit, or the agent's to write afresh. */
  entry: TreeEntry;
  /** Where the conversation continues from, and where it was. */
  at: number;
  previous: number | null;
  /** Whether a transcript ending at `latest` shows it, which a regenerating run may already have gone past. */
  landed: (latest: number | null) => boolean;
};

/** A picker for continuing a conversation or starting a separate one from any entry. */
function open(
  tree: ConversationTree,
  place: Place,
  mode: "tree" | "fork",
  go: (place: Place) => void,
  moved?: (move: Move) => Promise<void>,
) {
  const scope = new Scope();
  const dialog = <dialog class="dialog tree-dialog" /> as HTMLDialogElement;
  dialog.addEventListener("close", () => {
    scope.dispose();
    dialog.remove();
  });
  // a tap outside closes it too, for screens without escape
  dialog.addEventListener("click", (event) => {
    const box = dialog.getBoundingClientRect();
    const outside = event.clientX < box.left || event.clientX > box.right ||
      event.clientY < box.top || event.clientY > box.bottom;
    if (event.target === dialog && outside) dialog.close();
  });

  const byId = new Map(tree.entries.map((entry) => [entry.id, entry]));
  const here = place.at ?? tree.branches.find((b) => b.id === tree.head)?.tip ?? null;
  const path = new Set<number>();
  for (let id = here; id !== null; id = byId.get(id)?.parent ?? null) path.add(id);
  const newest = tree.entries.at(-1)?.id ?? 0;

  const failure = <p class="tree-error" role="alert" hidden /> as HTMLParagraphElement;
  const fail = (message: string) => {
    failure.textContent = message;
    failure.hidden = false;
  };
  let choosing = false;
  // the picker stays open until the transcript shows the choice, so it never closes onto the old one
  const choose = async (entry: TreeEntry, pending: Element, regenerate = false) => {
    if (choosing) return;
    failure.hidden = true;
    const how = regenerate ? "regenerate" : entry.kind === "pi.user" ? "edit" : "continue";
    const at = how === "continue" ? entry.id : entry.parent;
    if (at === null) return fail("nothing comes before this message to branch from");
    const text = how === "edit" ? { text: entry.text } : {};
    choosing = true;
    pending.setAttribute("aria-busy", "true");
    try {
      if (mode === "fork") {
        const { tid } = await api<{ tid: string }>(
          "POST",
          `/api/conversations/${place.tid}/fork`,
          { entryId: at },
        );
        dialog.close();
        return go({ tid, ...text });
      }
      if (how === "regenerate") {
        await api("POST", `/api/conversations/${place.tid}/regenerate`, { entryId: entry.id });
      } else {
        await api("POST", `/api/conversations/${place.tid}/branch`, { entryId: at });
      }
      await moved?.({
        how,
        entry,
        at,
        previous: here,
        landed: (latest) => latest === at || (latest !== null && latest > newest),
      });
      dialog.close();
      go({ tid: place.tid, ...text });
    } catch (error) {
      fail(error instanceof Error ? error.message : String(error));
    } finally {
      choosing = false;
      pending.removeAttribute("aria-busy");
    }
  };
  const row = ({ entry, lead, branches, onPath }: Row) => {
    const [label, text] = describe(entry);
    const element = (
      <div class="tree-entry">
        <button
          type="button"
          class={entry.id === here ? "tree-row here" : "tree-row"}
          dataset={{ id: String(entry.id), kind: entry.kind, error: String(entry.error ?? false) }}
          _onclick={() => choose(entry, element)}
        >
          <span class="lead">{lead}{onPath ? "• " : ""}</span>
          <span class="label">{label}</span>
          <span class="text">{text}</span>
          {branches.map((branch) => (
            <time class={branch.id === tree.head ? "branch current" : "branch"}>
              {when(new Date(branch.created))}
            </time>
          ))}
        </button>
        {mode === "tree" && entry.kind === "pi.assistant"
          ? (
            <button
              type="button"
              class="tree-regenerate"
              dataset={{ id: String(entry.id) }}
              disabled={entry.parent === null}
              aria-label={`regenerate: ${text}`}
              _onclick={() => choose(entry, element, true)}
            >
              regenerate
            </button>
          )
          : ""}
        {entry.time === undefined
          ? ""
          : (
            <time class="time" title={new Date(entry.time).toLocaleString()}>
              {when(new Date(entry.time), { seconds: true })}
            </time>
          )}
      </div>
    );
    return element;
  };

  // focus follows the reader: a filter that hides their row moves them to its nearest shown ancestor
  let selected = here;
  const list = <div class="tree-rows" /> as HTMLDivElement;
  list.addEventListener("focusin", (event) => {
    const id = (event.target as HTMLElement).dataset.id;
    if (id !== undefined) selected = Number(id);
  });
  list.addEventListener("keydown", (event) => {
    const rows = [...list.querySelectorAll<HTMLElement>(".tree-row")];
    const step = ({
      ArrowUp: -1,
      ArrowDown: 1,
      PageUp: -15,
      PageDown: 15,
      Home: -Infinity,
      End: Infinity,
    } as Record<string, number>)[event.key];
    if (step === undefined) return;
    event.preventDefault();
    const at = rows.findIndex((row) => Number(row.dataset.id) === selected);
    rows[Math.max(0, Math.min(rows.length - 1, at + step))]?.focus();
  });
  const refocus = () => {
    let id = selected;
    while (id !== null && !list.querySelector(`.tree-row[data-id="${id}"]`)) {
      id = byId.get(id)?.parent ?? null;
    }
    const target = id === null
      ? list.querySelector<HTMLElement>(".tree-row")
      : list.querySelector<HTMLElement>(`.tree-row[data-id="${id}"]`);
    target?.focus({ preventScroll: true });
    target?.scrollIntoView({ block: "center" });
  };

  const filter = new Signal<Filter>("messages");
  swap(scope, list, filter, (chosen) => layout(tree, FILTERS[chosen].shows, path).map(row));
  scope.add(filter.subscribe(() => queueMicrotask(refocus)).unsubscribe);

  dialog.append(
    <h2>{mode}</h2>,
    <p class="desc">
      {mode === "tree" ? "continue from any point" : "fork this conversation from any point"}
    </p>,
    <div class="tree-filters">
      {(Object.keys(FILTERS) as Filter[]).map((key) => (
        <label>
          <input
            type="radio"
            name="tree-filter"
            checked={key === filter.get()}
            _onchange={() => filter.set(key)}
          />
          {FILTERS[key].label}
        </label>
      ))}
    </div>,
    list,
    failure,
    <div class="hints">
      <span>
        <kbd>↑</kbd>
        <kbd>↓</kbd> move
      </span>
      <span>
        <kbd>enter</kbd> choose
      </span>
      <span>
        <kbd>esc</kbd> close
      </span>
    </div>,
  );
  document.body.append(dialog);
  dialog.showModal();
  refocus();
}

/** pi's `/tree`, `/fork` and `/clone`, for the view at `place`. */
export function branchCommands(
  place: Place,
  go: (place: Place) => void,
  moved?: (move: Move) => Promise<void>,
): CommandInfo[] {
  const show = (mode: "tree" | "fork") => async () => {
    const tree = await api<ConversationTree>("GET", `/api/conversations/${place.tid}/tree`);
    open(tree, place, mode, go, moved);
  };
  return [
    {
      name: "tree",
      extension: "berry",
      description: "continue from any point of conversation history",
      run: show("tree"),
    },
    {
      name: "fork",
      extension: "berry",
      description: "fork this conversation from any point of history",
      run: show("fork"),
    },
    {
      name: "clone",
      extension: "berry",
      description: "copy this conversation into a new one",
      run: async () => {
        const { tid } = await api<{ tid: string }>(
          "POST",
          `/api/conversations/${place.tid}/fork`,
          place.at === undefined ? {} : { entryId: place.at },
        );
        go({ tid });
      },
    },
  ];
}