webapp/lib/tree/parser.ts

import { layoutTreeNodes } from "@/lib/tree/layout";
import type {
  ParseBlock,
  ParseResult,
  TrajectoryJsonNode,
  TrajectoryTreeNode
} from "@/lib/tree/types";
import { validateTreeNodes } from "@/lib/tree/validate";

const startTokens = ["<|trajectory|>", "|trajectory|", "<trajectory>"];
const endTokens = ["<|end_of_trajectory|>", "</|end_of_trajectory|>", "</trajectory>", "</|trajectory|>"];
const beginTrajectoryTokens = ["<|begin_of_trajectory|>"];
const beginThoughtToken = "<|begin_of_thought|>";
const endThoughtTokens = ["<|end_of_thought|>", "</|end_of_thought>", "</|end_of_thought|>", "</end_of_thought>"];

function isRecord(value: unknown): value is Record<string, unknown> {
  return Boolean(value) && typeof value === "object" && !Array.isArray(value);
}

function asString(value: unknown, fallback = "") {
  return typeof value === "string" ? value : fallback;
}

function asStringArray(value: unknown) {
  return Array.isArray(value) ? value.map(String) : [];
}

function asRecord(value: unknown) {
  return isRecord(value) ? value : {};
}

function asArray(value: unknown) {
  return Array.isArray(value) ? value : [];
}

function firstString(...values: unknown[]) {
  for (const value of values) {
    if (typeof value === "string" && value.trim()) return value.trim();
    if (Array.isArray(value)) {
      for (const item of value) {
        if (typeof item === "string" && item.trim()) return item.trim();
      }
    }
  }
  return "";
}

function compactJson(value: unknown) {
  if (!isRecord(value) || Object.keys(value).length === 0) return "";
  try {
    return JSON.stringify(value);
  } catch {
    return String(value);
  }
}

function inferGoal(item: Record<string, unknown>) {
  const inputs = asRecord(item.inputs);
  const output = asRecord(item.output);
  const direct = firstString(item.goal);
  if (direct) return direct;

  const outputText = firstString(
    output.goal,
    output.plan,
    output.interpretation,
    output.derived_value,
    output.verified_result,
    output.answer,
    output.value,
    output.claims,
    output.note,
    output.status
  );
  const inputText = firstString(
    inputs.task,
    inputs.expression,
    inputs.equation,
    inputs.verification,
    inputs.answer,
    inputs.problem,
    inputs.operation,
    inputs.method
  );

  if (outputText && inputText) return `${inputText} -> ${outputText}`;
  if (outputText) return outputText;
  if (inputText) return inputText;

  return compactJson(output) || asString(item.kind, "step");
}

function normalizeJsonNode(value: unknown): TrajectoryJsonNode {
  const item = asRecord(value);
  const inputs = asRecord(item.inputs);
  const stateDelta = asRecord(item.state_delta);
  const output = asRecord(item.output);
  const retrieval = asRecord(item.retrieval);
  const verification = asRecord(item.verification);
  const branch = asRecord(item.branch);
  const stateSource = Object.keys(stateDelta).length > 0 ? stateDelta : output;

  return {
    ...item,
    v: Number(item.v ?? 1),
    id: asString(item.id),
    parent: item.parent === null || item.parent === undefined || item.parent === "" ? null : String(item.parent),
    kind: asString(item.kind, "step"),
    goal: inferGoal(item),
    skill: asStringArray(item.skill),
    inputs: {
      ...inputs,
      problem_refs: asStringArray(inputs.problem_refs),
      node_refs: asStringArray(inputs.node_refs)
    },
    state_delta: {
      ...stateDelta,
      claims: asArray(stateSource.claims),
      equations: asArray(stateSource.equations),
      constraints: asArray("constraints" in stateSource ? stateSource.constraints : item.constraints),
      definitions: asArray("definitions" in stateSource ? stateSource.definitions : item.definitions)
    },
    retrieval: {
      ...retrieval,
      needed: Boolean(retrieval.needed),
      source: asString(retrieval.source, "none"),
      queries: asStringArray(retrieval.queries),
      evidence_refs: asStringArray(retrieval.evidence_refs)
    },
    tool: "tool" in item ? item.tool : null,
    verification: {
      ...verification,
      kind: asString(verification.kind, "none"),
      status: asString(verification.status, "pending"),
      notes: asString(verification.notes)
    },
    branch: {
      ...branch,
      mode: asString(branch.mode, "main"),
      label: asString(branch.label)
    }
  };
}

