import type { SessionEntry, SessionTreeNode } from "@earendil-works/pi-coding-agent";
import { classifyOwnership } from "./ownership.ts";
import type { OwnershipState, SemanticNode, SerializedTree, TreeLabelerSettings } from "./types.ts";

function normalize(value: string): string {
  return value.replace(/[\u0000-\u001f\u007f]+/g, " ").replace(/\s+/g, " ").trim();
}

function redact(value: string): string {
  return value
    .replace(/-----BEGIN [^-]+PRIVATE KEY-----[\s\S]*?-----END [^-]+PRIVATE KEY-----/g, "[REDACTED PRIVATE KEY]")
    .replace(/\b(?:sk|ghp|github_pat|xox[baprs])[-_A-Za-z0-9]{16,}\b/g, "[REDACTED TOKEN]")
    .replace(/\bBearer\s+[A-Za-z0-9._~+\/-]+=*/gi, "Bearer [REDACTED]")
    .replace(/[A-Za-z]:\\(?:[^\s<>:"|?*]+\\)+/g, "[PATH]\\")
    .replace(/\/(?:Users|home)\/[^\s/]+\//g, "~/");
}

function contentText(content: unknown, includeToolNames: boolean): string {
  if (typeof content === "string") return content;
  if (!Array.isArray(content)) return "";
  const parts: string[] = [];
  for (const block of content) {
    if (!block || typeof block !== "object") continue;
    const item = block as { type?: string; text?: string; name?: string };
    if (item.type === "text" && typeof item.text === "string") parts.push(item.text);
    if (includeToolNames && item.type === "toolCall" && typeof item.name === "string") parts.push(`[tool: ${item.name}]`);
  }
  return parts.join(" ");
}

function describe(entry: SessionEntry, includeToolNames: boolean): { role: SemanticNode["role"]; text: string } | undefined {
  if (entry.type === "message") {
    const message = entry.message;
    if (message.role !== "user" && message.role !== "assistant") return undefined;
    return { role: message.role, text: contentText("content" in message ? message.content : "", includeToolNames) };
  }
  if (entry.type === "compaction") return { role: "compaction", text: entry.summary };
  if (entry.type === "branch_summary") return { role: "branch-summary", text: entry.summary };
  // Custom messages may contain hidden extension state or policy text; exclude them unless
  // a future explicit allowlist setting is added.
  return undefined;
}

export function buildSemanticTree(
  roots: SessionTreeNode[],
  activeIds: Set<string>,
  ownership: OwnershipState,
  settings: TreeLabelerSettings,
): SemanticNode[] {
  const nodes: SemanticNode[] = [];
  const stack: Array<{ treeNode: SessionTreeNode; semanticParentId: string | null; depth: number }> = [];
  for (let index = roots.length - 1; index >= 0; index--) stack.push({ treeNode: roots[index], semanticParentId: null, depth: 0 });
  while (stack.length) {
    const { treeNode, semanticParentId, depth } = stack.pop()!;
    const description = describe(treeNode.entry, settings.includeToolNames);
    const nextParent = description ? treeNode.entry.id : semanticParentId;
    const nextDepth = description ? depth + 1 : depth;
    if (description) {
      const fullExcerpt = redact(normalize(description.text));
      nodes.push({
        id: treeNode.entry.id,
        parentId: semanticParentId,
        role: description.role,
        excerpt: fullExcerpt.slice(0, settings.excerptChars),
        existingLabel: treeNode.label,
        ownership: classifyOwnership(treeNode.entry.id, ownership),
        active: activeIds.has(treeNode.entry.id),
        depth,
        childCount: 0,
        branchRoot: false,
        branchPoint: false,
        leaf: false,
        excerptTruncated: fullExcerpt.length > settings.excerptChars,
      });
    }
    for (let index = treeNode.children.length - 1; index >= 0; index--) {
      stack.push({ treeNode: treeNode.children[index], semanticParentId: nextParent, depth: nextDepth });
    }
  }

  const byId = new Map(nodes.map((node) => [node.id, node]));
  for (const node of nodes) {
    if (!node.parentId) continue;
    const parent = byId.get(node.parentId);
    if (parent) parent.childCount += 1;
  }
  for (const node of nodes) {
    node.branchPoint = node.childCount > 1;
    node.branchRoot = Boolean(node.parentId && (byId.get(node.parentId)?.childCount ?? 0) > 1);
    node.leaf = node.childCount === 0;
  }
  return nodes;
}

function priority(node: SemanticNode): number {
  return (node.existingLabel ? 80 : 0) + (node.active ? 30 : 0) + (node.branchRoot ? 50 : 0) + (node.branchPoint ? 45 : 0) + (node.leaf ? 35 : 0) + (node.role === "branch-summary" || node.role === "compaction" ? 20 : 0);
}

function flags(node: SemanticNode): string {
  return [
    node.active ? "active" : "",
    node.branchPoint ? "branch-point" : "",
    node.branchRoot ? "branch-root" : "",
    node.leaf ? "leaf" : "",
  ].filter(Boolean).join(",") || "linear";
}

export function serializeTree(nodes: SemanticNode[], settings: TreeLabelerSettings): SerializedTree {
  let pool = settings.exposure === "branch" ? nodes.filter((node) => node.active) : nodes;
  const limitsHit = new Set<SerializedTree["limitsHit"][number]>();
  if (pool.some((node) => node.excerptTruncated)) limitsHit.add("excerptChars");
  const order = new Map(nodes.map((node, index) => [node.id, index]));
  if (pool.length > settings.maxNodes) {
    limitsHit.add("maxNodes");
    const keep = new Set([...pool]
      .sort((a, b) => priority(b) - priority(a) || (order.get(b.id)! - order.get(a.id)!))
      .slice(0, settings.maxNodes)
      .map((node) => node.id));
    pool = pool.filter((node) => keep.has(node.id));
  }

  const exposedIds = new Set(pool.map((node) => node.id));
  const allById = new Map(nodes.map((node) => [node.id, node]));
  const nearestExposedParent = (node: SemanticNode): string | null => {
    let parentId = node.parentId;
    while (parentId && !exposedIds.has(parentId)) parentId = allById.get(parentId)?.parentId ?? null;
    return parentId;
  };
  const header = ["<tree>", "Each node is: ID | nearest-exposed-parent | role | flags | ownership | existing-label | excerpt"];
  const baseLines = pool.map((node) => `${node.id} | ${nearestExposedParent(node) ?? "ROOT"} | ${node.role} | ${flags(node)} | ${node.ownership} | ${JSON.stringify(node.existingLabel ?? "")} | `);
  const fixedLength = [...header, ...baseLines, "</tree>"].join("\n").length;
  if (fixedLength > settings.maxInputChars) {
    throw new Error(`The ${pool.length}-node tree skeleton exceeds maxInputChars=${settings.maxInputChars}. Lower maxNodes or use exposure \"branch\".`);
  }

  const richExcerpt = (node: SemanticNode): boolean => {
    if (settings.exposure === "tree" || settings.exposure === "branch") return true;
    if (settings.exposure === "balanced") {
      return node.active || node.branchRoot || node.branchPoint || node.leaf || Boolean(node.existingLabel);
    }
    const parent = node.parentId ? allById.get(node.parentId) : undefined;
    return node.active || Boolean(node.existingLabel) || Boolean(parent?.active && (node.branchRoot || node.leaf));
  };

  let remaining = settings.maxInputChars - fixedLength;
  const excerpts = new Map<string, string>();
  for (const node of [...pool].sort((a, b) => priority(b) - priority(a))) {
    const wanted = richExcerpt(node) ? node.excerpt : node.excerpt.slice(0, Math.min(100, settings.excerptChars));
    if (remaining <= 2) {
      if (wanted) limitsHit.add("maxInputChars");
      continue;
    }
    const clipped = wanted.slice(0, Math.max(0, remaining - 2));
    if (clipped.length < wanted.length) limitsHit.add("maxInputChars");
    if (clipped) {
      const encoded = JSON.stringify(clipped);
      excerpts.set(node.id, encoded);
      remaining -= encoded.length;
    }
  }
  const lines = [...header, ...pool.map((node, index) => `${baseLines[index]}${excerpts.get(node.id) ?? JSON.stringify("")}`), "</tree>"];
  const text = lines.join("\n");
  return { text, exposedIds, nodeCount: pool.length, omittedCount: nodes.length - pool.length, limitsHit: [...limitsHit] };
}