/**
 * Optimized ChatBedrockConverse wrapper that fixes content block merging for
 * streaming responses and adds support for latest @langchain/aws features:
 *
 * - Prompt caching support for Bedrock Converse API (Illuma feature)
 * - Application Inference Profiles (PR #9129)
 * - Service Tiers (Priority/Standard/Flex) (PR #9785) - requires AWS SDK 3.966.0+
 *
 * Bedrock's `@langchain/aws` library does not include an `index` property on content
 * blocks (unlike Anthropic/OpenAI), which causes LangChain's `_mergeLists` to append
 * each streaming chunk as a separate array entry instead of merging by index.
 *
 * This wrapper takes full ownership of the stream by directly interfacing with the
 * AWS SDK client (`this.client`) and using custom handlers from `./utils/` that
 * include `contentBlockIndex` in response_metadata for every delta type. It then
 * promotes `contentBlockIndex` to an `index` property on each content block
 * (mirroring Anthropic's pattern) and strips it from metadata to avoid
 * `_mergeDicts` conflicts.
 *
 * When multiple content block types are present (e.g. reasoning + text), text deltas
 * are promoted from strings to array form with `index` so they merge correctly once
 * the accumulated content is already an array.
 *
 * PROMPT CACHING:
 * When promptCache: true is set, this wrapper adds cachePoint markers to the tools array
 * to enable Bedrock prompt caching for tool definitions. This allows tool schemas to be
 * cached and reused across requests, reducing latency and costs.
 *
 * CACHE TOKEN EXTRACTION:
 * Cache token extraction is handled in `./utils/message_outputs.ts` and added
 * to usage_metadata with input_token_details (cacheReadInputTokens/cacheWriteInputTokens).
 */

import { ChatBedrockConverse } from '@langchain/aws';
import { ConverseStreamCommand } from '@aws-sdk/client-bedrock-runtime';
import { AIMessageChunk } from '@langchain/core/messages';
import type { BaseMessage } from '@langchain/core/messages';
import { ChatGenerationChunk, ChatResult } from '@langchain/core/outputs';
import type { CallbackManagerForLLMRun } from '@langchain/core/callbacks/manager';
import type { ChatBedrockConverseInput } from '@langchain/aws';
import {
  convertToConverseMessages,
  handleConverseStreamContentBlockStart,
  handleConverseStreamContentBlockDelta,
  handleConverseStreamMetadata,
} from './utils';

/**
 * Service tier type for Bedrock invocations.
 * Requires AWS SDK >= 3.966.0 to actually work.
 * @see https://docs.aws.amazon.com/bedrock/latest/userguide/service-tiers-inference.html
 */
export type ServiceTierType = 'priority' | 'default' | 'flex' | 'reserved';

/**
 * Extended input interface with additional features:
 * - promptCache: Enable Bedrock prompt caching for tool definitions
 * - applicationInferenceProfile: Use an inference profile ARN instead of model ID
 * - serviceTier: Specify service tier (Priority, Standard, Flex, Reserved)
 */
export interface IllumaBedrockConverseInput extends ChatBedrockConverseInput {
  /**
   * Enable Bedrock prompt caching for tool definitions.
   * When true, adds cachePoint markers to tools array.
   */
  promptCache?: boolean;

  /**
   * Application Inference Profile ARN to use for the model.
   * For example, "arn:aws:bedrock:eu-west-1:123456789102:application-inference-profile/fm16bt65tzgx"
   * When provided, this ARN will be used for the actual inference calls instead of the model ID.
   * Must still provide `model` as normal modelId to benefit from all the metadata.
   * @see https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-create.html
   */
  applicationInferenceProfile?: string;

  /**
   * Service tier for model invocation.
   * Specifies the processing tier type used for serving the request.
   * Supported values are 'priority', 'default', 'flex', and 'reserved'.
   *
   * - 'priority': Prioritized processing for lower latency
   * - 'default': Standard processing tier
   * - 'flex': Flexible processing tier with lower cost
   * - 'reserved': Reserved capacity for consistent performance
   *
   * If not provided, AWS uses the default tier.
   * Note: Requires AWS SDK >= 3.966.0 to work.
   * @see https://docs.aws.amazon.com/bedrock/latest/userguide/service-tiers-inference.html
   */
  serviceTier?: ServiceTierType;
}

