// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

/**
 * Main OpenClaw hook replay dispatcher.
 *
 * OpenClaw hook callbacks arrive here as lifecycle, LLM, model-timing, tool, and
 * subagent events. This class routes each event to focused replay modules and
 * owns fail-open behavior so observability never breaks the agent runtime.
 */
import type { NemoFlowHookBackendConfig } from "./config.js";
import { emitMark, toJsonRecord } from "./hook-replay/marks.js";
import { llmKey } from "./hook-replay/correlation.js";
import {
  emitUnpairedModelCallTimingMarks,
  recordBeforeMessageWrite,
  recordLlmInput,
  recordLlmOutput,
  recordModelCallEnded,
  recordModelCallStarted,
  replayAgentEndMessages,
  replayPendingLlmOutputsForSession,
} from "./hook-replay/llm.js";
import { guardBeforeToolCall, replayAfterToolCall } from "./hook-replay/tool.js";
import {
  createHookReplayState,
  drainSession,
  closeSessionRoot,
  deleteSession,
  ensureSession,
  resolveSessionKey,
  type HookReplayBackendState,
  type SessionLookupInput,
  type SessionState,
} from "./hook-replay/session.js";
import type { NemoFlowRuntimeModule } from "./modules.js";
import type {
  PluginHookAfterToolCallEvent,
  PluginHookAgentContext,
  PluginHookAgentEndEvent,
  PluginHookBeforeAgentFinalizeEvent,
  PluginHookBeforeMessageWriteContext,
  PluginHookBeforeMessageWriteEvent,
  PluginHookBeforeToolCallEvent,
  PluginHookGatewayContext,
  PluginHookGatewayStartEvent,
  PluginHookLlmInputEvent,
  PluginHookLlmOutputEvent,
  PluginHookModelCallEndedEvent,
  PluginHookModelCallStartedEvent,
  PluginHookSessionContext,
  PluginHookSessionEndEvent,
  PluginHookSessionStartEvent,
  PluginHookSubagentContext,
  PluginHookSubagentEndedEvent,
  PluginHookSubagentSpawnedEvent,
  PluginHookToolContext,
} from "./openclaw-hook-types.js";
import type { PluginLogger } from "openclaw/plugin-sdk/plugin-entry";
import type { JsonObject as JsonRecord } from "nemo-flow-node/typed";

export type HookReplayBackendOptions = {
  nf: NemoFlowRuntimeModule;
  config: NemoFlowHookBackendConfig;
  logger: PluginLogger;
  agentVersion: string;
};

/** Replays OpenClaw public hook events into NeMo Flow scopes, spans, and marks. */
export class HookReplayBackend {
  private readonly nf: NemoFlowRuntimeModule;
  private readonly config: NemoFlowHookBackendConfig;
  private readonly logger: PluginLogger;
  private readonly agentVersion: string;
  private readonly stateValue = createHookReplayState();
  private readonly warningCounts = new Map<string, number>();

  constructor(options: HookReplayBackendOptions) {
    this.nf = options.nf;
    this.config = options.config;
    this.logger = options.logger;
    this.agentVersion = options.agentVersion;
  }

  /** Return mutable replay state for tests and health snapshots. */
  state(): HookReplayBackendState {
    return this.stateValue;
  }

  /** Keep gateway_start registered even though session roots are created lazily. */
  onGatewayStart(_event: PluginHookGatewayStartEvent, _ctx: PluginHookGatewayContext): void {
    // Gateway events have no session root in the hook backend. Keep this hook
    // registered so later telemetry lifecycle can attach without changing the shell.
  }

  /** Open or alias an explicit OpenClaw session root. */
  onSessionStart(event: PluginHookSessionStartEvent, ctx: PluginHookSessionContext): void {
    this.ensureSession({
      sessionId: event.sessionId,
      sessionKey: event.sessionKey ?? ctx.sessionKey,
      agentId: ctx.agentId,
      source: "session_start",
      resumedFrom: event.resumedFrom,
    });

    // ensureSession opens the root scope and emits openclaw.session_start for both explicit and lazy sessions.
  }

