tree-marks

"use client"

// MARKS ON A TREE: tri-state include, exclude, neither, as a layer on ui/tree (its
// `layer` prop). The arithmetic is lib/tree-marks-core (what a row reads as once a
// mark covers its subtree, the lens and its grammar); this is the hands.
//
// MARKS ARE PRESSED LIKE FILES ARE SELECTED, so the hand needs nothing new: a click
// marks ONLY this row (its includes replaced; pressed again while it is the only one,
// none), ⌘-click (Ctrl elsewhere) adds it or takes it back, ⇧-click marks every row
// from the last one pressed, and ⌥-click is the negative, exclude (or take the exclude
// back); ⌘⌥ adds as "or" where the marks are a lens (`onMark` with `orWith`). Excludes
// are only set with ⌥ or the menu, and a click's "only this" never clears them. On the
// mark glyph always; on the whole row with `pressRows` (the rows ARE marks: a filter
// rail). The keyboard presses the same: Space or Enter with the same modifiers, `+`
// include, `-` exclude, Backspace or 0 clears. The row's menu starts with include and
// exclude, their presses taught beside them, and "or" with what it joins.
//
//   const layer = useTreeMarks({ marks, onMarksChange: setMarks, pressRows: true })
//   <Tree items={tags} label={(t) => t.name} layer={layer} />

import { Minus, Plus, Split } from "lucide-react"
import * as React from "react"
import { cn } from "@/lib/utils"
import {
  between,
  descendants,
  type Nodelike,
} from "@/registry/base-nova/lib/tree-core"
import {
  type Mark,
  type MarkChange,
  type MarkEntry,
  markStates,
  setMark,
  toggleMark,
} from "@/registry/base-nova/lib/tree-marks-core"
import type { MenuEntry } from "@/registry/base-nova/ui/menu"
import type {
  TreeApi,
  TreeLayer,
  TreeRowState,
} from "@/registry/base-nova/ui/tree"

export interface TreeMarksOptions {
  /** tri-state marks by id (include, exclude, absent) */
  marks: Readonly<Record<string, Mark>>
  onMarksChange?: (marks: Record<string, Mark>) => void
  /** the mark gesture whole, in place of onMarksChange, for marks that are more
   *  than a record (lib/tree-marks-core's Lens, with or-groups): its changes in
   *  order, and `only`: the includes start from none (a plain press, "only this";
   *  excludes stay) before the changes, else the changes go on top of what is marked */
  onMark?: (changes: readonly MarkChange[], only: boolean) => void
  /** what an "or" joins, by name ("people", "beach or mountains"), or null when
   *  nothing is included to join: the menu then offers [OR] with those names, and
   *  ⌘⌥-click joins them. "Or" is only ever offered with what it is or to. */
  orWith?: string | null
  /** a press on the row marks it instead of selecting, for lists whose rows are
   *  marks (a filter rail) */
  pressRows?: boolean
}

/** AND and OR, the lens's two joiners, as one badge of one size: they trade places
 *  under a press without moving anything. OR is lit (it changes what a group means,
 *  "any of these"), AND is the same badge, quiet. */
export function JoinBadge({
  word,
  className,
}: {
  word: "and" | "or"
  className?: string
}) {
  return (
    <span
      data-slot="join-badge"
      data-word={word}
      className={cn(
        "inline-flex h-4 w-7 shrink-0 items-center justify-center rounded-[4px] font-mono font-semibold text-[9px] tracking-wider ring-1 ring-inset",
        word === "or"
          ? "bg-amber-400/15 text-amber-300 ring-amber-400/35"
          : "bg-foreground/6 text-foreground/55 ring-foreground/15",
        className,
      )}
    >
      {word.toUpperCase()}
    </span>
  )
}

export const OrBadge = ({ className }: { className?: string }) => (
  <JoinBadge word="or" className={className} />
)

/** The mark column's glyph: + included, − excluded, a dot for a branch whose subtree
 *  disagrees (tinted by what is under it); neither is nothing at all, a faint + only
 *  under the hand (what a press would do). Inherited or derived marks draw at half
 *  strength. */