function nodeFromJson(value: unknown, rawJson: string, reasoning: string): TrajectoryTreeNode {
  const data = normalizeJsonNode(value);
  return {
    id: data.id,
    parent: data.parent,
    kind: data.kind,
    goal: data.goal,
    skill: data.skill,
    data,
    rawJson: JSON.stringify(data, null, 2) || rawJson,
    reasoning: reasoning.trim()
  };
}

function parseTrajectoryJson(rawJson: string) {
  const text = rawJson.trim().replace(/^\|+/, "").trim();
  const repairNotes: string[] = [];
  try {
    const parsed = JSON.parse(text);
    const objectRepair = repairTrajectoryObject(parsed);
    return {
      parsed: objectRepair.value,
      rawJson: objectRepair.repairNotes.length > 0 ? JSON.stringify(objectRepair.value) : text,
      repairNotes: objectRepair.repairNotes
    };
  } catch (error) {
    const fixed = fixJsonText(text);
    repairNotes.push(
      `fixed JSON syntax${error instanceof Error && error.message ? `: ${error.message}` : ""}`
    );
    const parsed = JSON.parse(fixed.text);
    repairNotes.push(...fixed.notes);
    const objectRepair = repairTrajectoryObject(parsed);
    repairNotes.push(...objectRepair.repairNotes);
    return {
      parsed: objectRepair.value,
      rawJson: JSON.stringify(objectRepair.value),
      repairNotes: Array.from(new Set(repairNotes))
    };
  }
}

function fixJsonText(rawJson: string) {
  let text = rawJson.trim().replace(/^\|+/, "").trim();
  const notes: string[] = [];

  const normalizedQuotes = text.replace(/[“”]/g, "\"").replace(/[‘’]/g, "'");
  if (normalizedQuotes !== text) {
    text = normalizedQuotes;
    notes.push("normalized smart quotes");
  }

  const fixedQuotedKeys = text.replace(/([,{]\s*)"([A-Za-z_][A-Za-z0-9_.-]*)\s*:\s*"/g, "$1\"$2\":\"");
  if (fixedQuotedKeys !== text) {
    text = fixedQuotedKeys;
    notes.push("repaired malformed quoted key");
  }

  const fixedUnquotedKeys = text.replace(/([,{]\s*)([A-Za-z_][A-Za-z0-9_.-]*)\s*:/g, "$1\"$2\":");
  if (fixedUnquotedKeys !== text) {
    text = fixedUnquotedKeys;
    notes.push("quoted unquoted object key");
  }

  const withoutTrailingCommas = text.replace(/,\s*([}\]])/g, "$1");
  if (withoutTrailingCommas !== text) {
    text = withoutTrailingCommas;
    notes.push("removed trailing comma");
  }

  const balanced = appendMissingJsonClosers(text);
  if (balanced !== text) {
    text = balanced;
    notes.push("added missing closing delimiter");
  }

  return { text, notes };
}

function appendMissingJsonClosers(text: string) {
  const stack: string[] = [];
  let inString = false;
  let escaped = false;

  for (const char of text) {
    if (inString) {
      if (escaped) {
        escaped = false;
      } else if (char === "\\") {
        escaped = true;
      } else if (char === "\"") {
        inString = false;
      }
      continue;
    }

    if (char === "\"") {
      inString = true;
    } else if (char === "{") {
      stack.push("}");
    } else if (char === "[") {
      stack.push("]");
    } else if (char === "}" || char === "]") {
      if (stack[stack.length - 1] === char) stack.pop();
    }
  }

  if (inString || stack.length === 0) return text;
  return text + stack.reverse().join("");
}

