archify-vscode-ext/src/model/format.ts

import { LineCounter, isMap, isScalar, parse as parseYaml, parseDocument } from 'yaml';

/**
 * The model file formats (see docs/model.md): `main.archify` with a `model`
 * block, `*.archify-registry` files with shared elements, and `ref`
 * components in diagrams. Pure functions, no VS Code imports, so they run
 * in unit tests and can be reused by other tools.
 */

export const MAIN_FILE = 'main.archify';
export const REGISTRY_EXT = '.archify-registry';

/** Component fields copied from a registry element into every diagram that refers to it. */
export const COPIED_FIELDS = ['type', 'label', 'sublabel', 'tag', 'icon', 'brand', 'sources'] as const;
/** The subset a diagram may override on a `ref` component. */
export const OVERRIDABLE_FIELDS = ['label', 'sublabel', 'tag'] as const;

export interface ModelMeta {
  id?: string;
  name?: string;
  description?: string;
}

export interface RegistryElement {
  id: string;
  type?: string;
  label?: string;
  sublabel?: string;
  tag?: string;
  icon?: unknown;
  brand?: string;
  sources?: unknown;
  description?: string;
  /** Diagram that explains the element, relative to the model root. */
  diagram?: string;
  /** Registry file that defines it, and the 0-based line of its key. */
  file: string;
  line: number;
}

export interface ModelProblem {
  message: string;
  severity: 'error' | 'warning';
  /** File the problem is about; defaults to the diagram being resolved. */
  file?: string;
  /** JSON pointer into that file's document, when it is a diagram. */
  pointer?: string;
  /** 0-based line, when known (registry files). */
  line?: number;
}

export interface Registry {
  elements: Map<string, RegistryElement>;
  files: string[];
  problems: ModelProblem[];
}

export function isMainFile(fileName: string): boolean {
  return (fileName.split(/[\\/]/).pop() ?? '') === MAIN_FILE;
}

export function isRegistryFile(fileName: string): boolean {
  return fileName.endsWith(REGISTRY_EXT);
}

type Obj = Record<string, unknown>;
const isObj = (v: unknown): v is Obj => v !== null && typeof v === 'object' && !Array.isArray(v);
const str = (v: unknown): string | undefined => (typeof v === 'string' && v.length > 0 ? v : undefined);

/** The `model` block of a `main.archify` document, if present. */
export function readModelMeta(doc: unknown): ModelMeta | undefined {
  if (!isObj(doc) || !isObj(doc.model)) return undefined;
  const { id, name, description } = doc.model;
  return { id: str(id), name: str(name), description: str(description) };
}

/** True when the document also carries a diagram (a `diagram_type`). */
export function hasDiagramBody(doc: unknown): boolean {
  return isObj(doc) && typeof doc.diagram_type === 'string';
}

/** Parses one registry file. Line numbers come from the YAML document. */
export function parseRegistry(text: string, file: string): { elements: RegistryElement[]; problems: ModelProblem[] } {
  const problems: ModelProblem[] = [];
  const elements: RegistryElement[] = [];
  // prettyErrors would put line numbers into the messages; the line counter gives them without.
  const lines = new LineCounter();
  const doc = parseDocument(text, { prettyErrors: false, uniqueKeys: false, lineCounter: lines });
  for (const error of doc.errors) {
    problems.push({ message: `Registry is not valid YAML: ${error.message.split('\n')[0]}`, severity: 'error', file, line: lines.linePos(error.pos[0]).line - 1 });
  }
  if (doc.errors.length) return { elements, problems };
  const lineOf = (offset: number | undefined) => (offset === undefined ? 0 : text.slice(0, offset).split('\n').length - 1);
  const root = doc.contents;
  if (!isMap(root)) {
    problems.push({ message: 'A registry must be a mapping with an "elements" key.', severity: 'error', file, line: 0 });
    return { elements, problems };
  }
  const elementsNode = root.get('elements', true);
  if (!isMap(elementsNode)) {
    problems.push({ message: 'A registry needs an "elements" mapping (id → element).', severity: 'error', file, line: 0 });
    return { elements, problems };
  }
  for (const pair of elementsNode.items) {
    const key = pair.key;
    const id = isScalar(key) ? String(key.value) : undefined;
    const line = isScalar(key) ? lineOf(key.range?.[0]) : 0;
    if (!id) continue;
    const value = isMap(pair.value) ? (pair.value.toJSON() as Obj) : undefined;
    if (!value) {
      problems.push({ message: `Element "${id}" must be a mapping (type, label, description, …).`, severity: 'error', file, line });
      continue;
    }
    elements.push({
      id,
      type: str(value.type),
      label: str(value.label),
      sublabel: str(value.sublabel),
      tag: str(value.tag),
      icon: value.icon,
      brand: str(value.brand),
      sources: value.sources,
      description: str(value.description),
      diagram: str(value.diagram),
      file,
      line,
    });
  }
  return { elements, problems };
}