function MarkGlyph({
  state,
  own,
  has,
  className,
}: MarkEntry & { className?: string }) {
  if (state === "mixed")
    return (
      <span
        className={cn(
          "size-1.5 rounded-full bg-current",
          has.include && has.exclude
            ? "text-foreground/55"
            : has.include
              ? "text-tone-ok/70"
              : "text-tone-bad/70",
        )}
      />
    )
  return (
    <span
      className={cn(
        "font-mono text-[13px] leading-none",
        state === "include" && "text-tone-ok",
        state === "exclude" && "text-tone-bad",
        state === null &&
          "text-foreground/35 opacity-0 group-hover/row:opacity-100",
        state !== null && !own && "opacity-55",
        className,
      )}
    >
      {state === "exclude" ? "−" : "+"}
    </span>
  )
}

const MARK_WORDS: Record<string, string> = {
  include: "included",
  exclude: "excluded",
  mixed: "mixed",
  null: "not marked",
}

/** The press that excludes: ⌥ alone (⌘ adds, ⌘⌥ adds as "or"). */
const excludes = (e: React.MouseEvent | React.KeyboardEvent) =>
  e.altKey && !e.metaKey && !e.ctrlKey

/** A tree's `layer`, for a tree of any node. */
export type TreeMarksLayer = <N extends Nodelike<N>>(
  tree: TreeApi<N>,
) => TreeLayer<N>

/** The marks as a tree's `layer`: what each row reads as, the mark column, the
 *  presses and keys, the menu's verbs (see the header). Call it where the tree is
 *  drawn, one per tree: it keeps the row a press last marked. */
