/* eslint-disable no-console */
// src/agents/AgentContext.ts
import { SystemMessage } from '@langchain/core/messages';
import { RunnableLambda } from '@langchain/core/runnables';
import type {
  UsageMetadata,
  BaseMessage,
  BaseMessageFields,
} from '@langchain/core/messages';
import type { RunnableConfig, Runnable } from '@langchain/core/runnables';
import type * as t from '@/types';
import type { createPruneMessages } from '@/messages';
import { createSchemaOnlyTools } from '@/tools/schema';
import { ContentTypes, Providers } from '@/common';
import { toJsonSchema } from '@/utils/schema';

/**
 * Encapsulates agent-specific state that can vary between agents in a multi-agent system
 */
export class AgentContext {
  /**
   * Create an AgentContext from configuration with token accounting initialization
   */
  static fromConfig(
    agentConfig: t.AgentInputs,
    tokenCounter?: t.TokenCounter,
    indexTokenCountMap?: Record<string, number>
  ): AgentContext {
    const {
      agentId,
      name,
      description,
      provider,
      clientOptions,
      tools,
      toolMap,
      toolEnd,
      toolRegistry,
      toolDefinitions,
      instructions,
      additional_instructions,
      streamBuffer,
      maxContextTokens,
      reasoningKey,
      useLegacyContent,
      dynamicContext,
      structuredOutput: structuredOutputCamel,

      structured_output: structuredOutputSnake,
      discoveredTools,
      summarizeCallback,
      persistedSummary,
      summarizationConfig,
      fileManifest,
      system_cache_blocks,
      instructions_cache_ttl,
    } = agentConfig;

    // Normalize structured output: support both camelCase and snake_case inputs
    // Priority: structuredOutput (camelCase) > structured_output (snake_case with enabled check)
    let structuredOutput: t.StructuredOutputConfig | undefined;
    if (structuredOutputCamel) {
      structuredOutput = structuredOutputCamel;
    } else if (
      structuredOutputSnake?.enabled === true &&
      structuredOutputSnake.schema != null
    ) {
      // Convert snake_case input to StructuredOutputConfig
      structuredOutput = {
        schema: structuredOutputSnake.schema,
        name: structuredOutputSnake.name,
        description: structuredOutputSnake.description,
        mode: structuredOutputSnake.mode,
        strict: structuredOutputSnake.strict,
      };
    }

    const agentContext = new AgentContext({
      agentId,
      name: name ?? agentId,
      description,
      provider,
      clientOptions,
      maxContextTokens,
      streamBuffer,
      tools,
      toolMap,
      toolRegistry,
      toolDefinitions,
      instructions,
      additionalInstructions: additional_instructions,
      reasoningKey,
      toolEnd,
      instructionTokens: 0,
      tokenCounter,
      useLegacyContent,
      dynamicContext,
      structuredOutput,
      discoveredTools,
      summarizeCallback,
      persistedSummary,
      summarizationConfig,
      fileManifest,
      systemCacheBlocks: system_cache_blocks,
      instructionsCacheTtl: instructions_cache_ttl,
    });

    /**
     * Track upstream-aligned subagent inputs on the context. `_sourceInputs`
     * preserves the original AgentInputs so SubagentExecutor can self-spawn
     * (`SubagentConfig.self === true`) without separate config; the other
     * two flow straight from agentConfig.
     */
    agentContext._sourceInputs = agentConfig;
    agentContext.subagentConfigs = agentConfig.subagentConfigs;
    agentContext.maxSubagentDepth = agentConfig.maxSubagentDepth;

    if (tokenCounter) {
      // Initialize system runnable BEFORE async tool token calculation
      // This ensures system message tokens are in instructionTokens before
      // updateTokenMapWithInstructions is called
      agentContext.initializeSystemRunnable();

      const tokenMap = indexTokenCountMap || {};
      agentContext.baseIndexTokenCountMap = { ...tokenMap };
      agentContext.indexTokenCountMap = tokenMap;
      agentContext.tokenCalculationPromise = agentContext
        .calculateInstructionTokens(tokenCounter)
        .then(() => {
          // Update token map with instruction tokens (includes system + tool tokens)
          agentContext.updateTokenMapWithInstructions(tokenMap);
        })
        .catch((err) => {
          console.error('Error calculating instruction tokens:', err);
        });
    } else if (indexTokenCountMap) {
      agentContext.baseIndexTokenCountMap = { ...indexTokenCountMap };
      agentContext.indexTokenCountMap = indexTokenCountMap;
    }

    return agentContext;
  }

