import { Inject, Logger } from '@nestjs/common';
import {
  WebSocketGateway,
  WebSocketServer,
  SubscribeMessage,
  OnGatewayConnection,
  OnGatewayDisconnect,
  type OnGatewayInit,
  ConnectedSocket,
  MessageBody,
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import type { AgentSessionEvent } from '@mariozechner/pi-coding-agent';
import {
  verifyDiscordIngressEnvelope,
  type DiscordIngressEnvelope,
  type DiscordIngressPayload,
} from '@mosaicstack/discord-plugin';
import type { Auth } from '@mosaicstack/auth';
import type { Brain } from '@mosaicstack/brain';
import { redactSensitiveContent } from '@mosaicstack/log';
import type {
  SetThinkingPayload,
  SlashCommandApprovalResultPayload,
  SlashCommandPayload,
  SystemReloadPayload,
  RoutingDecisionInfo,
  AbortPayload,
} from '@mosaicstack/types';
import { AgentService, type ConversationHistoryMessage } from '../agent/agent.service.js';
import { AUTH } from '../auth/auth.tokens.js';
import {
  scopeFromUser,
  type ActorTenantScope,
  type AuthenticatedUserLike,
} from '../auth/session-scope.js';
import { BRAIN } from '../brain/brain.tokens.js';
import { CommandRegistryService } from '../commands/command-registry.service.js';
import { CommandExecutorService } from '../commands/command-executor.service.js';
import { RoutingEngineService } from '../agent/routing/routing-engine.service.js';
import { v4 as uuid } from 'uuid';
import { ChatSocketMessageDto } from './chat.dto.js';
import { validateDiscordServiceToken, validateSocketSession } from './chat.gateway-auth.js';
import { DiscordReplayProtector } from '../plugin/discord-replay-protector.js';

/** Per-client state tracking streaming accumulation for persistence. */
interface ClientSession {
  conversationId: string;
  cleanup: () => void;
  /** Accumulated assistant response text for the current turn. */
  assistantText: string;
  /** Tool calls observed during the current turn. */
  toolCalls: Array<{ toolCallId: string; toolName: string; args: unknown; isError: boolean }>;
  /** Tool calls in-flight (started but not ended yet). */
  pendingToolCalls: Map<string, { toolName: string; args: unknown }>;
  /** Server-derived owner/tenant scope for this socket's conversation attachment. */
  scope: ActorTenantScope;
  /** Last routing decision made for this session (M4-008) */
  lastRoutingDecision?: RoutingDecisionInfo;
}

/**
 * Per-conversation model overrides set via /model command (M4-007).
 * Keyed by conversationId, value is the model name to use.
 */
const modelOverrides = new Map<string, string>();

function isDiscordIngressEnvelope(value: unknown): value is DiscordIngressEnvelope {
  if (typeof value !== 'object' || value === null) return false;
  const envelope = value as { payload?: unknown; signature?: unknown };
  if (
    typeof envelope.signature !== 'string' ||
    typeof envelope.payload !== 'object' ||
    envelope.payload === null
  ) {
    return false;
  }
  const payload = envelope.payload as Record<string, unknown>;
  return [
    payload['correlationId'],
    payload['messageId'],
    payload['guildId'],
    payload['channelId'],
    payload['userId'],
    payload['conversationId'],
    payload['content'],
  ].every((field: unknown): boolean => typeof field === 'string');
}

function isChatSocketMessage(value: unknown): value is ChatSocketMessageDto {
  if (typeof value !== 'object' || value === null) return false;
  const payload = value as { content?: unknown; conversationId?: unknown };
  return (
    typeof payload.content === 'string' &&
    (payload.conversationId === undefined || typeof payload.conversationId === 'string')
  );
}

@WebSocketGateway({
  cors: {
    origin: process.env['GATEWAY_CORS_ORIGIN'] ?? 'http://localhost:3000',
  },
  namespace: '/chat',
})
export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server!: Server;

  private readonly logger = new Logger(ChatGateway.name);
  private readonly clientSessions = new Map<string, ClientSession>();
  /** Raw stream fragments are kept in memory only until they are safe to redact and emit. */
  private readonly textEgressBuffers = new Map<string, string>();
  private readonly thinkingEgressBuffers = new Map<string, string>();
  private readonly discordReplayProtector = new DiscordReplayProtector();

  constructor(
    @Inject(AgentService) private readonly agentService: AgentService,
    @Inject(AUTH) private readonly auth: Auth,
    @Inject(BRAIN) private readonly brain: Brain,
    @Inject(CommandRegistryService) private readonly commandRegistry: CommandRegistryService,
    @Inject(CommandExecutorService) private readonly commandExecutor: CommandExecutorService,
    @Inject(RoutingEngineService) private readonly routingEngine: RoutingEngineService,
  ) {}

  afterInit(): void {
    this.logger.log('Chat WebSocket gateway initialized');
  }

  async handleConnection(client: Socket): Promise<void> {
    const serviceToken = client.handshake.auth['discordServiceToken'];
    if (validateDiscordServiceToken(serviceToken, process.env['DISCORD_SERVICE_TOKEN'])) {
      client.data.discordService = true;
      this.logger.log(`Authenticated Discord service connected: ${client.id}`);
      return;
    }

    const session = await validateSocketSession(client.handshake.headers, this.auth);
    if (!session) {
      this.logger.warn(`Rejected unauthenticated WebSocket client: ${client.id}`);
      client.disconnect();
      return;
    }

    client.data.user = session.user;
    client.data.session = session.session;
    this.logger.log(`Client connected: ${client.id}`);
    client.emit('commands:manifest', { manifest: this.commandRegistry.getManifest() });
  }

  handleDisconnect(client: Socket): void {
    this.logger.log(`Client disconnected: ${client.id}`);
    const session = this.clientSessions.get(client.id);
    if (session) {
      session.cleanup();
      this.agentService.removeChannel(
        session.conversationId,
        `websocket:${client.id}`,
        session.scope,
      );
      this.clientSessions.delete(client.id);
    }
    this.textEgressBuffers.delete(client.id);
    this.thinkingEgressBuffers.delete(client.id);
  }

  private getClientScope(client: Socket): ActorTenantScope | null {
    const user = client.data.user as AuthenticatedUserLike | undefined;
    if (!user?.id) return null;
    return scopeFromUser(user);
  }

  private modelOverrideKey(conversationId: string, scope: ActorTenantScope): string {
    return `${scope.tenantId}:${scope.userId}:${conversationId}`;
  }

  private scopesEqual(a: ActorTenantScope, b: ActorTenantScope): boolean {
    return a.userId === b.userId && a.tenantId === b.tenantId;
  }

  @SubscribeMessage('message')
  async handleMessage(
    @ConnectedSocket() client: Socket,
    @MessageBody() rawData: unknown,
  ): Promise<void> {
    let discordIngress: DiscordIngressPayload | null = null;
    let data: ChatSocketMessageDto;
    if (client.data.discordService) {
      if (!isDiscordIngressEnvelope(rawData)) {
        this.logger.warn(`Rejected malformed Discord ingress from ${client.id}`);
        return;
      }
      discordIngress = this.resolveDiscordIngress(client, rawData);
      if (!discordIngress) return;
      data = { conversationId: discordIngress.conversationId, content: discordIngress.content };
    } else {
      if (!isChatSocketMessage(rawData)) {
        this.logger.warn(`Rejected malformed chat message from ${client.id}`);
        return;
      }
      data = rawData;
    }
    const conversationId = data.conversationId ?? uuid();
    const discordServiceUserId = process.env['DISCORD_SERVICE_USER_ID'];
    if (discordIngress && !discordServiceUserId) {
      this.logger.warn(
        `Rejected Discord ingress without configured service owner from ${client.id}`,
      );
      return;
    }
    const scope = discordIngress
      ? {
          userId: discordServiceUserId!,
          tenantId: process.env['DISCORD_SERVICE_TENANT_ID'] ?? discordServiceUserId!,
        }
      : this.getClientScope(client);
    if (!scope) {
      client.emit('error', { conversationId, error: 'Authenticated user scope is required.' });
      return;
    }
    const userId = scope.userId;
    const correlationId = discordIngress?.correlationId;

    this.logger.log(
      `Message from ${client.id} in conversation ${conversationId}${correlationId ? ` correlation=${correlationId}` : ''}`,
    );

    // Ensure agent session exists for this conversation
    let sessionRoutingDecision: RoutingDecisionInfo | undefined;
    try {
      let agentSession = this.agentService.getSession(conversationId, scope);
      if (!agentSession) {
        // When resuming an existing conversation, load prior messages to inject as context (M1-004)
        const conversationHistory = await this.loadConversationHistory(conversationId, userId);

        // M5-004: Check if there's an existing sessionId bound to this conversation
        let existingSessionId: string | undefined;
        if (userId) {
          existingSessionId = await this.getConversationSessionId(conversationId, userId);
          if (existingSessionId) {
            this.logger.log(
              `Resuming existing sessionId=${existingSessionId} for conversation=${conversationId}`,
            );
          }
        }

        // Determine provider/model via routing engine or per-session /model override (M4-012 / M4-007)
        let resolvedProvider = data.provider;
        let resolvedModelId = data.modelId;

        const modelOverride = modelOverrides.get(this.modelOverrideKey(conversationId, scope));
        if (modelOverride) {
          // /model override bypasses routing engine (M4-007)
          resolvedModelId = modelOverride;
          this.logger.log(
            `Using /model override "${modelOverride}" for conversation=${conversationId}`,
          );
        } else if (!resolvedProvider && !resolvedModelId) {
          // No explicit provider/model from client — use routing engine (M4-012)
          try {
            const routingDecision = await this.routingEngine.resolve(data.content, userId);
            resolvedProvider = routingDecision.provider;
            resolvedModelId = routingDecision.model;
            sessionRoutingDecision = {
              model: routingDecision.model,
              provider: routingDecision.provider,
              ruleName: routingDecision.ruleName,
              reason: routingDecision.reason,
            };
            this.logger.log(
              `Routing decision for conversation=${conversationId}: ${routingDecision.provider}/${routingDecision.model} (rule="${routingDecision.ruleName}")`,
            );
          } catch (routingErr) {
            this.logger.warn(
              `Routing engine failed for conversation=${conversationId}, using defaults`,
              routingErr instanceof Error ? routingErr.message : String(routingErr),
            );
          }
        }

        // M5-004: Use existingSessionId as sessionId when available (session reuse)
        const sessionIdToCreate = existingSessionId ?? conversationId;
        agentSession = await this.agentService.createSession(sessionIdToCreate, {
          provider: resolvedProvider,
          modelId: resolvedModelId,
          agentConfigId: data.agentId,
          userId,
          tenantId: scope.tenantId,
          conversationHistory: conversationHistory.length > 0 ? conversationHistory : undefined,
        });

        if (conversationHistory.length > 0) {
          this.logger.log(
            `Loaded ${conversationHistory.length} prior messages for conversation=${conversationId}`,
          );
        }
      }
    } catch (err) {
      this.logger.error(
        `Session creation failed for client=${client.id}, conversation=${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
      client.emit('error', {
        conversationId,
        error: 'Failed to start agent session. Please try again.',
      });
      return;
    }

    // Ensure conversation record exists in the DB before persisting messages
    // M5-004: Also bind the sessionId to the conversation record
    if (userId) {
      await this.ensureConversation(conversationId, userId);
      await this.bindSessionToConversation(conversationId, userId, conversationId);
    }

    // M5-007: Count the user message
    this.agentService.recordMessage(conversationId);

    // Persist the user message
    if (userId) {
      try {
        await this.brain.conversations.addMessage(
          {
            conversationId,
            role: 'user',
            content: redactSensitiveContent(data.content).content,
            metadata: {
              timestamp: new Date().toISOString(),
              ...(correlationId
                ? {
                    correlationId,
                    discordMessageId: discordIngress?.messageId,
                    discordUserId: discordIngress?.userId,
                  }
                : {}),
              classifications: redactSensitiveContent(data.content).classifications,
            },
          },
          userId,
        );
      } catch (err) {
        this.logger.error(
          `Failed to persist user message for conversation=${conversationId}`,
          err instanceof Error ? err.stack : String(err),
        );
      }
    }

    // Always clean up previous listener to prevent leak
    const existing = this.clientSessions.get(client.id);
    if (existing) {
      existing.cleanup();
    }

    // Subscribe to agent events and relay to client
    const cleanup = this.agentService.onEvent(
      conversationId,
      (event: AgentSessionEvent) => {
        this.relayEvent(client, conversationId, event);
      },
      scope,
    );

    // Preserve routing decision from the existing client session if we didn't get a new one
    const prevClientSession = this.clientSessions.get(client.id);
    const routingDecisionToStore = sessionRoutingDecision ?? prevClientSession?.lastRoutingDecision;

    this.clientSessions.set(client.id, {
      conversationId,
      cleanup,
      assistantText: '',
      toolCalls: [],
      pendingToolCalls: new Map(),
      scope,
      lastRoutingDecision: routingDecisionToStore,
    });

    // Track channel connection
    this.agentService.addChannel(conversationId, `websocket:${client.id}`, scope);

    // Send session info so the client knows the model/provider (M4-008: include routing decision)
    // Include agentName when a named agent config is active (M5-001)
    {
      const agentSession = this.agentService.getSession(conversationId, scope);
      if (agentSession) {
        const piSession = agentSession.piSession;
        client.emit('session:info', {
          conversationId,
          provider: agentSession.provider,
          modelId: agentSession.modelId,
          thinkingLevel: piSession.thinkingLevel,
          availableThinkingLevels: piSession.getAvailableThinkingLevels(),
          ...(agentSession.agentName ? { agentName: agentSession.agentName } : {}),
          ...(routingDecisionToStore ? { routingDecision: routingDecisionToStore } : {}),
        });
      }
    }

    // Send acknowledgment
    client.emit('message:ack', {
      conversationId,
      messageId: uuid(),
      ...(correlationId
        ? {
            correlationId,
            discordMessageId: discordIngress?.messageId,
            discordUserId: discordIngress?.userId,
          }
        : {}),
    });

    // Dispatch to agent
    try {
      await this.agentService.prompt(conversationId, data.content, scope);
    } catch (err) {
      this.logger.error(
        `Agent prompt failed for client=${client.id}, conversation=${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
      client.emit('error', {
        conversationId,
        error: 'The agent failed to process your message. Please try again.',
      });
    }
  }

  @SubscribeMessage('set:thinking')
  handleSetThinking(
    @ConnectedSocket() client: Socket,
    @MessageBody() data: SetThinkingPayload,
  ): void {
    const scope = this.getClientScope(client);
    if (!scope) {
      client.emit('error', {
        conversationId: data.conversationId,
        error: 'Authenticated user scope is required.',
      });
      return;
    }

    const session = this.agentService.getSession(data.conversationId, scope);
    if (!session) {
      client.emit('error', {
        conversationId: data.conversationId,
        error: 'No active session for this conversation.',
      });
      return;
    }

    const validLevels = session.piSession.getAvailableThinkingLevels();
    if (!validLevels.includes(data.level as never)) {
      client.emit('error', {
        conversationId: data.conversationId,
        error: `Invalid thinking level "${data.level}". Available: ${validLevels.join(', ')}`,
      });
      return;
    }

    session.piSession.setThinkingLevel(data.level as never);
    this.logger.log(
      `Thinking level set to "${data.level}" for conversation ${data.conversationId}`,
    );

    client.emit('session:info', {
      conversationId: data.conversationId,
      provider: session.provider,
      modelId: session.modelId,
      thinkingLevel: session.piSession.thinkingLevel,
      availableThinkingLevels: session.piSession.getAvailableThinkingLevels(),
      ...(session.agentName ? { agentName: session.agentName } : {}),
    });
  }

  @SubscribeMessage('abort')
  async handleAbort(
    @ConnectedSocket() client: Socket,
    @MessageBody() data: AbortPayload,
  ): Promise<void> {
    const conversationId = data.conversationId;
    this.logger.log(`Abort requested by ${client.id} for conversation ${conversationId}`);

    const scope = this.getClientScope(client);
    if (!scope) {
      client.emit('error', { conversationId, error: 'Authenticated user scope is required.' });
      return;
    }

    const session = this.agentService.getSession(conversationId, scope);
    if (!session) {
      client.emit('error', {
        conversationId,
        error: 'No active session to abort.',
      });
      return;
    }

    try {
      await session.piSession.abort();
      this.logger.log(`Agent session ${conversationId} aborted successfully`);
    } catch (err) {
      this.logger.error(
        `Failed to abort session ${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
      client.emit('error', {
        conversationId,
        error: 'Failed to abort the agent operation.',
      });
    }
  }

  @SubscribeMessage('command:execute')
  async handleCommandExecute(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: SlashCommandPayload,
  ): Promise<void> {
    const scope = this.getClientScope(client);
    if (!scope) {
      client.emit('command:result', {
        command: payload.command,
        conversationId: payload.conversationId,
        success: false,
        message: 'Authenticated user scope is required.',
      });
      return;
    }

    const result = await this.commandExecutor.execute(payload, scope);
    client.emit('command:result', result);
  }

  @SubscribeMessage('command:approve')
  async handleCommandApproval(
    @ConnectedSocket() client: Socket,
    @MessageBody() payload: SlashCommandPayload,
  ): Promise<void> {
    const scope = this.getClientScope(client);
    const approval = scope ? await this.commandExecutor.createApproval(payload, scope) : null;
    const result: SlashCommandApprovalResultPayload = approval
      ? {
          command: payload.command,
          conversationId: payload.conversationId,
          success: true,
          approvalId: approval.approvalId,
          expiresAt: approval.expiresAt,
        }
      : {
          command: payload.command,
          conversationId: payload.conversationId,
          success: false,
          message: 'Not authorized to approve this command.',
        };
    client.emit('command:approval', result);
  }

  broadcastReload(payload: SystemReloadPayload): void {
    this.server.emit('system:reload', payload);
    this.logger.log('Broadcasted system:reload to all connected clients');
  }

  /**
   * Set a per-conversation model override (M4-007 / M5-002).
   * When set, the routing engine is bypassed and the specified model is used.
   * Pass null to clear the override and resume automatic routing.
   * M5-005: Emits session:info to clients subscribed to this conversation when a model is set.
   * M5-007: Records a model switch in session metrics.
   */
  setModelOverride(
    conversationId: string,
    modelName: string | null,
    scope: ActorTenantScope,
  ): void {
    const key = this.modelOverrideKey(conversationId, scope);
    if (modelName) {
      modelOverrides.set(key, modelName);
      this.logger.log(`Model override set: conversation=${conversationId} model="${modelName}"`);

      // M5-002: Update the live session's modelId so session:info reflects the new model immediately
      this.agentService.updateSessionModel(conversationId, modelName, scope);

      // M5-005: Broadcast session:info to all clients subscribed to this conversation
      this.broadcastSessionInfo(conversationId, scope);
    } else {
      modelOverrides.delete(key);
      this.logger.log(`Model override cleared: conversation=${conversationId}`);
    }
  }

  /**
   * Return the active model override for a conversation, or undefined if none.
   */
  getModelOverride(conversationId: string, scope: ActorTenantScope): string | undefined {
    return modelOverrides.get(this.modelOverrideKey(conversationId, scope));
  }

  /**
   * M5-005: Broadcast session:info to all clients currently subscribed to a conversation.
   * Called on model or agent switch to ensure the TUI TopBar updates immediately.
   */
  broadcastSessionInfo(
    conversationId: string,
    scope: ActorTenantScope,
    extra?: { agentName?: string; routingDecision?: RoutingDecisionInfo },
  ): void {
    const agentSession = this.agentService.getSession(conversationId, scope);
    if (!agentSession) return;

    const piSession = agentSession.piSession;
    const resolvedAgentName = extra?.agentName ?? agentSession.agentName;
    const payload = {
      conversationId,
      provider: agentSession.provider,
      modelId: agentSession.modelId,
      thinkingLevel: piSession.thinkingLevel,
      availableThinkingLevels: piSession.getAvailableThinkingLevels(),
      ...(resolvedAgentName ? { agentName: resolvedAgentName } : {}),
      ...(extra?.routingDecision ? { routingDecision: extra.routingDecision } : {}),
    };

    // Emit to all clients currently subscribed to this conversation
    for (const [clientId, session] of this.clientSessions) {
      if (session.conversationId === conversationId && this.scopesEqual(session.scope, scope)) {
        const socket = this.server.sockets.sockets.get(clientId);
        if (socket?.connected) {
          socket.emit('session:info', payload);
        }
      }
    }
  }

  /**
   * Ensure a conversation record exists in the DB.
   * Creates it if absent — safe to call concurrently since a duplicate insert
   * would fail on the PK constraint and be caught here.
   */
  private resolveDiscordIngress(
    client: Socket,
    envelope: DiscordIngressEnvelope,
  ): DiscordIngressPayload | null {
    const payload = verifyDiscordIngressEnvelope(
      envelope,
      process.env['DISCORD_SERVICE_TOKEN'] ?? '',
      {
        guildIds: this.readDiscordAllowlist('DISCORD_ALLOWED_GUILD_IDS'),
        channelIds: this.readDiscordAllowlist('DISCORD_ALLOWED_CHANNEL_IDS'),
        userIds: this.readDiscordAllowlist('DISCORD_ALLOWED_USER_IDS'),
      },
    );
    if (!payload) {
      this.logger.warn(`Rejected invalid Discord ingress envelope from ${client.id}`);
      return null;
    }
    if (!this.discordReplayProtector.claim(payload.messageId)) {
      this.logger.warn(
        `Rejected replayed Discord message=${payload.messageId} correlation=${payload.correlationId}`,
      );
      return null;
    }
    return payload;
  }

  private readDiscordAllowlist(name: string): string[] {
    return (process.env[name] ?? '')
      .split(',')
      .map((id: string): string => id.trim())
      .filter((id: string): boolean => id.length > 0);
  }

  private async ensureConversation(conversationId: string, userId: string): Promise<void> {
    try {
      const existing = await this.brain.conversations.findById(conversationId, userId);
      if (!existing) {
        await this.brain.conversations.create({
          id: conversationId,
          userId,
        });
      }
    } catch (err) {
      this.logger.error(
        `Failed to ensure conversation record for conversation=${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
    }
  }

  /**
   * M5-004: Bind the agent sessionId to the conversation record in the DB.
   * Updates the sessionId column so future resumes can reuse the session.
   */
  private async bindSessionToConversation(
    conversationId: string,
    userId: string,
    sessionId: string,
  ): Promise<void> {
    try {
      await this.brain.conversations.update(conversationId, userId, { sessionId });
    } catch (err) {
      this.logger.error(
        `Failed to bind sessionId=${sessionId} to conversation=${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
    }
  }

  /**
   * M5-004: Retrieve the sessionId bound to a conversation, if any.
   * Returns undefined when the conversation does not exist or has no bound session.
   */
  private async getConversationSessionId(
    conversationId: string,
    userId: string,
  ): Promise<string | undefined> {
    try {
      const conv = await this.brain.conversations.findById(conversationId, userId);
      return conv?.sessionId ?? undefined;
    } catch (err) {
      this.logger.error(
        `Failed to get sessionId for conversation=${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
      return undefined;
    }
  }

  /**
   * Load prior conversation messages from DB for context injection on session resume (M1-004).
   * Returns an empty array when no history exists, the conversation is not owned by the user,
   * or userId is not provided.
   */
  private async loadConversationHistory(
    conversationId: string,
    userId: string | undefined,
  ): Promise<ConversationHistoryMessage[]> {
    if (!userId) return [];

    try {
      const messages = await this.brain.conversations.findMessages(conversationId, userId);
      if (messages.length === 0) return [];

      return messages.map((msg) => ({
        role: msg.role as 'user' | 'assistant' | 'system',
        content: msg.content,
        createdAt: msg.createdAt,
      }));
    } catch (err) {
      this.logger.error(
        `Failed to load conversation history for conversation=${conversationId}`,
        err instanceof Error ? err.stack : String(err),
      );
      return [];
    }
  }

  private appendAndFlushRedactedEgress(
    client: Socket,
    conversationId: string,
    eventName: 'agent:text' | 'agent:thinking',
    buffers: Map<string, string>,
    delta: string,
  ): void {
    buffers.set(client.id, `${buffers.get(client.id) ?? ''}${delta}`);
    this.flushRedactedEgress(client, conversationId, eventName, buffers, false);
  }

  /**
   * Holds any suffix that could become a secret, email, or phone number after a
   * later stream chunk. This avoids relying on downstream redaction after data
   * has already reached the socket.
   */
  private flushRedactedEgress(
    client: Socket,
    conversationId: string,
    eventName: 'agent:text' | 'agent:thinking',
    buffers: Map<string, string>,
    final: boolean,
  ): void {
    const buffered = buffers.get(client.id) ?? '';
    const releaseLength = final ? buffered.length : this.safeRedactionPrefixLength(buffered);
    const released = buffered.slice(0, releaseLength);
    const pending = buffered.slice(releaseLength);

    if (pending) {
      buffers.set(client.id, pending);
    } else {
      buffers.delete(client.id);
    }

    if (released) {
      client.emit(eventName, {
        conversationId,
        text: redactSensitiveContent(released).content,
      });
    }
  }

  private safeRedactionPrefixLength(content: string): number {
    let retainedFrom = content.length;

    // Retain the current token because it may become a split secret or email.
    const token = /(?:^|\s)(\S*)$/.exec(content);
    if (token) {
      const matched = token[0] ?? '';
      const trailingToken = token[1] ?? '';
      retainedFrom = token.index + matched.length - trailingToken.length;
    }

    // The secret classifier accepts whitespace around ':' and '=', so preserve
    // a pending label until its value and delimiter are both complete.
    const secretLabel = /(?:api[_-]?key|token|password|secret)\s*[:=]\s*$/i.exec(content);
    if (secretLabel) {
      retainedFrom = Math.min(retainedFrom, secretLabel.index);
    }

    // Phone numbers can contain whitespace and punctuation; preserve the full
    // trailing numeric candidate until a non-phone character establishes a boundary.
    const phone = /(?:^|[^A-Za-z0-9_])(\+?\d[\d(). -]*)$/.exec(content);
    if (phone) {
      const matched = phone[0] ?? '';
      const trailingPhoneCandidate = phone[1] ?? '';
      retainedFrom = Math.min(
        retainedFrom,
        phone.index + matched.length - trailingPhoneCandidate.length,
      );
    }

    return retainedFrom;
  }

  private relayEvent(client: Socket, conversationId: string, event: AgentSessionEvent): void {
    if (!client.connected) {
      this.logger.warn(
        `Dropping event ${event.type} for disconnected client=${client.id}, conversation=${conversationId}`,
      );
      return;
    }

    switch (event.type) {
      case 'agent_start': {
        // Reset accumulation buffers for the new turn
        const cs = this.clientSessions.get(client.id);
        if (cs) {
          cs.assistantText = '';
          cs.toolCalls = [];
          cs.pendingToolCalls.clear();
        }
        this.textEgressBuffers.set(client.id, '');
        this.thinkingEgressBuffers.set(client.id, '');
        client.emit('agent:start', { conversationId });
        break;
      }

      case 'agent_end': {
        // Gather usage stats from the Pi session
        const activeClientSession = this.clientSessions.get(client.id);
        const agentSession = activeClientSession
          ? this.agentService.getSession(conversationId, activeClientSession.scope)
          : undefined;
        const piSession = agentSession?.piSession;
        const stats = piSession?.getSessionStats();
        const contextUsage = piSession?.getContextUsage();

        const usagePayload = stats
          ? {
              provider: agentSession?.provider ?? 'unknown',
              modelId: agentSession?.modelId ?? 'unknown',
              thinkingLevel: piSession?.thinkingLevel ?? 'off',
              tokens: stats.tokens,
              cost: stats.cost,
              context: {
                percent: contextUsage?.percent ?? null,
                window: contextUsage?.contextWindow ?? 0,
              },
            }
          : undefined;

        this.flushRedactedEgress(
          client,
          conversationId,
          'agent:text',
          this.textEgressBuffers,
          true,
        );
        this.flushRedactedEgress(
          client,
          conversationId,
          'agent:thinking',
          this.thinkingEgressBuffers,
          true,
        );
        client.emit('agent:end', {
          conversationId,
          usage: usagePayload,
        });

        // M5-007: Accumulate token usage in session metrics
        if (stats?.tokens) {
          this.agentService.recordTokenUsage(conversationId, {
            input: stats.tokens.input ?? 0,
            output: stats.tokens.output ?? 0,
            cacheRead: stats.tokens.cacheRead ?? 0,
            cacheWrite: stats.tokens.cacheWrite ?? 0,
            total: stats.tokens.total ?? 0,
          });
        }

        // Persist the assistant message with metadata
        const cs = this.clientSessions.get(client.id);
        const userId = (client.data.user as { id: string } | undefined)?.id;
        if (cs && userId && cs.assistantText.trim().length > 0) {
          const metadata: Record<string, unknown> = {
            timestamp: new Date().toISOString(),
            model: agentSession?.modelId ?? 'unknown',
            provider: agentSession?.provider ?? 'unknown',
            toolCalls: cs.toolCalls,
          };

          if (stats?.tokens) {
            metadata['tokenUsage'] = {
              input: stats.tokens.input,
              output: stats.tokens.output,
              cacheRead: stats.tokens.cacheRead,
              cacheWrite: stats.tokens.cacheWrite,
              total: stats.tokens.total,
            };
          }

          this.brain.conversations
            .addMessage(
              {
                conversationId,
                role: 'assistant',
                content: redactSensitiveContent(cs.assistantText).content,
                metadata: {
                  ...metadata,
                  classifications: redactSensitiveContent(cs.assistantText).classifications,
                },
              },
              userId,
            )
            .catch((err: unknown) => {
              this.logger.error(
                `Failed to persist assistant message for conversation=${conversationId}`,
                err instanceof Error ? err.stack : String(err),
              );
            });

          // Reset accumulation
          cs.assistantText = '';
          cs.toolCalls = [];
          cs.pendingToolCalls.clear();
        }
        break;
      }

      case 'message_update': {
        const assistantEvent = event.assistantMessageEvent;
        if (assistantEvent.type === 'text_delta') {
          // Keep raw stream material in memory only; persist and emit only redacted text.
          const cs = this.clientSessions.get(client.id);
          if (cs) {
            cs.assistantText += assistantEvent.delta;
          }
          this.appendAndFlushRedactedEgress(
            client,
            conversationId,
            'agent:text',
            this.textEgressBuffers,
            assistantEvent.delta,
          );
        } else if (assistantEvent.type === 'thinking_delta') {
          this.appendAndFlushRedactedEgress(
            client,
            conversationId,
            'agent:thinking',
            this.thinkingEgressBuffers,
            assistantEvent.delta,
          );
        }
        break;
      }

      case 'tool_execution_start': {
        // Track pending tool call for later recording
        const cs = this.clientSessions.get(client.id);
        if (cs) {
          cs.pendingToolCalls.set(event.toolCallId, {
            toolName: event.toolName,
            args: event.args,
          });
        }
        client.emit('agent:tool:start', {
          conversationId,
          toolCallId: event.toolCallId,
          toolName: event.toolName,
        });
        break;
      }

      case 'tool_execution_end': {
        // Finalise tool call record
        const cs = this.clientSessions.get(client.id);
        if (cs) {
          const pending = cs.pendingToolCalls.get(event.toolCallId);
          cs.toolCalls.push({
            toolCallId: event.toolCallId,
            toolName: event.toolName,
            args: pending?.args ?? null,
            isError: event.isError,
          });
          cs.pendingToolCalls.delete(event.toolCallId);
        }
        client.emit('agent:tool:end', {
          conversationId,
          toolCallId: event.toolCallId,
          toolName: event.toolName,
          isError: event.isError,
        });
        break;
      }
    }
  }
}