export function useTreeMarks({
  marks,
  onMarksChange,
  onMark,
  orWith,
  pressRows = false,
}: TreeMarksOptions): TreeMarksLayer {
  /** The row a press last marked: a ⇧-press marks every row from it. */
  const markAnchor = React.useRef<string | null>(null)
  // What every row reads as, computed again only when the forest or the marks change.
  const cache = React.useRef<{
    items: unknown
    childrenOf: unknown
    marks: Readonly<Record<string, Mark>>
    states: Map<string, MarkEntry>
  } | null>(null)
  return <N extends Nodelike<N>>(tree: TreeApi<N>): TreeLayer<N> => {
    const { items, rows, rowAt, index, childrenOf, isDisabled, setActive } =
      tree
    const c = cache.current
    const markOf =
      c && c.items === items && c.childrenOf === childrenOf && c.marks === marks
        ? c.states
        : markStates(items, childrenOf, marks)
    cache.current = { items, childrenOf, marks, states: markOf }

    /** A row's next mark: `mark` toggled against its own (or set, `toggle` false; or
     *  cleared with null), its subtree made whole. `or` joins the last group. */
    const changeOf = (
      n: N,
      mark: Mark | null,
      or: boolean,
      toggle = true,
    ): MarkChange => {
      const state = markOf.get(n.id)?.state ?? null
      return {
        id: n.id,
        mark: mark === null ? null : toggle ? toggleMark(state, mark) : mark,
        or,
        clear: descendants(n, childrenOf),
      }
    }
    /** A gesture's changes, applied at once: `only` starts the includes from none. */
    const apply = (changes: MarkChange[], only: boolean) => {
      if (changes.length === 0 && !only) return
      if (onMark) return onMark(changes, only)
      const start = only
        ? Object.fromEntries(
            Object.entries(marks).filter(([, m]) => m === "exclude"),
          )
        : marks
      onMarksChange?.(
        changes.reduce<Record<string, Mark>>(
          (acc, c) => setMark(acc, c.id, c.mark, c.clear),
          { ...start },
        ),
      )
    }
    const canMark = !!onMarksChange || !!onMark
    /** Toggle one mark on a row (or clear it with null), on top of the rest: the
     *  menu's verbs and the keys. */
    const markRow = (n: N, mark: Mark | null, or = false) => {
      if (!canMark || isDisabled(n)) return
      apply([changeOf(n, mark, or)], false)
    }
    /** What a press means: only this row, ⌘ adds, ⇧ a range, ⌥ excludes, ⌘⌥ adds
     *  as "or". */
    const press = (e: React.MouseEvent | React.KeyboardEvent, n: N) => {
      if (!canMark || isDisabled(n)) return
      const mod = e.metaKey || e.ctrlKey
      const from = markAnchor.current
      if (e.shiftKey && from !== null && rowAt.has(from)) {
        // A range: every row from the last pressed, included, one "or" group (onMark).
        const ids = between(
          rows,
          rowAt.get(from) as number,
          rowAt.get(n.id) as number,
        )
        const range = ids
          .map((id) => (index.get(id) as { node: N }).node)
          .filter((x) => !isDisabled(x))
        apply(
          range.map((x, i) => changeOf(x, "include", !!onMark && i > 0, false)),
          false,
        )
        return
      }
      markAnchor.current = n.id
      if (e.altKey && !mod) return apply([changeOf(n, "exclude", false)], false)
      if (mod)
        return apply(
          [changeOf(n, "include", e.altKey && !!onMark && !!orWith)],
          false,
        )
      // Only this row; pressed again while it is the only one, none.
      const sole =
        markOf.get(n.id)?.state === "include" &&
        Object.entries(marks).every(([id, m]) => m !== "include" || id === n.id)
      apply(sole ? [] : [changeOf(n, "include", false, false)], true)
    }

    return {
      pressRows,
      stateOf: (id) => markOf.get(id)?.state ?? null,
      rowData: (id) => {
        const m = markOf.get(id)
        return {
          "data-mark": m?.state ?? undefined,
          "data-own-mark": (m?.own && m.state) || undefined,
        }
      },
      labelClass:
        "group-data-[mark=exclude]/row:text-foreground/45 group-data-[mark=exclude]/row:line-through",
      column: (n, text) => {
        const m = markOf.get(n.id)
        if (!m) return null
        return (
          <button
            type="button"
            tabIndex={-1}
            data-slot="tree-mark"
            aria-label={`${text}: ${MARK_WORDS[String(m.state)]}`}
            onClick={(e) => {
              e.stopPropagation()
              setActive(n.id)
              press(e, n)
            }}
            onDoubleClick={(e) => e.stopPropagation()}
            className="ml-1.5 grid size-5 shrink-0 place-items-center rounded-sm hover:bg-foreground/10"
          >
            <MarkGlyph state={m.state} own={m.own} has={m.has} />
          </button>
        )
      },
      press,
      key: (e, n) => {
        if (e.key === " ") {
          e.preventDefault()
          press(e, n)
        } else if (e.key === "+" || e.key === "=") {
          e.preventDefault()
          markRow(n, "include")
        } else if (e.key === "-") {
          e.preventDefault()
          markRow(n, "exclude")
        } else if (e.key === "0" || e.key === "Backspace") {
          e.preventDefault()
          markRow(n, null)
        } else return false
        return true
      },
      menu: canMark
        ? (n, state: TreeRowState) => [
            ...markEntries(
              state.mark === "mixed" ? null : (state.mark as Mark | null),
              (next) => markRow(n, next),
            ),
            ...(onMark && orWith && state.mark !== "include"
              ? [
                  {
                    icon: <Split className="text-amber-300" />,
                    label: (
                      <span className="flex min-w-0 items-center gap-1.5">
                        <OrBadge />
                        <span className="truncate text-foreground/70">
                          {orWith}
                        </span>
                      </span>
                    ),
                    hint: "mod+alt+click",
                    onSelect: () => markRow(n, "include", true),
                  } satisfies MenuEntry,
                ]
              : []),
          ]
        : undefined,
    }
  }
}

// ── a row's gestures, for any list of markable rows: the tree's, and a list that is
// not a tree (a timeline's layer heads) wearing the same ones, so a click, a ⌘-click,
// a right-click and the ⋯ mean one thing everywhere.

/** What a press on a row's own mark means, for a row marked alone (no list to be
 *  "only" in): ⌥ excludes (or takes the exclude back), any other press includes (or
 *  takes the include back), as the tree's ⌥ and ⌘ do. */
export function pressMark(
  e: React.MouseEvent | React.KeyboardEvent,
  state: Mark | null,
): Mark | null {
  return toggleMark(state, excludes(e) ? "exclude" : "include")
}

/** The mark verbs a row menu starts with, as the tree words them: exclude (or include
 *  again), with the press that does each drawn on the right. */