function repairTrajectoryObject(value: unknown) {
  if (!isRecord(value)) return { value, repairNotes: [] };
  const repaired: Record<string, unknown> = { ...value };
  const stateDelta = asRecord(repaired.state_delta);
  if (Object.keys(stateDelta).length === 0) return { value: repaired, repairNotes: [] };

  const nextStateDelta: Record<string, unknown> = { ...stateDelta };
  const repairNotes: string[] = [];
  for (const key of ["retrieval", "tool", "verification", "branch"]) {
    if (!(key in nextStateDelta)) continue;
    if (!(key in repaired)) repaired[key] = nextStateDelta[key];
    delete nextStateDelta[key];
    repairNotes.push(`moved "${key}" out of state_delta`);
  }

  if (repairNotes.length > 0) repaired.state_delta = nextStateDelta;
  return { value: repaired, repairNotes };
}

function titleFromNodes(nodes: TrajectoryTreeNode[]) {
  const root = nodes.find((node) => node.parent === null) ?? nodes[0];
  const title = root?.goal?.trim() || root?.kind?.trim() || "Trajectory Tree";
  return title.length > 90 ? `${title.slice(0, 87)}...` : title;
}

function parseWrapped(rawText: string): ParseBlock[] {
  const blocks: ParseBlock[] = [];
  let cursor = 0;
  let previousContentEnd = 0;
  let index = 0;

  while (cursor < rawText.length) {
    const startMatch = findNextToken(rawText, startTokens, cursor);
    if (!startMatch) break;
    const [, currentReasoning] = splitPrecedingReasoning(rawText, previousContentEnd, startMatch.index);
    const jsonStart = startMatch.index + startMatch.token.length;
    const endMatch = findNextToken(rawText, endTokens, jsonStart);
    const jsonSpan = findJsonObjectSpan(rawText, jsonStart);
    const boundaryMatch = findNextToken(rawText, startTokens, jsonStart);
    if (!endMatch && !jsonSpan && !boundaryMatch) {
      if (blocks.length === 0) {
        blocks.push({
          index,
          rawJson: rawText.slice(jsonStart).trim(),
          reasoning: "",
          node: null,
          error: "Missing trajectory end token."
        });
      }
      break;
    }

    const useJsonSpan =
      Boolean(jsonSpan) &&
      (!endMatch || (jsonSpan!.end <= endMatch.index && (!boundaryMatch || endMatch.index < boundaryMatch.index)));
    const useEndToken = Boolean(endMatch) && (!boundaryMatch || endMatch!.index < boundaryMatch.index);
    const contentEnd = useJsonSpan
      ? jsonSpan!.end
      : useEndToken
        ? endMatch!.index + endMatch!.token.length
        : boundaryMatch!.index;
    const nextStart = findNextToken(rawText, startTokens, contentEnd);
    const rawFragment = rawText
      .slice(jsonStart, useJsonSpan ? jsonSpan!.end : useEndToken ? endMatch!.index : boundaryMatch!.index)
      .trim();
    const rawJson = extractJsonPrefix(rawFragment);
    const reasoningEnd = nextStart?.index ?? rawText.length;
    const trailingReasoning = trailingReasoningForPrevious(rawText.slice(contentEnd, reasoningEnd));
    const reasoning = currentReasoning || trailingReasoning;

    try {
      const parsed = parseTrajectoryJson(rawJson);
      const clean = cleanReasoning(reasoning);
      blocks.push({
        index,
        rawJson: parsed.rawJson,
        reasoning: clean,
        node: nodeFromJson(parsed.parsed, parsed.rawJson, clean),
        repairNotes: parsed.repairNotes
      });
    } catch (error) {
      blocks.push({
        index,
        rawJson,
        reasoning,
        node: null,
        error: error instanceof Error ? error.message : String(error)
      });
    }

    previousContentEnd = contentEnd;
    cursor = reasoningEnd;
    index += 1;
  }

  return blocks;
}

