import {
  BedrockRuntimeClient,
  ConverseCommand,
  ConverseStreamCommand,
} from '@aws-sdk/client-bedrock-runtime';
import type {
  Message as BedrockMessage,
  ContentBlock,
  ToolConfiguration,
  Tool as BedrockTool,
  ConverseCommandInput,
  ConverseStreamCommandInput,
} from '@aws-sdk/client-bedrock-runtime';
import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient, ModelToolCall, ToolDefinition, Message, MessageContentPart } from './types.js';

export interface BedrockClientConfig {
  model: string;
  region?: string;
  maxTokens?: number;
  /** AWS access key ID (if not using default credential chain). */
  accessKeyId?: string;
  /** AWS secret access key (if not using default credential chain). */
  secretAccessKey?: string;
}

export class BedrockClient implements ModelClient {
  private client: BedrockRuntimeClient;
  private model: string;
  private defaultMaxTokens: number;

  constructor(config: BedrockClientConfig) {
    const clientConfig: Record<string, unknown> = {
      region: config.region ?? process.env.AWS_REGION ?? 'us-east-1',
    };

    if (config.accessKeyId && config.secretAccessKey) {
      clientConfig.credentials = {
        accessKeyId: config.accessKeyId,
        secretAccessKey: config.secretAccessKey,
      };
    }

    this.client = new BedrockRuntimeClient(clientConfig);
    this.model = config.model;
    this.defaultMaxTokens = config.maxTokens ?? 4096;
  }

  async chat(request: ChatRequest): Promise<ChatResponse> {
    const messages = convertMessages(request.messages);

    const params: ConverseCommandInput = {
      modelId: this.model,
      messages,
      inferenceConfig: {
        maxTokens: request.maxTokens ?? this.defaultMaxTokens,
      },
    };

    if (request.system) {
      params.system = [{ text: request.system }];
    }

    if (request.tools && request.tools.length > 0) {
      params.toolConfig = convertTools(request.tools);
    }

    const command = new ConverseCommand(params);
    const response = await this.client.send(command);

    // Extract text and tool_use content from the response
    const outputContent = response.output?.message?.content ?? [];
    const textParts: string[] = [];
    const toolCalls: ModelToolCall[] = [];

    for (const block of outputContent) {
      if ('text' in block && block.text !== undefined) {
        textParts.push(block.text);
      }
      if ('toolUse' in block && block.toolUse !== undefined) {
        toolCalls.push({
          id: block.toolUse.toolUseId ?? `bedrock_${Date.now()}`,
          name: block.toolUse.name ?? '',
          args: block.toolUse.input as unknown,
        });
      }
    }

    const content = textParts.join('');

    // Map stop reason
    let stopReason: string = 'end_turn';
    if (response.stopReason === 'max_tokens') {stopReason = 'max_tokens';}
    else if (response.stopReason === 'tool_use') {stopReason = 'tool_use';}
    else if (response.stopReason === 'end_turn') {stopReason = 'end_turn';}
    else if (response.stopReason) {stopReason = response.stopReason;}

    return {
      content,
      stopReason,
      usage: {
        inputTokens: response.usage?.inputTokens ?? 0,
        outputTokens: response.usage?.outputTokens ?? 0,
      },
      ...(toolCalls.length > 0 ? { toolCalls } : {}),
    };
  }

  async *chatStream(request: ChatRequest): AsyncIterable<ChatStreamEvent> {
    const messages = convertMessages(request.messages);

    const params: ConverseStreamCommandInput = {
      modelId: this.model,
      messages,
      inferenceConfig: {
        maxTokens: request.maxTokens ?? this.defaultMaxTokens,
      },
    };

    if (request.system) {
      params.system = [{ text: request.system }];
    }

    if (request.tools && request.tools.length > 0) {
      params.toolConfig = convertTools(request.tools);
    }

    try {
      const command = new ConverseStreamCommand(params);
      const response = await this.client.send(command);

      let inputTokens = 0;
      let outputTokens = 0;

      if (response.stream) {
        for await (const event of response.stream) {
          if (event.contentBlockDelta?.delta && 'text' in event.contentBlockDelta.delta && event.contentBlockDelta.delta.text) {
            yield { type: 'content', content: event.contentBlockDelta.delta.text };
          }

          if (event.metadata?.usage) {
            inputTokens = event.metadata.usage.inputTokens ?? inputTokens;
            outputTokens = event.metadata.usage.outputTokens ?? outputTokens;
          }
        }
      }

      yield {
        type: 'done',
        usage: { inputTokens, outputTokens },
      };
    } catch (error) {
      yield {
        type: 'error',
        error: error instanceof Error ? error : new Error(String(error)),
      };
    }
  }
}

function convertMessages(messages: Message[]): BedrockMessage[] {
  return messages.map(m => {
    const role = m.role === 'assistant' ? 'assistant' as const : 'user' as const;

    if (typeof m.content === 'string') {
      return { role, content: [{ text: m.content }] as ContentBlock[] };
    }

    // Multimodal content: convert each part
    const blocks: ContentBlock[] = m.content.map((part: MessageContentPart) => {
      if (part.type === 'text') {
        return { text: part.text } as ContentBlock;
      }
      if (part.type === 'image') {
        // Image part — Bedrock uses { image: { format, source: { bytes } } }
        if (part.source.type === 'base64' && part.source.data) {
          return {
            image: {
              format: part.source.media_type.split('/')[1] as 'jpeg' | 'png' | 'gif' | 'webp',
              source: { bytes: Uint8Array.from(atob(part.source.data), c => c.charCodeAt(0)) },
            },
          } as unknown as ContentBlock;
        }
        // URL images not natively supported by Bedrock — fall back to text placeholder
        return { text: `[Image: ${part.source.url ?? 'unsupported'}]` } as ContentBlock;
      }
      // Audio — Bedrock doesn't support native audio input; use transcript fallback
      if (part.type === 'audio') {
        if (part.source.transcript) {
          return { text: `[Voice message]: ${part.source.transcript}` } as ContentBlock;
        }
        return { text: '[Audio message received but no transcript available]' } as ContentBlock;
      }
      return { text: JSON.stringify(part) } as ContentBlock;
    });

    return { role, content: blocks };
  });
}

function convertTools(tools: ToolDefinition[]): ToolConfiguration {
  return {
    tools: tools.map(t => ({
      toolSpec: {
        name: t.name,
        description: t.description,
        inputSchema: {
          json: t.input_schema as Record<string, unknown>,
        },
      },
    } as BedrockTool)),
  };
}