export function markEntries(
  state: Mark | null,
  mark: (next: Mark | null) => void,
): MenuEntry[] {
  return [
    state === "include"
      ? {
          icon: <Minus />,
          label: "leave out",
          hint: "mod+click",
          onSelect: () => mark(null),
        }
      : {
          icon: <Plus className="text-tone-ok" />,
          label: state === "exclude" ? "include again" : "include",
          hint: "mod+click",
          onSelect: () => mark("include"),
        },
    state === "exclude"
      ? {
          icon: <Minus />,
          label: "stop excluding",
          hint: "alt+click",
          onSelect: () => mark(null),
        }
      : {
          icon: <Minus className="text-tone-bad" />,
          label: "exclude",
          hint: "alt+click",
          onSelect: () => mark("exclude"),
        },
  ]
}

/** A mark as a row draws it: + included, − excluded, a faint + under the hand when
 *  neither. */
export function MarkSign({
  state,
  className,
}: {
  state: Mark | null
  className?: string
}) {
  return (
    <MarkGlyph
      state={state}
      own
      has={{ include: state === "include", exclude: state === "exclude" }}
      className={className}
    />
  )
}

// MARKS ON A TREE, framework-free: a tri-state mark per node (include, exclude,
// neither) and what each row reads as once marks cover their subtrees, and the LENS,
// marks with a grammar (groups ANDed, OR inside a group, excludes after), the way a
// filter rail or a tag query reads them. ui/tree-marks draws them on ui/tree; a
// server or a test can use them alone. Built on lib/tree-core's nodes and children.

import type { Children, Nodelike } from "@/registry/base-nova/lib/tree-core"

// ── tri-state marks (include, exclude, neither)

export type Mark = "include" | "exclude"
/** What a row reads as: one mark across its whole subtree, nothing, or mixed. */
export type MarkState = Mark | "mixed" | null

/** One mark a gesture changes: the row, its next mark (null: none), whether it joins
 *  the last group as "or" (a Lens's), and its descendants, cleared. */
export interface MarkChange {
  id: string
  mark: Mark | null
  or: boolean
  clear: readonly string[]
}

export interface MarkEntry {
  state: MarkState
  /** set on the node itself, not inherited from a branch or derived from children */
  own: boolean
  /** which marks the subtree holds: what a mixed branch is mixed of */
  has: { include: boolean; exclude: boolean }
}

const INCLUDE = 1
const EXCLUDE = 2
const NONE = 4

/** Every node's mark as a row draws it. A mark covers its descendants (a branch
 *  included includes what is under it) and EXCLUSION WINS (nothing under an excluded
 *  branch reads as included, the lab's rule). An unmarked branch reads as its
 *  children, and any branch whose subtree does not agree with itself reads `mixed`.
 *  `own` says the mark was set on the node itself, not inherited or derived. */
export function markStates<N extends Nodelike<N>>(
  roots: readonly N[],
  children: Children<N>,
  marks: Readonly<Record<string, Mark>>,
): Map<string, MarkEntry> {
  const out = new Map<string, MarkEntry>()
  const walk = (n: N, inherited: Mark | null): number => {
    const mine = marks[n.id]
    const eff: Mark | null =
      inherited === "exclude" ? "exclude" : (mine ?? inherited)
    const kids = children(n) ?? []
    // An unmarked branch reads as its children: all of them excluded is excluded.
    let mask =
      eff === "include"
        ? INCLUDE
        : eff === "exclude"
          ? EXCLUDE
          : kids.length
            ? 0
            : NONE
    for (const k of kids) mask |= walk(k, eff)
    const state: MarkState =
      mask === INCLUDE
        ? "include"
        : mask === EXCLUDE
          ? "exclude"
          : mask === NONE
            ? null
            : "mixed"
    out.set(n.id, {
      state,
      own: mine !== undefined,
      has: { include: (mask & INCLUDE) > 0, exclude: (mask & EXCLUDE) > 0 },
    })
    return mask
  }
  for (const r of roots) walk(r, null)
  return out
}

/** A mark set on a node, its descendants' own marks cleared so the subtree reads as
 *  one thing (a click on a mixed branch makes it whole). null clears the node's own. */