/** Merges parsed registry files; an id defined twice is an error on the second definition. */
export function mergeRegistries(parts: { elements: RegistryElement[]; problems: ModelProblem[]; file: string }[]): Registry {
  const elements = new Map<string, RegistryElement>();
  const problems: ModelProblem[] = [];
  for (const part of parts) {
    problems.push(...part.problems);
    for (const element of part.elements) {
      const existing = elements.get(element.id);
      if (existing) {
        problems.push({
          message: `Element "${element.id}" is defined twice (also in ${existing.file}, line ${existing.line + 1}).`,
          severity: 'error',
          file: element.file,
          line: element.line,
        });
        continue;
      }
      elements.set(element.id, element);
    }
  }
  return { elements, files: parts.map((p) => p.file), problems };
}

/** Where a diagram attaches a decision: the diagram itself, a node or an edge. */
export type DecisionAttachment = { kind: 'diagram' } | { kind: 'node'; id: string } | { kind: 'edge'; from: string; to: string; label?: string };

export interface DecisionRef {
  id: string;
  pointer: string;
  attachedTo: DecisionAttachment;
}

/** Collections whose items may carry `decisions` (and `diagram` links), per diagram kind. */
const NODE_COLLECTIONS = ['components', 'nodes', 'participants', 'states'] as const;
const EDGE_COLLECTIONS = ['connections', 'edges', 'flows', 'transitions', 'messages'] as const;

export interface ResolvedDiagram {
  /** The diagram as Archify expects it: refs expanded, model block and links removed. */
  diagram: Obj;
  problems: ModelProblem[];
  /** Component index → the registry element it came from. */
  refs: Map<number, RegistryElement>;
  /** Node id → diagram path (relative to the model root) it links to. */
  links: Map<string, string>;
  /** Registry element ids used by this diagram. */
  used: Set<string>;
  /** Node id → registry description, for the Elements panel. */
  descriptions: Map<string, string>;
  /** Decisions the diagram points at, with where they are attached. */
  decisionRefs: DecisionRef[];
}

export interface ResolveOptions {
  /** Undefined when the diagram is not inside a model. */
  registry?: Registry;
  /** Known decision ids of the model; undefined outside a model or when unknown. */
  decisions?: Set<string>;
  fileName: string;
}