  /** Close one explicit OpenClaw session and export its ATIF artifact. */
  async onSessionEnd(event: PluginHookSessionEndEvent, ctx: PluginHookSessionContext): Promise<void> {
    const session = this.ensureSession({
      sessionId: event.sessionId,
      sessionKey: event.sessionKey ?? ctx.sessionKey,
      agentId: ctx.agentId,
      source: "lazy_session",
    });

    if (!session) {
      return;
    }

    await this.closeSession(session, sessionEndSummary(event));
  }

  /** Buffer an LLM request snapshot until a matching response or trajectory replay arrives. */
  onLlmInput(event: PluginHookLlmInputEvent, ctx: PluginHookAgentContext): void {
    recordLlmInput(this.sessionManager(), event, ctx);
  }

  /** Replay an LLM output immediately or keep it briefly for a late input snapshot. */
  onLlmOutput(event: PluginHookLlmOutputEvent, ctx: PluginHookAgentContext): void {
    recordLlmOutput(this.sessionManager(), event, ctx);
  }

  /** Record provider-call start timing when OpenClaw exposes a call id. */
  onModelCallStarted(event: PluginHookModelCallStartedEvent, ctx: PluginHookAgentContext): void {
    recordModelCallStarted(this.sessionManager(), event, ctx);
  }

  /** Record provider-call completion timing for later LLM-span correlation. */
  onModelCallEnded(event: PluginHookModelCallEndedEvent, ctx: PluginHookAgentContext): void {
    recordModelCallEnded(this.sessionManager(), event, ctx);
  }

  /** Replay a finished OpenClaw tool call as a NeMo Flow tool span or blocked mark. */
  onAfterToolCall(event: PluginHookAfterToolCallEvent, ctx: PluginHookToolContext): void {
    replayAfterToolCall(this.sessionManager(), event, ctx);
  }

  /** Run conditional-execution guardrails before OpenClaw invokes a tool. */
  async onBeforeToolCall(event: PluginHookBeforeToolCallEvent, ctx: PluginHookToolContext): Promise<void> {
    await guardBeforeToolCall(this.sessionManager(), event, ctx);
  }

  /** Capture assistant message writes that may contain the clearest provider output. */
  onBeforeMessageWrite(event: PluginHookBeforeMessageWriteEvent, ctx: PluginHookBeforeMessageWriteContext): void {
    recordBeforeMessageWrite(this.sessionManager(), event, ctx);
  }

  /** Finalize one agent run, replaying message-write trajectory when needed. */
  onAgentEnd(event: PluginHookAgentEndEvent, ctx: PluginHookAgentContext): void {
    const session = this.ensureSession({
      sessionId: ctx.sessionId,
      sessionKey: ctx.sessionKey,
      runId: event.runId ?? ctx.runId,
      agentId: ctx.agentId,
      source: "lazy_session",
    });

    if (!session) {
      return;
    }

    const finalOutput = replayAgentEndMessages(this.sessionManager(), event, ctx, session);
    if (finalOutput && (!session.finalOutput || "content" in finalOutput)) {
      session.finalOutput = finalOutput;
    }

    this.emitSessionMark(
      "openclaw.agent_end",
      session,
      toJsonRecord({
        runId: event.runId ?? ctx.runId,
        success: event.success,
        error: event.error,
        durationMs: event.durationMs,
        messageCount: event.messages.length,
      }),
    );
  }

  /** Remember the last assistant text before OpenClaw finalizes the response. */
  onBeforeAgentFinalize(event: PluginHookBeforeAgentFinalizeEvent, ctx: PluginHookAgentContext): void {
    const session = this.ensureSession({
      sessionId: event.sessionId,
      sessionKey: event.sessionKey ?? ctx.sessionKey,
      runId: event.runId ?? ctx.runId,
      agentId: ctx.agentId,
      source: "lazy_session",
    });

    if (!session) {
      return;
    }

    if (typeof event.lastAssistantMessage === "string" && event.lastAssistantMessage.length > 0) {
      session.finalOutput = toJsonRecord({
        content: event.lastAssistantMessage,
        source: "openclaw.before_agent_finalize",
        runId: event.runId ?? ctx.runId,
      });
    }

    this.emitSessionMark(
      "openclaw.before_agent_finalize",
      session,
      toJsonRecord({
        runId: event.runId ?? ctx.runId,
        turnId: event.turnId,
        provider: event.provider,
        model: event.model,
        cwd: event.cwd,
        transcriptPath: event.transcriptPath,
        stopHookActive: event.stopHookActive,
        messageCount: event.messages?.length,
      }),
    );
  }

