import { randomUUID } from "node:crypto";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { isDeltaKind } from "./metrics.ts";
import { parseAgentObservation, type AgentObservation, type AgentTarget } from "./agents.ts";

export interface AgentBridgeOptions extends AgentTarget {
  source: string;
  /** Unique child invocation, not just its reusable role/name. */
  agentId: string;
  /** Deliver to the PARENT bus or an explicitly provided transport. No transport is installed here. */
  emit: (observation: AgentObservation) => void;
}

/** Opt-in inline extension for a child Pi SDK session. It neither creates a child nor calls a provider.
 * Clocks are local to the child. Only allowlisted scalar metadata leaves it.
 */
export function createAgentMetricsBridge(options: AgentBridgeOptions): (pi: ExtensionAPI) => void {
  // Snapshot the target now: callers must never retarget an in-flight child after branch/session changes.
  const { source, agentId, sessionId, anchorId, toolCallId, emit } = options;
  return pi => {
    let current: AgentObservation | undefined;
    let start: number | undefined;
    let ended = false;
    function publish(): void {
      if (!current) return;
      const parsed = parseAgentObservation(current);
      if (!parsed) throw new Error("Invalid Pi Metrics agent bridge identifiers or measurements.");
      emit(parsed);
      current.sequence++;
    }
    function elapsed(): number | null { return start === undefined ? null : performance.now() - start; }
    function interrupt(): void {
      if (!current) return;
      current.state = "interrupted";
      current.elapsedMs = elapsed();
      publish();
      current = undefined;
    }
    function begin(known: boolean, provider: string | null, model: string | null): void {
      interrupt();
      start = known ? performance.now() : undefined;
      ended = false;
      current = { version: 1, source, agentId, sessionId, anchorId, toolCallId, requestId: randomUUID(), sequence: 0,
        provider, model, state: "running", timingBasis: known ? "context_with_system" : "unavailable",
        elapsedMs: known ? 0 : null, ttftMs: null, responseMs: null, outputTokens: null, reasoningTokens: null, stopReason: null };
      publish();
    }
    pi.on("context_with_system", (_event, ctx) => begin(true, ctx.model?.provider ?? null, ctx.model?.id ?? null));
    pi.on("message_update", event => {
      const delta = event.assistantMessageEvent;
      if (!current || ended || current.ttftMs !== null || !isDeltaKind(delta.type) || !("delta" in delta) || !delta.delta) return;
      current.elapsedMs = elapsed();
      current.ttftMs = current.elapsedMs;
      publish();
    });
    pi.on("message_end", (event, ctx) => {
      if (event.message.role !== "assistant") return;
      if (!current) begin(false, ctx.model?.provider ?? null, ctx.model?.id ?? null);
      // Freeze BEFORE tools. Final usage must wait for turn_end transformations.
      current!.responseMs = elapsed();
      ended = true;
    });
    pi.on("turn_end", event => {
      if (!current || !ended || event.message.role !== "assistant") return;
      const m = event.message;
      const count = (n: unknown): number | null => typeof n === "number" && Number.isSafeInteger(n) && n >= 0 ? n : null;
      let output = count(m.usage?.output), reasoning = count(m.usage?.reasoning);
      if (output === null || (reasoning !== null && reasoning > output)) { output = null; reasoning = null; }
      Object.assign(current, { state: "complete", elapsedMs: elapsed(), provider: m.provider, model: m.responseModel ?? m.model,
        outputTokens: output, reasoningTokens: reasoning, stopReason: m.stopReason });
      publish();
      current = undefined;
    });
    pi.on("agent_settled", interrupt);
    pi.on("session_shutdown", interrupt);
  };
}