/** Strips `decisions` keys from a diagram and collects what they point at. Mutates `diagram`. */
function takeDecisionRefs(diagram: Obj, options: ResolveOptions, problems: ModelProblem[]): DecisionRef[] {
  const refs: DecisionRef[] = [];
  const take = (owner: Obj, pointer: string, attachedTo: DecisionAttachment) => {
    if (!('decisions' in owner)) return;
    const raw = owner.decisions;
    delete owner.decisions;
    const ids = Array.isArray(raw) ? raw : [raw];
    ids.forEach((value, i) => {
      const id = value === undefined || value === null ? undefined : String(value);
      if (!id) return;
      const at = `${pointer}/decisions/${i}`;
      if (typeof value === 'number') {
        problems.push({ message: `Decision id ${value} is a number to YAML (0001 becomes 1). Quote it ("${value}") or use words.`, severity: 'error', pointer: at });
        return;
      }
      if (!options.registry) {
        problems.push({ message: `Decision "${id}" cannot be resolved: ${options.fileName.split(/[\\/]/).pop()} is not inside a model (no ${MAIN_FILE} above it).`, severity: 'error', pointer: at });
      } else if (options.decisions && !options.decisions.has(id)) {
        problems.push({ message: `Unknown decision "${id}": no .archify-decision file of this model defines it.`, severity: 'error', pointer: at });
      }
      refs.push({ id, pointer: at, attachedTo });
    });
  };
  take(diagram, '', { kind: 'diagram' });
  for (const key of NODE_COLLECTIONS) {
    const items = diagram[key];
    if (!Array.isArray(items)) continue;
    diagram[key] = items.map((item, index) => {
      if (!isObj(item)) return item;
      const copy = { ...item };
      take(copy, `/${key}/${index}`, { kind: 'node', id: str(copy.id) ?? str(copy.ref) ?? String(index) });
      return copy;
    });
  }
  for (const key of EDGE_COLLECTIONS) {
    const items = diagram[key];
    if (!Array.isArray(items)) continue;
    diagram[key] = items.map((item, index) => {
      if (!isObj(item)) return item;
      const copy = { ...item };
      take(copy, `/${key}/${index}`, { kind: 'edge', from: str(copy.from) ?? '', to: str(copy.to) ?? '', label: str(copy.label) });
      return copy;
    });
  }
  return refs;
}

/**
 * Expands `ref` components from the registry and strips what Archify does
 * not know. Component order and indexes are kept, so Archify's diagnostics
 * (JSON pointers) still point at the right component in the source file.
 */
export function resolveRefs(doc: unknown, options: ResolveOptions): ResolvedDiagram {
  const problems: ModelProblem[] = [];
  const refs = new Map<number, RegistryElement>();
  const links = new Map<string, string>();
  const used = new Set<string>();
  const descriptions = new Map<string, string>();
  if (!isObj(doc)) return { diagram: {}, problems, refs, links, used, descriptions, decisionRefs: [] };
  const diagram: Obj = { ...doc };
  delete diagram.model;
  const decisionRefs = takeDecisionRefs(diagram, options, problems);
  const { registry } = options;
  const components = Array.isArray(diagram.components) ? diagram.components : undefined;
  if (!components) return { diagram, problems, refs, links, used, descriptions, decisionRefs };

  diagram.components = components.map((raw, index) => {
    if (!isObj(raw)) return raw;
    const ref = str(raw.ref);
    const linkFromComponent = str(raw.diagram);
    if (!ref) {
      const { diagram: _link, ...plain } = raw;
      const id = str(plain.id);
      if (id && linkFromComponent) links.set(id, linkFromComponent);
      return plain;
    }
    const { ref: _ref, diagram: _link, ...own } = raw;
    const element = registry?.elements.get(ref);
    const id = str(own.id) ?? ref;
    if (!element) {
      problems.push({
        message: registry
          ? `Unknown element "${ref}": no registry of this model defines it.`
          : `"${ref}" cannot be resolved: ${options.fileName.split(/[\\/]/).pop()} is not inside a model (no ${MAIN_FILE} above it).`,
        severity: 'error',
        pointer: `/components/${index}/ref`,
      });
      return { type: 'external', label: ref, ...own, id };
    }
    used.add(ref);
    refs.set(index, element);
    if (element.description) descriptions.set(id, element.description.trim());
    for (const key of Object.keys(own)) {
      if ((COPIED_FIELDS as readonly string[]).includes(key) && !(OVERRIDABLE_FIELDS as readonly string[]).includes(key)) {
        problems.push({
          message: `"${key}" comes from the registry element "${ref}" and cannot be changed in a diagram.`,
          severity: 'error',
          pointer: `/components/${index}/${key}`,
        });
      }
    }
    const merged: Obj = { id };
    for (const key of COPIED_FIELDS) if (element[key] !== undefined) merged[key] = element[key];
    for (const [key, value] of Object.entries(own)) {
      if (key === 'id') continue;
      if (!(COPIED_FIELDS as readonly string[]).includes(key) || (OVERRIDABLE_FIELDS as readonly string[]).includes(key)) merged[key] = value;
    }
    const link = linkFromComponent ?? element.diagram;
    if (link) links.set(id, link);
    return merged;
  });
  return { diagram, problems, refs, links, used, descriptions, decisionRefs };
}