  /** Attach subagent spawn metadata to the requester session when possible. */
  onSubagentSpawned(event: PluginHookSubagentSpawnedEvent, ctx: PluginHookSubagentContext): void {
    const session =
      this.ensureSession({
        requesterSessionKey: ctx.requesterSessionKey,
        source: "lazy_session",
      }) ??
      this.ensureSession({
        childSessionKey: ctx.childSessionKey ?? event.childSessionKey,
        runId: ctx.runId ?? event.runId,
        agentId: event.agentId,
        source: "lazy_session",
      });

    if (!session) {
      return;
    }

    this.emitSessionMark(
      "openclaw.subagent_spawned",
      session,
      toJsonRecord({
        runId: event.runId,
        childSessionKey: event.childSessionKey,
        agentId: event.agentId,
        label: event.label,
        mode: event.mode,
        threadRequested: event.threadRequested,
      }),
    );
  }

  /** Attach subagent completion metadata to the requester or child session. */
  onSubagentEnded(event: PluginHookSubagentEndedEvent, ctx: PluginHookSubagentContext): void {
    const session =
      this.ensureSession({
        requesterSessionKey: ctx.requesterSessionKey,
        source: "lazy_session",
      }) ??
      this.ensureSession({
        childSessionKey: ctx.childSessionKey ?? event.targetSessionKey,
        runId: ctx.runId ?? event.runId,
        source: "lazy_session",
      });

    if (!session) {
      return;
    }

    this.emitSessionMark(
      "openclaw.subagent_ended",
      session,
      toJsonRecord({
        runId: event.runId ?? ctx.runId,
        targetSessionKey: event.targetSessionKey,
        targetKind: event.targetKind,
        reason: event.reason,
        outcome: event.outcome,
        error: event.error,
        endedAt: event.endedAt,
        sendFarewell: event.sendFarewell,
        accountId: event.accountId,
      }),
    );
  }

  /** Drain all active sessions when the OpenClaw gateway is stopping. */
  async drainForGatewayStop(reason?: string): Promise<void> {
    await this.closeAllSessions({ reason: reason ?? "gateway_stop" });
  }

  /** Close one session selected by a runtime lifecycle cleanup hook. */
  async cleanupSession(input: SessionLookupInput & { reason: string }): Promise<void> {
    const key = resolveSessionKey(this.stateValue, input);
    if (!key) {
      return;
    }

    const session = this.stateValue.sessions.get(key);
    if (!session) {
      return;
    }

    await this.closeSession(session, { reason: input.reason });
  }

  /** Stop the backend and close every active session. */
  async stop(reason: string): Promise<void> {
    await this.closeAllSessions({ reason });
  }

  /** Run replay code with bounded warning logs and no exception escape. */
  safeReplay(label: string, session: SessionState | undefined, emit: () => void): void {
    try {
      emit();
    } catch (error) {
      this.stateValue.counters.replayErrors += 1;
      this.logBoundedWarn(
        `safe-replay:${label}`,
        `nemo-flow replay failed: label=${label} session=${session?.sessionId ?? "unknown"} error=${toMessage(error)}`,
      );
    }
  }

  /** Async variant of safeReplay for hooks that need export or cleanup awaits. */
  async safeReplayAsync(
    label: string,
    session: SessionState | undefined,
    emit: () => Promise<void>,
  ): Promise<void> {
    try {
      await emit();
    } catch (error) {
      this.stateValue.counters.replayErrors += 1;
      this.logBoundedWarn(
        `safe-replay:${label}`,
        `nemo-flow async replay failed: label=${label} session=${session?.sessionId ?? "unknown"} error=${toMessage(error)}`,
      );
    }
  }