  /** Agent identifier */
  agentId!: string;
  /** Human-readable name for this agent (used in handoff context). Falls back to agentId if not provided. */
  name?: string;
  /** Description of what this agent does (used to enrich handoff tool descriptions). */
  description?: string;
  /** Provider for this specific agent */
  provider!: Providers;
  /** Client options for this agent */
  clientOptions?: t.ClientOptions;
  /** Token count map indexed by message position */
  indexTokenCountMap: Record<string, number | undefined> = {};
  /** Canonical pre-run token map used to restore token accounting on reset */
  baseIndexTokenCountMap: Record<string, number> = {};
  /** Maximum context tokens for this agent */
  maxContextTokens?: number;
  /** Current usage metadata for this agent */
  currentUsage?: Partial<UsageMetadata>;
  /** Prune messages function configured for this agent */
  pruneMessages?: ReturnType<typeof createPruneMessages>;
  /** Token counter function for this agent */
  tokenCounter?: t.TokenCounter;
  /** Instructions/system message token count */
  instructionTokens: number = 0;
  /** The amount of time that should pass before another consecutive API call */
  streamBuffer?: number;
  /** Last stream call timestamp for rate limiting */
  lastStreamCall?: number;
  /** Tools available to this agent */
  tools?: t.GraphTools;
  /** Graph-managed tools (e.g., handoff tools created by MultiAgentGraph) that bypass event-driven dispatch */
  graphTools?: t.GraphTools;
  /** Tool map for this agent */
  toolMap?: t.ToolMap;
  /**
   * Tool definitions registry (includes deferred and programmatic tool metadata).
   * Used for tool search and programmatic tool calling.
   */
  toolRegistry?: t.LCToolRegistry;
  /**
   * Serializable tool definitions for event-driven execution.
   * When provided, ToolNode operates in event-driven mode.
   */
  toolDefinitions?: t.LCTool[];
  /** Set of tool names discovered via tool search (to be loaded) */
  discoveredToolNames: Set<string> = new Set();
  /** Instructions for this agent */
  instructions?: string;
  /** Additional instructions for this agent */
  additionalInstructions?: string;
  /**
   * Dynamic context that changes per-request (e.g., current time, user info).
   * This is NOT included in the system message to preserve cache.
   * Instead, it's injected as a user message at the start of the conversation.
   */
  dynamicContext?: string;
  /** Reasoning key for this agent */
  reasoningKey: 'reasoning_content' | 'reasoning' = 'reasoning_content';
  /** Last token for reasoning detection */
  lastToken?: string;
  /** Token type switch state */
  tokenTypeSwitch?: 'reasoning' | 'content';
  /** Tracks how many reasoning→text transitions have occurred (ensures unique post-reasoning step keys) */
  reasoningTransitionCount = 0;
  /** Current token type being processed */
  currentTokenType: ContentTypes.TEXT | ContentTypes.THINK | 'think_and_text' =
    ContentTypes.TEXT;
  /** Whether tools should end the workflow */
  toolEnd: boolean = false;
  /** Cached system runnable (created lazily) */
  private cachedSystemRunnable?: Runnable<
    BaseMessage[],
    (BaseMessage | SystemMessage)[],
    RunnableConfig<Record<string, unknown>>
  >;
  /** Whether system runnable needs rebuild (set when discovered tools change) */
  private systemRunnableStale: boolean = true;
  /** Cached system message token count (separate from tool tokens) */
  private systemMessageTokens: number = 0;
  /** Promise for token calculation initialization */
  tokenCalculationPromise?: Promise<void>;
  /** Format content blocks as strings (for legacy compatibility) */
  useLegacyContent: boolean = false;
  /** Detailed per-tool token breakdown for admin tracking */
  private toolsDetail: Array<{ name: string; tokens: number }> = [];
  /** Total tool tokens (sum of all toolsDetail) */
  private toolTokensTotal: number = 0;
  /** Per-prompt token breakdown for detailed admin reporting */
  private promptBreakdown: {
    branding: number;
    toolRouting: number;
    agentInstructions: number;
    mcpInstructions: number;
    artifacts: number;
    memory: number;
  } = {
    branding: 0,
    toolRouting: 0,
    agentInstructions: 0,
    mcpInstructions: 0,
    artifacts: 0,
    memory: 0,
  };
  /**
   * Handoff context when this agent receives control via handoff.
   * Contains source and parallel execution info for system message context.
   */
  handoffContext?: {
    /** Source agent that transferred control */
    sourceAgentName: string;
    /** Names of sibling agents executing in parallel (empty if sequential) */
    parallelSiblings: string[];
  };
  /**
   * Structured output configuration.
   * When set, the agent will return a validated JSON response
   * instead of streaming text.
   */
  structuredOutput?: t.StructuredOutputConfig;
  /** Optional callback to summarize discarded messages during context pruning */
  summarizeCallback?: (messages: BaseMessage[]) => Promise<string | undefined>;
  /** Pre-existing summary loaded from persistent storage, injected into context on new turns */
  persistedSummary?: string;
  /** Summarization configuration controlling trigger strategy, reserve ratio, and EMA calibration */
  summarizationConfig?: t.SummarizationConfig;
  /** Lightweight file manifest for file-aware compaction (IDs and names only, no content) */
  fileManifest?: t.FileManifestEntry[];
  /**
   * Workspace-shared system-message tiers. When set, each entry becomes a
   * separate text block in the SystemMessage with its own cachePoint /
   * cache_control marker, BEFORE the per-agent `instructions` block.
   * See {@link t.AgentInputs.system_cache_blocks} for full semantics.
   */
  systemCacheBlocks?: Array<{ text: string; ttl?: '5m' | '1h' }>;
  /** TTL hint for the per-agent instructions cache marker. Defaults to '5m'. */
  instructionsCacheTtl?: '5m' | '1h';
  /** Original AgentInputs used to create this context — used for self-spawn subagent resolution. */
  _sourceInputs?: t.AgentInputs;
  /** Subagent configurations for hierarchical delegation. */
  subagentConfigs?: t.SubagentConfig[];
  /** Maximum subagent nesting depth. */
  maxSubagentDepth?: number;