/**
 * Extended call options with serviceTier override support.
 */
export interface IllumaBedrockConverseCallOptions {
  serviceTier?: ServiceTierType;
}

export class IllumaBedrockConverse extends ChatBedrockConverse {
  /** Enable Bedrock prompt caching for tool definitions */
  promptCache: boolean;

  /** Application Inference Profile ARN to use instead of model ID */
  applicationInferenceProfile?: string;

  /** Service tier for model invocation */
  serviceTier?: ServiceTierType;

  constructor(fields?: IllumaBedrockConverseInput) {
    super(fields);
    this.promptCache = fields?.promptCache ?? false;
    this.applicationInferenceProfile = fields?.applicationInferenceProfile;
    this.serviceTier = fields?.serviceTier;

    // Fix: Force supportsToolChoiceValues for Claude models
    // The parent constructor checks `model.includes('claude-3')` but this fails when:
    // 1. Using applicationInferenceProfile ARNs (arn:aws:bedrock:...)
    // 2. Using different naming conventions (claude-4, claude-opus-4, etc.)
    // We need to ensure tool_choice is properly set for withStructuredOutput to work
    const modelName = (fields?.model ?? '').toLowerCase();
    const profileName = (
      fields?.applicationInferenceProfile ?? ''
    ).toLowerCase();
    const isClaudeModel =
      modelName.includes('claude') ||
      modelName.includes('anthropic') ||
      profileName.includes('claude') ||
      profileName.includes('anthropic');

    if (isClaudeModel && !this.supportsToolChoiceValues?.length) {
      // Claude models support all tool choice values
      this.supportsToolChoiceValues = ['auto', 'any', 'tool'];
    }
  }

  static lc_name(): string {
    return 'IllumaBedrockConverse';
  }

  /**
   * Get the model ID to use for API calls.
   * Returns applicationInferenceProfile if set, otherwise returns this.model.
   */
  protected getModelId(): string {
    return this.applicationInferenceProfile ?? this.model;
  }

  /**
   * Override invocationParams to:
   * 1. Add cachePoint to tools when promptCache is enabled
   * 2. Add serviceTier support
   *
   * CACHING STRATEGY: Separate cachePoints for core tools and MCP tools
   * - Core tools (web_search, execute_code, etc.) are stable → cache first
   * - MCP tools (have '_mcp_' in name) are dynamic → cache separately after
   * - This allows core tools to stay cached when MCP selection changes
   *
   * NOTE: Only Claude models support cachePoint - Nova and other models will reject it.
   */
  override invocationParams(
    options?: this['ParsedCallOptions'] & IllumaBedrockConverseCallOptions
  ): ReturnType<ChatBedrockConverse['invocationParams']> & {
    serviceTier?: { type: ServiceTierType };
  } {
    const params = super.invocationParams(options);

    // Add cachePoint to tools array if promptCache is enabled and tools exist
    /**
     * Bedrock requires all toolSpec.description fields to be non-empty strings.
     * Some tools (e.g., MCP-sourced or dynamically created) may have empty or
     * missing descriptions. Patch them here to avoid Bedrock validation errors.
     */
    if (params.toolConfig?.tools && Array.isArray(params.toolConfig.tools)) {
      for (const t of params.toolConfig.tools) {
        const spec = (t as { toolSpec?: { description?: string } }).toolSpec;
        if (spec && (!spec.description || spec.description === '')) {
          spec.description =
            spec.description ||
            `Tool: ${(spec as { name?: string }).name ?? 'unknown'}`;
        }
      }
    }

    // Only Claude models support cachePoint - check model name
    const modelId = this.model.toLowerCase();
    const isClaudeModel =
      modelId.includes('claude') || modelId.includes('anthropic');

    if (
      this.promptCache &&
      isClaudeModel &&
      params.toolConfig?.tools &&
      Array.isArray(params.toolConfig.tools) &&
      params.toolConfig.tools.length > 0
    ) {
      // Separate core tools from MCP tools
      // MCP tools have '_mcp_' in their name (e.g., 'search_emails_mcp_Google-Workspace')
      const coreTools: typeof params.toolConfig.tools = [];
      const mcpTools: typeof params.toolConfig.tools = [];

      for (const tool of params.toolConfig.tools) {
        // Check if tool has a name property with '_mcp_' pattern
        const toolName =
          (tool as { toolSpec?: { name?: string } }).toolSpec?.name ?? '';
        if (toolName.includes('_mcp_')) {
          mcpTools.push(tool);
        } else {
          coreTools.push(tool);
        }
      }

      // Build tools array with strategic cachePoints:
      // [CoreTool1, CoreTool2, cachePoint] + [MCPTool1, MCPTool2, cachePoint]
      const toolsWithCache: typeof params.toolConfig.tools = [];

      // Add core tools with cachePoint (if any)
      if (coreTools.length > 0) {
        toolsWithCache.push(...coreTools);
        toolsWithCache.push({ cachePoint: { type: 'default' } });
      }

      // Add MCP tools with their own cachePoint (if any)
      if (mcpTools.length > 0) {
        toolsWithCache.push(...mcpTools);
        toolsWithCache.push({ cachePoint: { type: 'default' } });
      }

      // If no tools at all (shouldn't happen but safety check)
      if (toolsWithCache.length === 0) {
        toolsWithCache.push({ cachePoint: { type: 'default' } });
      }

      params.toolConfig.tools = toolsWithCache;
    }

    // Add serviceTier support
    const serviceTierType = options?.serviceTier ?? this.serviceTier;

    return {
      ...params,
      serviceTier: serviceTierType ? { type: serviceTierType } : undefined,
    };
  }