  /** Emit spans/marks under the stored session scope stack and ATIF capture window. */
  emitCapturedUnderSession(label: string, session: SessionState, emit: () => void): void {
    this.safeReplay(label, session, () => {
      const previousStack = this.nf.currentScopeStack();
      try {
        this.nf.setThreadScopeStack(session.stack);
        emit();
      } finally {
        this.nf.setThreadScopeStack(previousStack);
      }
    });
  }

  /** Force any pending LLM outputs for a session to replay before closure. */
  replayPendingLlmOutputsForSession(session: SessionState, options: { allowPlaceholderRequest: boolean }): void {
    replayPendingLlmOutputsForSession(this.sessionManager(), session, options);
  }

  /** Emit model-call timing diagnostics that could not be paired with an LLM span. */
  emitUnpairedModelCallTimingMarks(session: SessionState): void {
    emitUnpairedModelCallTimingMarks(this.sessionManager(), session);
  }

  /** Create or resolve a session through the shared session manager facade. */
  private ensureSession(input: Parameters<typeof ensureSession>[1]): SessionState | undefined {
    return ensureSession(this.sessionManager(), input);
  }

  /** Drain, close, export, and delete one session. */
  private async closeSession(session: SessionState, summary: JsonRecord): Promise<void> {
    drainSession(this.sessionManager(), session);
    closeSessionRoot(this.sessionManager(), session, summary, session.finalOutput ?? summary);
    deleteSession(this.stateValue, session);
  }

  /** Emit a session-level OpenClaw lifecycle mark. */
  private emitSessionMark(name: string, session: SessionState, data: JsonRecord): void {
    this.emitCapturedUnderSession(name, session, () => {
      emitMark({
        nf: this.nf,
        state: this.stateValue,
        session,
        name,
        data,
      });
    });
  }

  /** Close every active session with the same lifecycle summary. */
  private async closeAllSessions(summary: JsonRecord): Promise<void> {
    for (const session of [...this.stateValue.sessions.values()]) {
      await this.closeSession(session, summary);
    }
  }

  /** Build the narrow manager interface consumed by focused replay modules. */
  private sessionManager() {
    return {
      nf: this.nf,
      config: this.config,
      logger: this.logger,
      state: this.stateValue,
      agentVersion: this.agentVersion,
      emitCapturedUnderSession: (label: string, session: SessionState, emit: () => void) =>
        this.emitCapturedUnderSession(label, session, emit),
      replayPendingLlmOutputsForSession: (session: SessionState, options: { allowPlaceholderRequest: boolean }) =>
        this.replayPendingLlmOutputsForSession(session, options),
      emitUnpairedModelCallTimingMarks: (session: SessionState) => this.emitUnpairedModelCallTimingMarks(session),
      logBoundedWarn: (key: string, message: string) => this.logBoundedWarn(key, message),
    };
  }

  /** Log one warning per key to avoid noisy repeated hook failures. */
  private logBoundedWarn(key: string, message: string): void {
    const count = this.warningCounts.get(key) ?? 0;
    this.warningCounts.set(key, count + 1);
    if (count === 0) {
      this.logger.warn?.(message);
    }
  }
}

export { llmKey };

/** Expose session-key resolution for tests without exporting the full session module. */
export function resolveBackendSessionKey(
  state: HookReplayBackendState,
  input: Parameters<typeof resolveSessionKey>[1],
): string | undefined {
  return resolveSessionKey(state, input);
}

/** Build the lifecycle summary stored as the session_end mark payload. */
function sessionEndSummary(event: PluginHookSessionEndEvent): JsonRecord {
  return toJsonRecord({
    sessionId: event.sessionId,
    sessionKey: event.sessionKey,
    messageCount: event.messageCount,
    durationMs: event.durationMs,
    reason: event.reason,
    sessionFile: event.sessionFile,
    transcriptArchived: event.transcriptArchived,
    nextSessionId: event.nextSessionId,
    nextSessionKey: event.nextSessionKey,
  });
}

/** Convert thrown values into stable log strings. */
function toMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}