  constructor({
    agentId,
    name,
    description,
    provider,
    clientOptions,
    maxContextTokens,
    streamBuffer,
    tokenCounter,
    tools,
    toolMap,
    toolRegistry,
    toolDefinitions,
    instructions,
    additionalInstructions,
    dynamicContext,
    reasoningKey,
    toolEnd,
    instructionTokens,
    useLegacyContent,
    structuredOutput,
    discoveredTools,
    summarizeCallback,
    persistedSummary,
    summarizationConfig,
    fileManifest,
    systemCacheBlocks,
    instructionsCacheTtl,
  }: {
    agentId: string;
    name?: string;
    description?: string;
    provider: Providers;
    clientOptions?: t.ClientOptions;
    maxContextTokens?: number;
    streamBuffer?: number;
    tokenCounter?: t.TokenCounter;
    tools?: t.GraphTools;
    toolMap?: t.ToolMap;
    toolRegistry?: t.LCToolRegistry;
    toolDefinitions?: t.LCTool[];
    instructions?: string;
    additionalInstructions?: string;
    dynamicContext?: string;
    reasoningKey?: 'reasoning_content' | 'reasoning';
    toolEnd?: boolean;
    instructionTokens?: number;
    useLegacyContent?: boolean;
    structuredOutput?: t.StructuredOutputConfig;
    discoveredTools?: string[];
    summarizeCallback?: (
      messages: BaseMessage[]
    ) => Promise<string | undefined>;
    persistedSummary?: string;
    summarizationConfig?: t.SummarizationConfig;
    fileManifest?: t.FileManifestEntry[];
    systemCacheBlocks?: Array<{ text: string; ttl?: '5m' | '1h' }>;
    instructionsCacheTtl?: '5m' | '1h';
  }) {
    this.agentId = agentId;
    this.name = name;
    this.description = description;
    this.provider = provider;
    this.clientOptions = clientOptions;
    this.maxContextTokens = maxContextTokens;
    this.streamBuffer = streamBuffer;
    this.tokenCounter = tokenCounter;
    this.tools = tools;
    this.toolMap = toolMap;
    this.toolRegistry = toolRegistry;
    this.toolDefinitions = toolDefinitions;
    this.instructions = instructions;
    this.additionalInstructions = additionalInstructions;
    this.dynamicContext = dynamicContext;
    this.structuredOutput = structuredOutput;
    this.summarizeCallback = summarizeCallback;
    this.persistedSummary = persistedSummary;
    this.summarizationConfig = summarizationConfig;
    this.fileManifest = fileManifest;
    if (systemCacheBlocks && systemCacheBlocks.length > 0) {
      this.systemCacheBlocks = systemCacheBlocks;
    }
    if (instructionsCacheTtl) {
      this.instructionsCacheTtl = instructionsCacheTtl;
    }
    if (reasoningKey) {
      this.reasoningKey = reasoningKey;
    }
    if (toolEnd !== undefined) {
      this.toolEnd = toolEnd;
    }
    if (instructionTokens !== undefined) {
      this.instructionTokens = instructionTokens;
    }

    this.useLegacyContent = useLegacyContent ?? false;

    if (discoveredTools && discoveredTools.length > 0) {
      for (const toolName of discoveredTools) {
        this.discoveredToolNames.add(toolName);
      }
    }
  }

  /**
   * Checks if structured output mode is enabled for this agent.
   * When enabled, the agent will use model.invoke() instead of streaming
   * and return a validated JSON response.
   */
  get isStructuredOutputMode(): boolean {
    // Runtime safety: schema can be null/undefined via API despite type saying required
    return (
      this.structuredOutput != null &&
      // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
      this.structuredOutput.schema != null
    );
  }

  /**
   * Gets the structured output schema with normalized defaults.
   * Returns undefined if structured output is not configured.
   */
  getStructuredOutputSchema(): Record<string, unknown> | undefined {
    if (!this.structuredOutput?.schema) {
      return undefined;
    }

    const schema = { ...this.structuredOutput.schema };

    // Ensure type is set
    if (schema.type == null && schema.properties != null) {
      schema.type = 'object';
    }

    // Add title from config name
    if (
      this.structuredOutput.name != null &&
      this.structuredOutput.name !== '' &&
      schema.title == null
    ) {
      schema.title = this.structuredOutput.name;
    }

    // Add description from config
    if (
      this.structuredOutput.description != null &&
      this.structuredOutput.description !== '' &&
      schema.description == null
    ) {
      schema.description = this.structuredOutput.description;
    }

    // Enable strict mode by default
    if (this.structuredOutput.strict !== false && schema.type === 'object') {
      schema.additionalProperties = schema.additionalProperties ?? false;
    }

    return schema;
  }