function cleanReasoning(value: string) {
  let cleaned = value.replaceAll(beginThoughtToken, "");
  for (const token of [...endThoughtTokens, ...beginTrajectoryTokens, ...endTokens]) {
    cleaned = cleaned.replaceAll(token, "");
  }
  return cleaned.replaceAll("<|im_end|>", "").replace(/^[\s|]+|[\s|]+$/g, "");
}

function splitPrecedingReasoning(text: string, lowerBound: number, start: number): [string, string] {
  const segment = text.slice(lowerBound, start);
  const begin = segment.lastIndexOf(beginThoughtToken);
  if (begin === -1) return [segment, ""];
  const endMatch = findNextToken(segment, endThoughtTokens, begin + beginThoughtToken.length);
  if (!endMatch) return [segment, ""];

  const afterEnd = segment.slice(endMatch.index + endMatch.token.length);
  if (stripDelimiterNoise(afterEnd)) return [segment, ""];

  return [segment.slice(0, begin), segment.slice(begin + beginThoughtToken.length, endMatch.index)];
}

function trailingReasoningForPrevious(segment: string) {
  const begin = segment.indexOf(beginThoughtToken);
  const trajectoryBegin = findNextToken(segment, beginTrajectoryTokens, 0);
  if (begin !== -1) return segment.slice(0, begin);
  if (trajectoryBegin) return segment.slice(0, trajectoryBegin.index);
  return segment;
}

function findNextToken(text: string, tokens: string[], start: number) {
  let match: { index: number; token: string } | null = null;
  for (const token of tokens) {
    const index = text.indexOf(token, start);
    if (index === -1) continue;
    if (!match || index < match.index) match = { index, token };
  }
  return match;
}

function findJsonObjectSpan(text: string, start: number) {
  let index = start;
  while (index < text.length && (/\s/.test(text[index]) || text[index] === "|")) index += 1;
  if (index >= text.length || text[index] !== "{") return null;

  let depth = 0;
  let inString = false;
  let escaped = false;

  for (let cursor = index; cursor < text.length; cursor += 1) {
    const char = text[cursor];
    if (inString) {
      if (escaped) {
        escaped = false;
      } else if (char === "\\") {
        escaped = true;
      } else if (char === "\"") {
        inString = false;
      }
      continue;
    }

    if (char === "\"") {
      inString = true;
    } else if (char === "{") {
      depth += 1;
    } else if (char === "}") {
      depth -= 1;
      if (depth === 0) return { start: index, end: cursor + 1 };
    }
  }

  return null;
}

function stripDelimiterNoise(value: string) {
  let cleaned = value;
  for (const token of [...beginTrajectoryTokens, ...endTokens]) {
    cleaned = cleaned.replaceAll(token, "");
  }
  return cleaned.replace(/^[\s|]+|[\s|]+$/g, "");
}

function extractJsonPrefix(rawJson: string) {
  const text = rawJson.trim().replace(/^\|+/, "").trim();
  let depth = 0;
  let inString = false;
  let escaped = false;
  let started = false;

  for (let index = 0; index < text.length; index += 1) {
    const char = text[index];
    if (!started) {
      if (char === "{") {
        started = true;
        depth = 1;
      } else if (!/\s/.test(char)) {
        break;
      }
      continue;
    }

    if (inString) {
      if (escaped) {
        escaped = false;
      } else if (char === "\\") {
        escaped = true;
      } else if (char === "\"") {
        inString = false;
      }
      continue;
    }

    if (char === "\"") {
      inString = true;
    } else if (char === "{") {
      depth += 1;
    } else if (char === "}") {
      depth -= 1;
      if (depth === 0) return text.slice(0, index + 1).trim();
    }
  }

  return text;
}

