import { useCallback, useEffect, useState } from "react";
import type { ChunkRow } from "../types";
import { chunkColor, effectiveScheme, tintStyle, type ColorScheme } from "../lib/colorScheme";

// DocumentChunkTree renders a document's chunk hierarchy (RFC AK), built from
// the flat query_chunks list by parent_id + position. Mirrors PathTree: a single
// expand Map owned by the tree, default-expanded, click-to-select.

export interface ChunkNode {
  row: ChunkRow;
  children: ChunkNode[];
}

// buildChunkTree groups chunks into parent → children (position-ordered). A
// chunk whose parent_id is absent or not in the set is a root (the document's
// root chunk; defensively also any orphan).
export function buildChunkTree(chunks: ChunkRow[]): ChunkNode[] {
  const byId = new Map<string, ChunkNode>();
  chunks.forEach((c) => byId.set(c.id, { row: c, children: [] }));
  const roots: ChunkNode[] = [];
  chunks.forEach((c) => {
    const node = byId.get(c.id)!;
    if (c.parent_id && byId.has(c.parent_id)) byId.get(c.parent_id)!.children.push(node);
    else roots.push(node);
  });
  const sortRec = (nodes: ChunkNode[]) => {
    nodes.sort((a, b) => a.row.position - b.row.position);
    nodes.forEach((n) => sortRec(n.children));
  };
  sortRec(roots);
  return roots;
}

// typeIcon returns a small glyph for a media chunk type (RFC BO) so the tree
// distinguishes image / diagram chunks at a glance; "" for a plain text chunk.
function typeIcon(type: string | undefined): string {
  if (type === "image") return "🖼";
  if (type === "mermaid") return "📊";
  return "";
}

// findChunkNode locates a node by id within a chunk forest.
export function findChunkNode(nodes: ChunkNode[], id: string): ChunkNode | undefined {
  for (const n of nodes) {
    if (n.row.id === id) return n;
    const hit = findChunkNode(n.children, id);
    if (hit) return hit;
  }
  return undefined;
}

export interface DocumentChunkTreeProps {
  tree: ChunkNode[];
  selectedId?: string;
  onSelect: (node: ChunkNode) => void;
  // RFC BN: when colorEnabled, chunk rows are tinted by status via `scheme`
  // (falls back to the default palette). Unset → neutral rows (prior behavior).
  colorEnabled?: boolean;
  scheme?: ColorScheme;
  // RFC BN P4: chunk ids that have at least one cross-reference edge — those
  // tiles get a 🔗 badge. Unset → no badges.
  linkedIds?: Set<string>;
}

export default function DocumentChunkTree({
  tree,
  selectedId,
  onSelect,
  colorEnabled,
  scheme,
  linkedIds,
}: DocumentChunkTreeProps) {
  const [expanded, setExpanded] = useState<Map<string, boolean>>(() => new Map());
  const toggle = useCallback((id: string) => {
    setExpanded((prev) => {
      const next = new Map(prev);
      next.set(id, prev.get(id) === false ? true : false);
      return next;
    });
  }, []);

  // Auto-expand ancestors of the selection.
  useEffect(() => {
    if (!selectedId) return;
    const parents = ancestorIds(tree, selectedId);
    if (parents.length === 0) return;
    setExpanded((prev) => {
      let mutated = false;
      const next = new Map(prev);
      for (const id of parents) {
        if (next.get(id) === false) {
          next.set(id, true);
          mutated = true;
        }
      }
      return mutated ? next : prev;
    });
  }, [tree, selectedId]);

  // Resolve the palette once (null → no tinting) and thread it down.
  const tintScheme = colorEnabled ? effectiveScheme(scheme) : null;

  if (tree.length === 0) {
    return (
      <div className="empty">
        <p>No chunks.</p>
      </div>
    );
  }
  return (
    <ul className="tree chunk-tree">
      {tree.map((n) => (
        <ChunkTreeNode
          key={n.row.id}
          node={n}
          expanded={expanded}
          toggle={toggle}
          selectedId={selectedId}
          onSelect={onSelect}
          tintScheme={tintScheme}
          linkedIds={linkedIds}
        />
      ))}
    </ul>
  );
}

function ancestorIds(roots: ChunkNode[], id: string): string[] {
  const path: string[] = [];
  const walk = (nodes: ChunkNode[], parents: string[]): boolean => {
    for (const n of nodes) {
      if (n.row.id === id) {
        path.push(...parents);
        return true;
      }
      if (walk(n.children, [...parents, n.row.id])) return true;
    }
    return false;
  };
  walk(roots, []);
  return path;
}

interface NodeProps {
  node: ChunkNode;
  expanded: Map<string, boolean>;
  toggle: (id: string) => void;
  selectedId?: string;
  onSelect: (node: ChunkNode) => void;
  tintScheme: ColorScheme | null;
  linkedIds?: Set<string>;
}

function ChunkTreeNode({ node, expanded, toggle, selectedId, onSelect, tintScheme, linkedIds }: NodeProps) {
  const hasChildren = node.children.length > 0;
  const isOpen = expanded.get(node.row.id) !== false; // default-expanded
  const isSelected = selectedId === node.row.id;
  const rowStyle = tintScheme ? tintStyle(chunkColor(node.row.status, tintScheme)) : undefined;
  const hasLink = linkedIds?.has(node.row.id);
  return (
    <li className={`node chunk-node ${isSelected ? "selected" : ""}`}>
      <div className="row chunk-row" style={rowStyle}>
        <button
          type="button"
          className="tree-caret"
          aria-label={isOpen ? "collapse" : "expand"}
          disabled={!hasChildren}
          onClick={(e) => {
            e.stopPropagation();
            if (hasChildren) toggle(node.row.id);
          }}
        >
          {hasChildren ? (isOpen ? "▼" : "▶") : "·"}
        </button>
        <button
          type="button"
          className="chunk-link"
          onClick={() => onSelect(node)}
          title={node.row.title}
        >
          {typeIcon(node.row.type) && (
            <span className="chunk-type-icon" aria-hidden>
              {typeIcon(node.row.type)}
            </span>
          )}
          <span className="chunk-title">{node.row.title || "(untitled)"}</span>
          {node.row.type && <span className="chunk-badge">{node.row.type}</span>}
          {node.row.status && <span className="chunk-badge chunk-status">{node.row.status}</span>}
          {hasLink && (
            <span className="chunk-link-badge" title="Has cross-reference links" aria-label="linked">
              🔗
            </span>
          )}
        </button>
      </div>
      {hasChildren && isOpen && (
        <ul className="children chunk-children">
          {node.children.map((c) => (
            <ChunkTreeNode
              key={c.row.id}
              node={c}
              expanded={expanded}
              toggle={toggle}
              selectedId={selectedId}
              onSelect={onSelect}
              tintScheme={tintScheme}
              linkedIds={linkedIds}
            />
          ))}
        </ul>
      )}
    </li>
  );
}