  /**
   * Resolves the structured output mode to a concrete method based on provider capabilities.
   *
   * Resolution logic:
   * - 'native' or 'auto' + supported provider → native constrained decoding
   * - 'native' + unsupported provider → fallback to 'functionCalling' with warning
   * - 'provider' → LangChain's jsonMode (existing behavior)
   * - 'tool' → function calling trick (existing behavior)
   * - 'auto' + unsupported provider → 'functionCalling'
   *
   * @returns The resolved method for withStructuredOutput, or 'native' for direct API usage
   */
  resolveStructuredOutputMode(): {
    method: t.ResolvedStructuredOutputMethod;
    warnings: string[];
  } {
    const mode = this.structuredOutput?.mode ?? 'auto';
    const warnings: string[] = [];

    // Providers that support native constrained decoding via LangChain
    const nativeProviders = new Set([
      Providers.ANTHROPIC,
      Providers.OPENAI,
      Providers.AZURE,
    ]);

    // Providers where LangChain supports jsonMode
    const jsonModeProviders = new Set([
      Providers.ANTHROPIC,
      Providers.OPENAI,
      Providers.AZURE,
    ]);

    switch (mode) {
      case 'native': {
        if (nativeProviders.has(this.provider)) {
          if (this.provider === Providers.ANTHROPIC) {
            return { method: 'jsonSchema', warnings };
          }
          // OpenAI/Azure
          return { method: 'jsonSchema', warnings };
        }
        // Fallback for unsupported providers
        warnings.push(
          `Native structured output is not supported for provider '${this.provider}'. Falling back to function calling.`
        );
        return { method: 'functionCalling', warnings };
      }

      case 'auto': {
        if (nativeProviders.has(this.provider)) {
          if (this.provider === Providers.ANTHROPIC) {
            return { method: 'jsonSchema', warnings };
          }
          // OpenAI/Azure
          return { method: 'jsonSchema', warnings };
        }
        // Default to function calling for all other providers
        return { method: undefined, warnings };
      }

      case 'provider': {
        if (this.provider === Providers.BEDROCK) {
          // Bedrock doesn't support jsonMode, fall back to functionCalling
          return { method: 'functionCalling', warnings };
        }
        if (jsonModeProviders.has(this.provider)) {
          return { method: 'jsonMode', warnings };
        }
        return { method: 'jsonMode', warnings };
      }

      case 'tool': {
        return { method: 'functionCalling', warnings };
      }

      default: {
        return { method: undefined, warnings };
      }
    }
  }

  /**
   * Builds instructions text for tools that are ONLY callable via programmatic code execution.
   * These tools cannot be called directly by the LLM but are available through the
   * run_tools_with_code tool.
   *
   * Includes:
   * - Code_execution-only tools that are NOT deferred
   * - Code_execution-only tools that ARE deferred but have been discovered via tool search
   */
  private buildProgrammaticOnlyToolsInstructions(): string {
    if (!this.toolRegistry) return '';

    const programmaticOnlyTools: t.LCTool[] = [];
    for (const [name, toolDef] of this.toolRegistry) {
      const allowedCallers = toolDef.allowed_callers ?? ['direct'];
      const isCodeExecutionOnly =
        allowedCallers.includes('code_execution') &&
        !allowedCallers.includes('direct');

      if (!isCodeExecutionOnly) continue;

      // Include if: not deferred OR deferred but discovered
      const isDeferred = toolDef.defer_loading === true;
      const isDiscovered = this.discoveredToolNames.has(name);
      if (!isDeferred || isDiscovered) {
        programmaticOnlyTools.push(toolDef);
      }
    }

    if (programmaticOnlyTools.length === 0) return '';

    const toolDescriptions = programmaticOnlyTools
      .map((tool) => {
        let desc = `- **${tool.name}**`;
        if (tool.description != null && tool.description !== '') {
          desc += `: ${tool.description}`;
        }
        if (tool.parameters) {
          desc += `\n  Parameters: ${JSON.stringify(tool.parameters, null, 2).replace(/\n/g, '\n  ')}`;
        }
        return desc;
      })
      .join('\n\n');

    return (
      '\n\n## Programmatic-Only Tools\n\n' +
      'The following tools are available exclusively through the `run_tools_with_code` tool. ' +
      'You cannot call these tools directly; instead, use `run_tools_with_code` with Python code that invokes them.\n\n' +
      toolDescriptions
    );
  }

  /**
   * Gets the system runnable, creating it lazily if needed.
   * Includes instructions, additional instructions, and programmatic-only tools documentation.
   * Only rebuilds when marked stale (via markToolsAsDiscovered).
   */
  get systemRunnable():
    | Runnable<
        BaseMessage[],
        (BaseMessage | SystemMessage)[],
        RunnableConfig<Record<string, unknown>>
      >
    | undefined {
    // Return cached if not stale
    if (!this.systemRunnableStale && this.cachedSystemRunnable !== undefined) {
      return this.cachedSystemRunnable;
    }

    // Stale or first access - rebuild
    const instructionsString = this.buildInstructionsString();
    this.cachedSystemRunnable = this.buildSystemRunnable(instructionsString);
    this.systemRunnableStale = false;
    return this.cachedSystemRunnable;
  }

