import { resultValueError } from "pi-agents/src/model/json-schema.js";
import type { AttributionManifest } from "../types.js";
import { attributionManifestSchema } from "./result-contracts.js";

const OPEN = "<!-- ULTRAPI_ATTRIBUTION\n";
const CLOSE = "\n-->";

export const ROOT_ATTRIBUTION_INSTRUCTION = `Before the final answer ends, append exactly one hidden comment ${OPEN}{"usedFactIds":[],"rejectedFactIds":[],"decisions":[],"changes":[]}${CLOSE}. Fill it with the root decision attribution; UltraPi removes it before display.`;

export function extractRootAttribution(text: string): { visibleText: string; manifest?: AttributionManifest; invalid?: true } {
  const start = text.lastIndexOf(OPEN);
  if (start < 0) return { visibleText: text };
  const end = text.indexOf(CLOSE, start + OPEN.length);
  if (end < 0) return { visibleText: text.slice(0, start).trimEnd(), invalid: true };
  const visibleText = `${text.slice(0, start)}${text.slice(end + CLOSE.length)}`.trimEnd();
  let parsed: unknown;
  try { parsed = JSON.parse(text.slice(start + OPEN.length, end)); }
  catch { return { visibleText, invalid: true }; }
  if (resultValueError(parsed, attributionManifestSchema())) return { visibleText, invalid: true };
  return { visibleText, manifest: parsed as AttributionManifest };
}