/** Native AWS Bedrock binding for AgentGuard Spend. */

import type { CapabilityTier, AttestationChainInput, CallContext, Provider, SpendPolicy, SpendScope } from '../types';
import {
  AgentGuardBlockedError,
  SpendGuard,
  type SpendGuardConfig,
  extractText,
  usageCountsFromObject,
  wrapUsageStream,
  type UsageCounts,
} from '../spend-guard';

export interface BedrockBindingOptions {
  policy: SpendPolicy;
  scope: SpendScope;
  capabilityClaim?: CapabilityTier;
  /** Optional DAG-TAT chain backing the capability claim (required when policy.attestation lists it). */
  attestation?: AttestationChainInput;
  config?: Omit<SpendGuardConfig, 'policy'>;
}

type BedrockCommand = {
  input?: Record<string, unknown>;
  constructor?: { name?: string };
};

export function withSpendGuardBedrock(
  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  client: any,
  opts: BedrockBindingOptions,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
): any {
  const guard = new SpendGuard({ policy: opts.policy, ...(opts.config ?? {}) });
  const originalSend = client?.send?.bind(client);
  if (!originalSend) throw new Error('withSpendGuardBedrock: expected client.send');

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  client.send = async (command: BedrockCommand, ...args: any[]) => {
    if (!isInvokeModelCommand(command)) return originalSend(command, ...args);

    const call = buildBedrockCall(command, guard, opts);
    const { decision } = await guard.decide(call);
    if (decision.action === 'block') throw new AgentGuardBlockedError(decision, opts.scope, opts.config?.locale);

    if (decision.action === 'downgrade' && decision.modelResolved !== decision.modelRequested && command.input) {
      command.input = { ...command.input, modelId: decision.modelResolved };
    }

    const response = await originalSend(command, ...args);
    if (isStreamingCommand(command) && response?.body) {
      return {
        ...response,
        body: wrapBedrockStream(response.body, guard, decision.decisionId, call),
      };
    }

    const usage = extractBedrockResponseUsage(response, call.model);
    if (usage) await guard.settleStreamUsage(decision.decisionId, usage.inputTokens, usage.outputTokens);
    return response;
  };

  client.__agentguard = guard;
  return client;
}

function isInvokeModelCommand(command: BedrockCommand): boolean {
  const name = command?.constructor?.name;
  return name === 'InvokeModelCommand' || name === 'InvokeModelWithResponseStreamCommand';
}

function isStreamingCommand(command: BedrockCommand): boolean {
  return command?.constructor?.name === 'InvokeModelWithResponseStreamCommand';
}

function buildBedrockCall(command: BedrockCommand, guard: SpendGuard, opts: BedrockBindingOptions): CallContext {
  const input = command.input ?? {};
  const modelId = String(input.modelId ?? 'unknown');
  const body = parseJsonBody(input.body);
  const inputText = extractText(body);
  const metadata = (body as Record<string, unknown> | null)?.metadata as Record<string, unknown> | undefined;
  return {
    provider: 'bedrock' as Provider,
    model: modelId,
    inputTokens: guard.estimateTokens(inputText),
    outputTokens: outputTokensFromBody(body),
    scope: opts.scope,
    capabilityClaim: opts.capabilityClaim,
    attestation: opts.attestation,
    label: stringValue(metadata?.label),
    workflowId: stringValue(metadata?.workflow_id) || stringValue(metadata?.workflowId),
    subagentId: stringValue(metadata?.subagent_id) || stringValue(metadata?.subagentId),
    callFingerprint: stringValue(metadata?.call_fingerprint) || stringValue(metadata?.callFingerprint),
    requestShape: { model: modelId, maxTokens: outputTokensFromBody(body), stream: isStreamingCommand(command), bodyKeys: Object.keys((body as Record<string, unknown> | null) ?? {}).sort() },
  };
}

function stringValue(value: unknown): string | undefined {
  return typeof value === 'string' && value.trim() ? value.trim() : undefined;
}