  /**
   * Explicitly initializes the system runnable.
   * Call this before async token calculation to ensure system message tokens are counted first.
   */
  initializeSystemRunnable(): void {
    if (this.systemRunnableStale || this.cachedSystemRunnable === undefined) {
      const instructionsString = this.buildInstructionsString();
      this.cachedSystemRunnable = this.buildSystemRunnable(instructionsString);
      this.systemRunnableStale = false;
    }
  }

  /**
   * Builds the raw instructions string (without creating SystemMessage).
   * Includes agent identity preamble and handoff context when available.
   */
  private buildInstructionsString(): string {
    const parts: string[] = [];

    /** Build agent identity and handoff context preamble */
    const identityPreamble = this.buildIdentityPreamble();
    if (identityPreamble) {
      parts.push(identityPreamble);
    }

    /** Add main instructions */
    if (this.instructions != null && this.instructions !== '') {
      parts.push(this.instructions);
    }

    /** Add additional instructions */
    if (
      this.additionalInstructions != null &&
      this.additionalInstructions !== ''
    ) {
      parts.push(this.additionalInstructions);
    }

    /** Add programmatic tools documentation */
    const programmaticToolsDoc = this.buildProgrammaticOnlyToolsInstructions();
    if (programmaticToolsDoc) {
      parts.push(programmaticToolsDoc);
    }

    return parts.join('\n\n');
  }

  /**
   * Builds the agent identity preamble including handoff context if present.
   * This helps the agent understand its role in the multi-agent workflow.
   */
  private buildIdentityPreamble(): string {
    if (!this.handoffContext) return '';

    const displayName = this.name ?? this.agentId;
    const { sourceAgentName, parallelSiblings } = this.handoffContext;
    const isParallel = parallelSiblings.length > 0;

    const lines: string[] = [];
    lines.push('## Subagent Context');
    lines.push('');
    lines.push(
      `You are "${displayName}", a subagent spawned by "${sourceAgentName}" for a specific task.`
    );

    if (isParallel) {
      lines.push(`Running in parallel with: ${parallelSiblings.join(', ')}.`);
    }

    lines.push('');
    lines.push('### Your Role');
    lines.push('- Complete your assigned task. That is your entire purpose.');
    lines.push(`- You are NOT "${sourceAgentName}". Do not try to be.`);
    lines.push('');
    lines.push('### Rules');
    lines.push('1. **Stay focused** — Do your assigned task, nothing else.');
    lines.push(
      '2. **Complete the task** — Your final message will be automatically reported back.'
    );
    lines.push(
      '3. **Be autonomous** — Execute directly without asking for user confirmation. Use your tools and best judgment.'
    );
    lines.push(
      '4. **No placeholders** — Never generate fake or placeholder data. If you cannot retrieve real data, say so.'
    );
    lines.push(
      '5. **No side effects** — Do not send messages, emails, or notifications unless explicitly tasked to do so.'
    );
    lines.push('');
    lines.push('### Output');
    lines.push('When complete, your final response should include:');
    lines.push('- What you accomplished or found');
    lines.push('- Any relevant details the orchestrator should know');
    lines.push('- Keep it concise but informative');

    return lines.join('\n');
  }

