import type { Spec } from "@json-render/core";

export type AnswerTaskLayout = {
  overview: string[];
  analysis: string[];
  evidence: string[];
  followUp: string[];
};

type ContentRole = keyof AnswerTaskLayout;

/** Keep authored comparisons and mixed sections intact. Only compose independent roots. */
export function getAnswerTaskLayout(spec: Spec): AnswerTaskLayout | null {
  const root = spec.elements[spec.root];
  if (root.type !== "Stack" || root.props.columns !== "one" || root.props.intent === "comparison") return null;
  const roleFor = (id: string): ContentRole => {
    const node = spec.elements[id];
    if (node.type === "Section") {
      const roles: Record<string, ContentRole> = { result: "overview", analysis: "analysis", evidence: "evidence", follow_up: "followUp" };
      const role = typeof node.props.role === "string" ? roles[node.props.role] : undefined;
      if (role) return role;
    }
    if (node.type === "Evidence") return "evidence";
    if (node.type === "FollowUpForm") return "followUp";
    if (["Chart", "Table", "MapLocation"].includes(node.type)) return "analysis";
    if (node.children?.length) {
      const roles = new Set(node.children.map(roleFor));
      if (roles.size === 1) return [...roles][0];
    }
    return "overview";
  };
  const groups: AnswerTaskLayout = { overview: [], analysis: [], evidence: [], followUp: [] };
  for (const id of root.children ?? []) groups[roleFor(id)].push(id);
  return groups;
}