function parseRawJson(rawText: string): ParseBlock[] {
  const parsed = parseTrajectoryJson(rawText).parsed;
  const candidates = Array.isArray(parsed)
    ? parsed
    : isRecord(parsed) && Array.isArray(parsed.nodes)
      ? parsed.nodes
      : [parsed];

  return candidates.map((candidate, index) => {
    const rawJson = JSON.stringify(candidate, null, 2);
    return {
      index,
      rawJson,
      reasoning: "",
      node: nodeFromJson(candidate, rawJson, "")
    };
  });
}

export function parseTrajectoryText(rawText: string): ParseResult {
  const trimmed = rawText.trim();
  let blocks: ParseBlock[] = [];

  if (!trimmed) {
    const validation = validateTreeNodes([]);
    return { ok: false, title: "Untitled Trajectory Tree", rawText, blocks: [], nodes: [], validation };
  }

  if (startTokens.some((token) => trimmed.includes(token))) {
    blocks = parseWrapped(trimmed);
  } else {
    try {
      blocks = parseRawJson(trimmed);
    } catch (error) {
      blocks = [
        {
          index: 0,
          rawJson: trimmed,
          reasoning: "",
          node: null,
          error: error instanceof Error ? error.message : String(error)
        }
      ];
    }
  }

  const parsedNodes = repairParentReferences(blocks.flatMap((block) => (block.node ? [block.node] : [])));
  const nodes = layoutTreeNodes(parsedNodes, { preservePositions: false });
  const validation = validateTreeNodes(nodes);
  const parseErrors = blocks.filter((block) => block.error);
  const repairedBlocks = blocks.filter((block) => block.node && block.repairNotes && block.repairNotes.length > 0);
  if (parseErrors.length > 0) {
    validation.issues.unshift(
      ...parseErrors.map((block) => ({
        level: "error" as const,
        code: "invalid_json",
        message: `Block ${block.index + 1}: ${block.error}`
      }))
    );
  }
  if (repairedBlocks.length > 0) {
    validation.issues.unshift(
      ...repairedBlocks.map((block) => ({
        level: "warning" as const,
        code: "json_repaired",
        nodeId: block.node?.id,
        message: `Block ${block.index + 1}: repaired JSON (${block.repairNotes?.join("; ")}).`
      }))
    );
  }

  return {
    ok: validation.ok && parseErrors.length === 0,
    title: titleFromNodes(nodes),
    rawText,
    blocks,
    nodes,
    validation: {
      ...validation,
      ok: validation.ok && parseErrors.length === 0
    }
  };
}

function repairParentReferences(nodes: TrajectoryTreeNode[]) {
  const knownIds = new Set<string>();
  const latestByKind = new Map<string, string>();
  const latestByBranchId = new Map<string, string>();
  const latestByBranchLabel = new Map<string, string>();

  return nodes.map((node) => {
    let nextNode = node;
    if (node.parent && !knownIds.has(node.parent)) {
      const parentKey = node.parent.trim().toLowerCase();
      const repairedParent =
        latestByKind.get(parentKey) ??
        latestByBranchId.get(parentKey) ??
        latestByBranchLabel.get(parentKey);
      if (repairedParent) {
        nextNode = {
          ...node,
          parent: repairedParent,
          data: {
            ...node.data,
            parent: repairedParent
          },
          rawJson: JSON.stringify(
            {
              ...node.data,
              parent: repairedParent
            },
            null,
            2
          )
        };
      }
    }

    knownIds.add(nextNode.id);
    if (nextNode.kind.trim()) latestByKind.set(nextNode.kind.trim().toLowerCase(), nextNode.id);
    const branch = asRecord(nextNode.data.branch);
    const branchId = asString(branch.id).trim().toLowerCase();
    const branchLabel = asString(branch.label).trim().toLowerCase();
    if (branchId) latestByBranchId.set(branchId, nextNode.id);
    if (branchLabel) latestByBranchLabel.set(branchLabel, nextNode.id);
    return nextNode;
  });
}

export function parseNodeRawJson(rawJson: string, reasoning: string) {
  const parsed = JSON.parse(rawJson);
  return nodeFromJson(parsed, rawJson, reasoning);
}