  /**
   * Build system runnable from pre-built instructions string.
   * Only called when content has actually changed.
   */
  private buildSystemRunnable(
    instructionsString: string
  ):
    | Runnable<
        BaseMessage[],
        (BaseMessage | SystemMessage)[],
        RunnableConfig<Record<string, unknown>>
      >
    | undefined {
    if (!instructionsString) {
      // Remove previous tokens if we had a system message before
      this.instructionTokens -= this.systemMessageTokens;
      this.systemMessageTokens = 0;
      return undefined;
    }

    let finalInstructions: string | BaseMessageFields = instructionsString;

    /**
     * Tiered system-message assembly. When `systemCacheBlocks` is set, build
     * a multi-block SystemMessage so each tier has its own cache breakpoint:
     *
     *   [tier1_block_1][tier1_cachePoint]...[tier1_block_N][tier1_cachePoint]
     *   [instructions][instructions_cachePoint]
     *
     * Forward-prefix-hash means agents in the same workspace whose tier-1
     * bytes are identical share the platform cache entry; only the per-agent
     * `instructions` block invalidates per-agent.
     *
     * - Bedrock: emit `cachePoint: { type: 'default' }` blocks (handled by
     *   `convertSystemMessageToConverseMessage`).
     * - Anthropic: emit `cache_control: { type: 'ephemeral' }` on each text
     *   block. (TTL hints are dropped for Anthropic — the SDK currently
     *   only supports `'ephemeral'`. Bedrock cachePoint has no TTL knob.)
     */
    const hasTieredCache =
      Array.isArray(this.systemCacheBlocks) &&
      this.systemCacheBlocks.length > 0;
    const isBedrock = this.provider === Providers.BEDROCK;
    const isAnthropic = this.provider === Providers.ANTHROPIC;
    const anthropicCacheEnabled =
      isAnthropic &&
      (this.clientOptions as t.AnthropicClientOptions | undefined)
        ?.promptCache === true;
    /**
     * Bedrock cachePoint is Claude-only — Nova/Llama/Titan reject it.
     * Mirrors the model check in `IllumaBedrockConverse.invocationParams`
     * (src/llm/bedrock/index.ts:186-189).
     */
    const bedrockCacheEnabled =
      isBedrock &&
      (() => {
        const opts = this.clientOptions as
          | (t.BedrockAnthropicClientOptions & { model?: string })
          | undefined;
        const modelId = (opts?.model ?? '').toLowerCase();
        const isClaudeModel =
          modelId.includes('claude') || modelId.includes('anthropic');
        return opts?.promptCache === true && isClaudeModel;
      })();

    if (hasTieredCache && (bedrockCacheEnabled || anthropicCacheEnabled)) {
      /**
       * Anthropic / Bedrock cap cache breakpoints at 4 per request. The
       * lib already emits one for tools and up to two for messages, so we
       * have at most 1 left for the system block. We spend it on Tier 1
       * (the workspace-shared bytes) — per-agent Tier 2 caching is still
       * achieved implicitly via the tools breakpoint that follows, since
       * cache lookups are forward-prefix-hash.
       *
       * Tier 1 may itself be emitted as multiple text blocks (one per
       * `systemCacheBlocks` entry); only the LAST gets the cache marker,
       * the rest are plain text inside the same cached prefix.
       */
      const contentBlocks: Array<Record<string, unknown>> = [];
      const tier1Blocks = this.systemCacheBlocks!.filter((b) => b?.text);
      tier1Blocks.forEach((block, idx) => {
        const isLast = idx === tier1Blocks.length - 1;
        if (bedrockCacheEnabled) {
          contentBlocks.push({ type: 'text', text: block.text });
          if (isLast) {
            contentBlocks.push({ cachePoint: { type: 'default' } });
          }
        } else if (isLast) {
          contentBlocks.push({
            type: 'text',
            text: block.text,
            cache_control: { type: 'ephemeral' },
          });
        } else {
          contentBlocks.push({ type: 'text', text: block.text });
        }
      });
      if (instructionsString) {
        // No cache marker on the trailing per-agent block — tools'
        // breakpoint covers it.
        contentBlocks.push({ type: 'text', text: instructionsString });
      }
      finalInstructions = {
        content: contentBlocks as unknown as BaseMessageFields['content'],
      };
    } else if (anthropicCacheEnabled) {
      // Legacy single-block Anthropic caching (preserved for back-compat).
      finalInstructions = {
        content: [
          {
            type: 'text',
            text: instructionsString,
            cache_control: { type: 'ephemeral' },
          },
        ],
      };
    }

    const systemMessage = new SystemMessage(finalInstructions);

    // Update token counts (subtract old, add new)
    if (this.tokenCounter) {
      this.instructionTokens -= this.systemMessageTokens;
      this.systemMessageTokens = this.tokenCounter(systemMessage);
      this.instructionTokens += this.systemMessageTokens;
    }

    return RunnableLambda.from((messages: BaseMessage[]) => {
      return [systemMessage, ...messages];
    }).withConfig({ runName: 'prompt' });
  }

  /**
   * Reset context for a new run
   */
  reset(): void {
    this.instructionTokens = 0;
    this.systemMessageTokens = 0;
    this.toolsDetail = [];
    this.toolTokensTotal = 0;
    this.cachedSystemRunnable = undefined;
    this.systemRunnableStale = true;
    this.lastToken = undefined;
    this.indexTokenCountMap = { ...this.baseIndexTokenCountMap };
    this.currentUsage = undefined;
    this.pruneMessages = undefined;
    this.lastStreamCall = undefined;
    this.tokenTypeSwitch = undefined;
    this.reasoningTransitionCount = 0;
    this.currentTokenType = ContentTypes.TEXT;
    this.discoveredToolNames.clear();
    this.handoffContext = undefined;

    if (this.tokenCounter) {
      this.initializeSystemRunnable();
      const baseTokenMap = { ...this.baseIndexTokenCountMap };
      this.indexTokenCountMap = baseTokenMap;
      this.tokenCalculationPromise = this.calculateInstructionTokens(
        this.tokenCounter
      )
        .then(() => {
          this.updateTokenMapWithInstructions(baseTokenMap);
        })
        .catch((err) => {
          console.error('Error calculating instruction tokens:', err);
        });
    } else {
      this.tokenCalculationPromise = undefined;
    }
  }

  /**
   * Update the token count map with instruction tokens
   */
  updateTokenMapWithInstructions(baseTokenMap: Record<string, number>): void {
    if (this.instructionTokens > 0) {
      // Shift all indices by the instruction token count
      const shiftedMap: Record<string, number> = {};
      for (const [key, value] of Object.entries(baseTokenMap)) {
        const index = parseInt(key, 10);
        if (!isNaN(index)) {
          shiftedMap[String(index)] =
            value + (index === 0 ? this.instructionTokens : 0);
        }
      }
      this.indexTokenCountMap = shiftedMap;
    } else {
      this.indexTokenCountMap = { ...baseTokenMap };
    }
  }