export function setMark(
  marks: Readonly<Record<string, Mark>>,
  id: string,
  mark: Mark | null,
  clear: readonly string[],
): Record<string, Mark> {
  const next = { ...marks }
  for (const d of clear) delete next[d]
  if (mark) next[id] = mark
  else delete next[id]
  return next
}

/** A mark gesture on one row is a TOGGLE of one mark, never a cycle: include
 *  (neither and include trade places, excluded becomes included), or
 *  exclude the same way. A mixed row takes the mark whole. */
export const toggleMark = (state: MarkState, mark: Mark): Mark | null =>
  state === mark ? null : mark

// ── the lens: marks with a grammar (TAGS.md's)

/** A query over marked ids. `groups` are ANDed, the ids inside one group ORed (OR
 *  binds tighter: `people (beach | mountains)`), and every `exclude` is AND NOT after
 *  them. Ordered: the order a person built it in is the order it reads. */
export interface Lens {
  groups: readonly (readonly string[])[]
  exclude: readonly string[]
}

export const NO_LENS: Lens = { groups: [], exclude: [] }

/** Every included id, in reading order. */
export const lensTerms = (lens: Lens): string[] => lens.groups.flat()

export const lensEmpty = (lens: Lens): boolean =>
  lens.groups.length === 0 && lens.exclude.length === 0

/** The lens as plain marks, what a tree draws. */
export function lensMarks(lens: Lens): Record<string, Mark> {
  const out: Record<string, Mark> = {}
  for (const id of lensTerms(lens)) out[id] = "include"
  for (const id of lens.exclude) out[id] = "exclude"
  return out
}

/** Plain marks as a lens: each include its own group (all ANDed), in `order`. */
export function lensOf(
  order: readonly string[],
  marks: Readonly<Record<string, Mark>>,
): Lens {
  return {
    groups: order.filter((id) => marks[id] === "include").map((id) => [id]),
    exclude: order.filter((id) => marks[id] === "exclude"),
  }
}

/** A mark gesture on the lens: `id` and the `clear`ed ids (its descendants) leave
 *  wherever they were, then `id` comes back included, as its own ANDed group or, with
 *  `or`, ORed into the last group; or excluded at the end; or not at all. */
export function markLens(
  lens: Lens,
  id: string,
  mark: Mark | null,
  or: boolean,
  clear: readonly string[] = [],
): Lens {
  const gone = new Set([id, ...clear])
  const groups = lens.groups
    .map((g) => g.filter((t) => !gone.has(t)))
    .filter((g) => g.length > 0)
  const exclude = lens.exclude.filter((t) => !gone.has(t))
  const last = groups.at(-1)
  if (mark === "include")
    return {
      groups:
        or && last
          ? [...groups.slice(0, -1), [...last, id]]
          : [...groups, [id]],
      exclude,
    }
  if (mark === "exclude") return { groups, exclude: [...exclude, id] }
  return { groups, exclude }
}

/** Flip the joiner between included term `at` and the next, in reading order: an
 *  "or" inside a group splits it there into two ANDed groups, an "and" between two
 *  groups merges them into one. */
export function flipJoin(lens: Lens, at: number): Lens {
  const { groups } = lens
  let start = 0
  for (let k = 0; k < groups.length; k++) {
    const g = groups[k] as readonly string[]
    const end = start + g.length - 1
    if (at >= start && at < end) {
      const cut = at - start + 1
      return {
        ...lens,
        groups: [
          ...groups.slice(0, k),
          g.slice(0, cut),
          g.slice(cut),
          ...groups.slice(k + 1),
        ],
      }
    }
    if (at === end) {
      const next = groups[k + 1]
      if (!next) throw new Error(`flipJoin: no term after ${at}`)
      return {
        ...lens,
        groups: [
          ...groups.slice(0, k),
          [...g, ...next],
          ...groups.slice(k + 2),
        ],
      }
    }
    start = end + 1
  }
  throw new Error(`flipJoin: no term ${at}`)
}

/** The lens in its text form: `people (beach | mountains) -color/red`. */
export function lensText(lens: Lens): string {
  return [
    ...lens.groups.map((g) => (g.length > 1 ? `(${g.join(" | ")})` : g[0])),
    ...lens.exclude.map((id) => `-${id}`),
  ].join(" ")
}