function wrapBedrockStream(
  stream: AsyncIterable<unknown>,
  guard: SpendGuard,
  decisionId: string,
  call: CallContext,
): AsyncIterable<unknown> {
  const tracker = new BedrockStreamUsageTracker(guard, call);
  return wrapUsageStream(stream, {
    onChunk: (event) => tracker.observe(event),
    settle: (partial, reason) => {
      const usage = tracker.usage();
      return guard.settleStreamUsage(decisionId, usage.inputTokens, usage.outputTokens, { partial, reason });
    },
  });
}

class BedrockStreamUsageTracker {
  private input: number | null = null;
  private output: number | null = null;
  private content = '';

  constructor(private guard: SpendGuard, private call: CallContext) {}

  observe(event: unknown): void {
    const payload = bedrockEventPayload(event);
    if (!payload) return;
    const usage = usageFromBedrockPayload(payload);
    if (usage) {
      this.input = usage.inputTokens;
      this.output = usage.outputTokens;
    }
    this.content += textFromBedrockPayload(payload);
  }

  usage(): UsageCounts {
    return {
      inputTokens: this.input ?? this.call.inputTokens,
      outputTokens: this.output ?? this.guard.estimateTokens(this.content),
    };
  }
}

function extractBedrockResponseUsage(response: unknown, modelId: string): UsageCounts | null {
  const body = (response as { body?: unknown } | null)?.body;
  const parsed = parseJsonBody(body);
  const usage = usageFromBedrockPayload(parsed);
  if (usage) return usage;
  if (modelId.startsWith('amazon.nova.')) return usageCountsFromObject((parsed as Record<string, unknown> | null)?.usage);
  return null;
}

function usageFromBedrockPayload(payload: unknown): UsageCounts | null {
  if (!payload || typeof payload !== 'object') return null;
  const obj = payload as Record<string, unknown>;
  const direct = usageCountsFromObject(obj.usage);
  if (direct) return direct;

  const bedrockMetrics = obj['amazon-bedrock-invocationMetrics'];
  if (bedrockMetrics && typeof bedrockMetrics === 'object') {
    const metrics = bedrockMetrics as Record<string, unknown>;
    const input = metrics.inputTokenCount;
    const output = metrics.outputTokenCount;
    if (Number.isSafeInteger(input) && Number.isSafeInteger(output)) {
      return { inputTokens: input as number, outputTokens: output as number };
    }
  }

  const nested = usageCountsFromObject((obj.message as Record<string, unknown> | undefined)?.usage);
  if (nested) return nested;
  return null;
}

function textFromBedrockPayload(payload: unknown): string {
  if (!payload || typeof payload !== 'object') return '';
  const obj = payload as Record<string, unknown>;
  const delta = obj.delta as Record<string, unknown> | undefined;
  if (typeof delta?.text === 'string') return delta.text;
  const contentBlockDelta = obj.contentBlockDelta as Record<string, unknown> | undefined;
  const nestedDelta = contentBlockDelta?.delta as Record<string, unknown> | undefined;
  if (typeof nestedDelta?.text === 'string') return nestedDelta.text;
  const output = obj.output as Record<string, unknown> | undefined;
  if (output) return extractText(output);
  return '';
}

function outputTokensFromBody(body: unknown): number {
  if (!body || typeof body !== 'object') return 512;
  const obj = body as Record<string, unknown>;
  const value = obj.max_tokens ?? obj.maxTokens ?? obj.max_new_tokens ??
    (obj.inferenceConfig as Record<string, unknown> | undefined)?.maxTokens;
  return Number.isSafeInteger(value) && (value as number) >= 0 ? value as number : 512;
}

function parseJsonBody(body: unknown): unknown {
  if (!body) return {};
  try {
    if (typeof body === 'string') return JSON.parse(body);
    if (body instanceof Uint8Array) return JSON.parse(new TextDecoder().decode(body));
    if (typeof Buffer !== 'undefined' && Buffer.isBuffer(body)) return JSON.parse(body.toString('utf8'));
    if (typeof body === 'object') return body;
  } catch {
    return {};
  }
  return {};
}

function bedrockEventPayload(event: unknown): unknown {
  const obj = event as Record<string, unknown> | null;
  const chunk = obj?.chunk as Record<string, unknown> | undefined;
  const bytes = chunk?.bytes ?? chunk?.payload ?? obj?.payload;
  return parseJsonBody(bytes);
}