  /**
   * Calculate tool tokens and add to instruction tokens
   * Note: System message tokens are calculated during systemRunnable creation
   * Also tracks per-tool token breakdown for admin reporting
   */
  async calculateInstructionTokens(
    tokenCounter: t.TokenCounter
  ): Promise<void> {
    let toolTokens = 0;
    // Track names to avoid double-counting when a tool appears in both
    // this.tools (bound StructuredTool instances) and this.toolDefinitions
    // (MCP / event-driven schemas).
    const countedToolNames = new Set<string>();

    // Reset per-tool breakdown
    this.toolsDetail = [];

    // Count tokens for bound tools (StructuredTool instances with .schema)
    if (this.tools && this.tools.length > 0) {
      for (const tool of this.tools) {
        const genericTool = tool as Record<string, unknown>;
        if (
          genericTool.schema != null &&
          typeof genericTool.schema === 'object'
        ) {
          const toolName = (genericTool.name as string | undefined) ?? '';
          const jsonSchema = toJsonSchema(
            genericTool.schema,
            toolName,
            (genericTool.description as string | undefined) ?? ''
          );
          const tokens = tokenCounter(
            new SystemMessage(JSON.stringify(jsonSchema))
          );
          if (toolName) {
            countedToolNames.add(toolName);
          }

          // Track per-tool breakdown
          this.toolsDetail.push({ name: toolName || 'unknown', tokens });
          toolTokens += tokens;
        }
      }
    }

    // Count tokens for tool definitions (MCP / event-driven tools).
    // These are sent to the provider API as tool schemas alongside bound tools.
    // Both can be populated simultaneously (graph tools + MCP tools).
    if (this.toolDefinitions && this.toolDefinitions.length > 0) {
      for (const def of this.toolDefinitions) {
        if (countedToolNames.has(def.name)) {
          continue; // Already counted via this.tools
        }
        const schema = {
          name: def.name,
          description: def.description ?? '',
          parameters: def.parameters ?? {},
        };
        const defTokens = tokenCounter(
          new SystemMessage(JSON.stringify(schema))
        );
        this.toolsDetail.push({
          name: def.name || 'unknown',
          tokens: defTokens,
        });
        toolTokens += defTokens;
      }
    }

    // Store total tool tokens for breakdown reporting
    this.toolTokensTotal = toolTokens;

    // Add tool tokens to existing instruction tokens (which may already include system message tokens)
    this.instructionTokens += toolTokens;
  }

  /**
   * Set the per-prompt token breakdown for detailed admin tracking.
   * Called by the client after assembling all prompt components.
   * @param breakdown - Object with token counts per prompt component
   */
  setPromptBreakdown(breakdown: {
    branding?: number;
    toolRouting?: number;
    agentInstructions?: number;
    mcpInstructions?: number;
    artifacts?: number;
    memory?: number;
  }): void {
    if (breakdown.branding !== undefined)
      this.promptBreakdown.branding = breakdown.branding;
    if (breakdown.toolRouting !== undefined)
      this.promptBreakdown.toolRouting = breakdown.toolRouting;
    if (breakdown.agentInstructions !== undefined)
      this.promptBreakdown.agentInstructions = breakdown.agentInstructions;
    if (breakdown.mcpInstructions !== undefined)
      this.promptBreakdown.mcpInstructions = breakdown.mcpInstructions;
    if (breakdown.artifacts !== undefined)
      this.promptBreakdown.artifacts = breakdown.artifacts;
    if (breakdown.memory !== undefined)
      this.promptBreakdown.memory = breakdown.memory;
  }

  /**
   * Get a detailed breakdown of context tokens for admin reporting.
   * This provides visibility into what's consuming the input token budget.
   * @returns ContextBreakdown object with per-component token counts
   */
  getContextBreakdown(): {
    instructions: number;
    artifacts: number;
    tools: number;
    toolCount: number;
    toolContext: number;
    total: number;
    toolsDetail: Array<{ name: string; tokens: number }>;
    toolContextDetail: Array<{ name: string; tokens: number }>;
    prompts?: {
      branding: number;
      toolRouting: number;
      agentInstructions: number;
      mcpInstructions: number;
      artifacts: number;
      memory: number;
    };
  } {
    // Calculate sum of prompt components
    const promptsTotal =
      this.promptBreakdown.branding +
      this.promptBreakdown.toolRouting +
      this.promptBreakdown.agentInstructions +
      this.promptBreakdown.mcpInstructions +
      this.promptBreakdown.artifacts +
      this.promptBreakdown.memory;

    return {
      // System message tokens (instructions + additional_instructions)
      instructions: this.systemMessageTokens,
      // Artifacts prompt tokens
      artifacts: this.promptBreakdown.artifacts,
      // Total tool definition tokens
      tools: this.toolTokensTotal,
      // Number of tools
      toolCount: this.toolsDetail.length,
      // Tool context/usage instructions (currently embedded in system message)
      toolContext: 0,
      // Total tracked context tokens
      total: this.instructionTokens,
      // Per-tool token breakdown
      toolsDetail: [...this.toolsDetail],
      // Tool context detail (currently not tracked separately)
      toolContextDetail: [],
      // Per-prompt breakdown (only include if any prompts were tracked)
      prompts: promptsTotal > 0 ? { ...this.promptBreakdown } : undefined,
    };
  }