/** Group of a diagram inside a model: '' for the root folder, else the folder path with '/' separators. */
export function groupOf(relativePath: string): string {
  const parts = relativePath.split(/[\\/]/);
  parts.pop();
  return parts.join('/');
}

/** A node as a diagram declares it: a `ref` to the registry, or an inline element. */
export interface DiagramNode {
  id: string;
  ref?: string;
  type?: string;
  label?: string;
  sublabel?: string;
  brand?: string;
}

/** A connection, edge, flow, transition or message between two nodes of a diagram. */
export interface DiagramEdge {
  from: string;
  to: string;
  label?: string;
}

export interface DiagramSummary {
  relativePath: string;
  group: string;
  kind?: string;
  title?: string;
  used: string[];
  /** Decisions the diagram points at. */
  decisions: DecisionRef[];
  /** Every node the diagram draws, shared or inline, in file order. */
  nodes: DiagramNode[];
  /** Every edge between two of those nodes. */
  edges: DiagramEdge[];
  parseError?: string;
}

/** What the model page needs to know about one diagram file. */
export function summarizeDiagram(text: string, relativePath: string): DiagramSummary {
  let doc: unknown;
  try {
    doc = parseYaml(text, { prettyErrors: false, uniqueKeys: false });
  } catch (error) {
    return { relativePath, group: groupOf(relativePath), used: [], decisions: [], nodes: [], edges: [], parseError: error instanceof Error ? error.message.split('\n')[0] : String(error) };
  }
  const obj = isObj(doc) ? doc : {};
  const meta = isObj(obj.meta) ? obj.meta : {};
  const components = Array.isArray(obj.components) ? obj.components : [];
  const used = [...new Set(components.map((c) => (isObj(c) ? str(c.ref) : undefined)).filter((r): r is string => !!r))];
  const decisions = takeDecisionRefs({ ...obj }, { registry: { elements: new Map(), files: [], problems: [] }, fileName: relativePath }, []);
  const nodes: DiagramNode[] = [];
  const seen = new Set<string>();
  for (const key of NODE_COLLECTIONS) {
    const items = obj[key];
    if (!Array.isArray(items)) continue;
    for (const item of items) {
      if (!isObj(item)) continue;
      const ref = str(item.ref);
      const id = str(item.id) ?? ref;
      if (!id || seen.has(id)) continue;
      seen.add(id);
      nodes.push({ id, ref, type: str(item.type), label: str(item.label), sublabel: str(item.sublabel), brand: str(item.brand) });
    }
  }
  const edges: DiagramEdge[] = [];
  for (const key of EDGE_COLLECTIONS) {
    const items = obj[key];
    if (!Array.isArray(items)) continue;
    for (const item of items) {
      if (!isObj(item)) continue;
      const from = str(item.from);
      const to = str(item.to);
      if (!from || !to || !seen.has(from) || !seen.has(to)) continue;
      edges.push({ from, to, label: str(item.label) });
    }
  }
  return { relativePath, group: groupOf(relativePath), kind: str(obj.diagram_type), title: str(meta.title), used, decisions, nodes, edges };
}

/** Groups in display order: the root group first, then folders in path order. */
export function sortedGroups(summaries: DiagramSummary[]): { group: string; diagrams: DiagramSummary[] }[] {
  const byGroup = new Map<string, DiagramSummary[]>();
  for (const s of summaries) byGroup.set(s.group, [...(byGroup.get(s.group) ?? []), s]);
  const groups = [...byGroup.keys()].sort((a, b) => (a === '' ? -1 : b === '' ? 1 : a.localeCompare(b)));
  return groups.map((group) => ({ group, diagrams: byGroup.get(group)!.sort((a, b) => a.relativePath.localeCompare(b.relativePath)) }));
}