  /**
   * Override _generateNonStreaming to use applicationInferenceProfile as modelId.
   * Uses the same model-swapping pattern as streaming for consistency.
   */
  override async _generateNonStreaming(
    messages: BaseMessage[],
    options: this['ParsedCallOptions'] & IllumaBedrockConverseCallOptions,
    runManager?: CallbackManagerForLLMRun
  ): Promise<ChatResult> {
    const originalModel = this.model;
    if (
      this.applicationInferenceProfile != null &&
      this.applicationInferenceProfile !== ''
    ) {
      this.model = this.applicationInferenceProfile;
    }

    try {
      return await super._generateNonStreaming(messages, options, runManager);
    } finally {
      this.model = originalModel;
    }
  }

  /**
   * Own the stream end-to-end so we have direct access to every
   * `contentBlockDelta.contentBlockIndex` from the AWS SDK.
   *
   * This replaces the parent's implementation which strips contentBlockIndex
   * from text and reasoning deltas, making it impossible to merge correctly.
   */
  override async *_streamResponseChunks(
    messages: BaseMessage[],
    options: this['ParsedCallOptions'] & IllumaBedrockConverseCallOptions,
    runManager?: CallbackManagerForLLMRun
  ): AsyncGenerator<ChatGenerationChunk> {
    const { converseMessages, converseSystem } =
      convertToConverseMessages(messages);
    const params = this.invocationParams(options);

    let { streamUsage } = this;
    if ((options as Record<string, unknown>).streamUsage !== undefined) {
      streamUsage = (options as Record<string, unknown>).streamUsage as boolean;
    }

    const modelId = this.getModelId();

    const command = new ConverseStreamCommand({
      modelId,
      messages: converseMessages,
      system: converseSystem,
      ...(params as Record<string, unknown>),
    });

    const response = await this.client.send(command, {
      abortSignal: options.signal,
    });

    if (!response.stream) {
      return;
    }

    const seenBlockIndices = new Set<number>();

    for await (const event of response.stream) {
      if (event.contentBlockStart != null) {
        const startChunk = handleConverseStreamContentBlockStart(
          event.contentBlockStart
        );
        if (startChunk != null) {
          const idx = event.contentBlockStart.contentBlockIndex;
          if (idx != null) {
            seenBlockIndices.add(idx);
          }
          yield this.enrichChunk(startChunk, seenBlockIndices);
        }
      } else if (event.contentBlockDelta != null) {
        const deltaChunk = handleConverseStreamContentBlockDelta(
          event.contentBlockDelta
        );

        const idx = event.contentBlockDelta.contentBlockIndex;
        if (idx != null) {
          seenBlockIndices.add(idx);
        }

        yield this.enrichChunk(deltaChunk, seenBlockIndices);

        await runManager?.handleLLMNewToken(
          deltaChunk.text,
          undefined,
          undefined,
          undefined,
          undefined,
          { chunk: deltaChunk }
        );
      } else if (event.metadata != null) {
        yield handleConverseStreamMetadata(event.metadata, { streamUsage });
      } else if (event.contentBlockStop != null) {
        const stopIdx = event.contentBlockStop.contentBlockIndex;
        if (stopIdx != null) {
          seenBlockIndices.add(stopIdx);
        }
      } else {
        yield new ChatGenerationChunk({
          text: '',
          message: new AIMessageChunk({
            content: '',
            response_metadata: event,
          }),
        });
      }
    }
  }