  /**
   * Gets the tool registry for deferred tools (for tool search).
   * @param onlyDeferred If true, only returns tools with defer_loading=true
   * @returns LCToolRegistry with tool definitions
   */
  getDeferredToolRegistry(onlyDeferred: boolean = true): t.LCToolRegistry {
    const registry: t.LCToolRegistry = new Map();

    if (!this.toolRegistry) {
      return registry;
    }

    for (const [name, toolDef] of this.toolRegistry) {
      if (!onlyDeferred || toolDef.defer_loading === true) {
        registry.set(name, toolDef);
      }
    }

    return registry;
  }

  /**
   * Sets the handoff context for this agent.
   * Call this when the agent receives control via handoff from another agent.
   * Marks system runnable as stale to include handoff context in system message.
   * @param sourceAgentName - Name of the agent that transferred control
   * @param parallelSiblings - Names of other agents executing in parallel with this one
   */
  setHandoffContext(sourceAgentName: string, parallelSiblings: string[]): void {
    this.handoffContext = { sourceAgentName, parallelSiblings };
    this.systemRunnableStale = true;
  }

  /**
   * Clears any handoff context.
   * Call this when resetting the agent or when handoff context is no longer relevant.
   */
  clearHandoffContext(): void {
    if (this.handoffContext) {
      this.handoffContext = undefined;
      this.systemRunnableStale = true;
    }
  }

  /**
   * Marks tools as discovered via tool search.
   * Discovered tools will be included in the next model binding.
   * Only marks system runnable stale if NEW tools were actually added.
   * @param toolNames - Array of discovered tool names
   * @returns true if any new tools were discovered
   */
  markToolsAsDiscovered(toolNames: string[]): boolean {
    let hasNewDiscoveries = false;
    for (const name of toolNames) {
      if (!this.discoveredToolNames.has(name)) {
        this.discoveredToolNames.add(name);
        hasNewDiscoveries = true;
      }
    }
    if (hasNewDiscoveries) {
      this.systemRunnableStale = true;
    }
    return hasNewDiscoveries;
  }

  /**
   * Gets tools that should be bound to the LLM.
   * In event-driven mode (toolDefinitions present, tools empty), creates schema-only tools.
   * Otherwise filters tool instances based on:
   * 1. Non-deferred tools with allowed_callers: ['direct']
   * 2. Discovered tools (from tool search)
   * @returns Array of tools to bind to model
   */
  getToolsForBinding(): t.GraphTools | undefined {
    /** Event-driven mode: create schema-only tools from definitions */
    if (this.toolDefinitions && this.toolDefinitions.length > 0) {
      return this.getEventDrivenToolsForBinding();
    }

    /** Traditional mode: filter actual tool instances */
    const filtered =
      !this.tools || !this.toolRegistry
        ? this.tools
        : this.filterToolsForBinding(this.tools);

    if (this.graphTools && this.graphTools.length > 0) {
      return [...(filtered ?? []), ...this.graphTools];
    }

    return filtered;
  }

  /** Creates schema-only tools from toolDefinitions for event-driven mode, merged with native tools */
  private getEventDrivenToolsForBinding(): t.GraphTools {
    if (!this.toolDefinitions) {
      return this.graphTools ?? [];
    }

    const defsToInclude = this.toolDefinitions.filter((def) => {
      const allowedCallers = def.allowed_callers ?? ['direct'];
      if (!allowedCallers.includes('direct')) {
        return false;
      }
      if (
        def.defer_loading === true &&
        !this.discoveredToolNames.has(def.name)
      ) {
        return false;
      }
      return true;
    });

    const schemaTools = createSchemaOnlyTools(defsToInclude) as t.GraphTools;

    const allTools = [...schemaTools];

    if (this.graphTools && this.graphTools.length > 0) {
      allTools.push(...this.graphTools);
    }

    if (this.tools && this.tools.length > 0) {
      allTools.push(...this.tools);
    }

    return allTools;
  }

  /** Filters tool instances for binding based on registry config */
  private filterToolsForBinding(tools: t.GraphTools): t.GraphTools {
    return tools.filter((tool) => {
      if (!('name' in tool)) {
        return true;
      }

      const toolDef = this.toolRegistry?.get(tool.name);
      if (!toolDef) {
        return true;
      }

      if (this.discoveredToolNames.has(tool.name)) {
        const allowedCallers = toolDef.allowed_callers ?? ['direct'];
        return allowedCallers.includes('direct');
      }

      const allowedCallers = toolDef.allowed_callers ?? ['direct'];
      return (
        allowedCallers.includes('direct') && toolDef.defer_loading !== true
      );
    });
  }
}