  /**
   * Inject `index` on content blocks for proper merge behaviour, then strip
   * `contentBlockIndex` from response_metadata to prevent `_mergeDicts` conflicts.
   *
   * Text string content is promoted to array form only when the stream contains
   * multiple content block indices (e.g. reasoning at index 0, text at index 1),
   * ensuring text merges correctly with the already-array accumulated content.
   */
  enrichChunk(
    chunk: ChatGenerationChunk,
    seenBlockIndices: Set<number>
  ): ChatGenerationChunk {
    const message = chunk.message;
    if (!(message instanceof AIMessageChunk)) {
      return chunk;
    }

    const metadata = message.response_metadata as Record<string, unknown>;
    const blockIndex = this.extractContentBlockIndex(metadata);
    const hasMetadataIndex = blockIndex != null;

    let content: AIMessageChunk['content'] = message.content;
    let contentModified = false;

    if (Array.isArray(content) && blockIndex != null) {
      content = content.map((block) =>
        typeof block === 'object' && !('index' in block)
          ? { ...block, index: blockIndex }
          : block
      );
      contentModified = true;
    } else if (
      typeof content === 'string' &&
      content !== '' &&
      blockIndex != null &&
      seenBlockIndices.size > 1
    ) {
      content = [{ type: 'text', text: content, index: blockIndex }];
      contentModified = true;
    }

    if (!contentModified && !hasMetadataIndex) {
      return chunk;
    }

    const cleanedMetadata = hasMetadataIndex
      ? (this.removeContentBlockIndex(metadata) as Record<string, unknown>)
      : metadata;

    return new ChatGenerationChunk({
      text: chunk.text,
      message: new AIMessageChunk({
        ...message,
        content,
        response_metadata: cleanedMetadata,
      }),
      generationInfo: chunk.generationInfo,
    });
  }

  /**
   * Extract `contentBlockIndex` from the top level of response_metadata.
   * Our custom handlers always place it at the top level.
   */
  private extractContentBlockIndex(
    metadata: Record<string, unknown>
  ): number | undefined {
    if (
      'contentBlockIndex' in metadata &&
      typeof metadata.contentBlockIndex === 'number'
    ) {
      return metadata.contentBlockIndex;
    }
    return undefined;
  }

  private removeContentBlockIndex(obj: unknown): unknown {
    if (obj === null || obj === undefined) {
      return obj;
    }

    if (Array.isArray(obj)) {
      return obj.map((item) => this.removeContentBlockIndex(item));
    }

    if (typeof obj === 'object') {
      const cleaned: Record<string, unknown> = {};
      for (const [key, value] of Object.entries(obj)) {
        if (key !== 'contentBlockIndex') {
          cleaned[key] = this.removeContentBlockIndex(value);
        }
      }
      return cleaned;
    }

    return obj;
  }
}

/** @deprecated Use IllumaBedrockConverse. Alias kept for backward compatibility with existing imports. */
export const CustomChatBedrockConverse = IllumaBedrockConverse;

export type { ChatBedrockConverseInput };