import type { SpendGuard } from '../spend-guard';
import { createFrameworkGuard, modelFromUnknown, objectRecord, preflightFrameworkCall, recordFrameworkReceipt, redactDataPlane, settleFrameworkCall, type FrameworkAdapterOptions, type FrameworkPreflight } from './common';

export type AgentGuardAiSdkMiddleware = {
  specificationVersion: 'v3';
  transformParams?: (args: { type: 'generate' | 'stream'; params: unknown; model: unknown }) => Promise<unknown> | unknown;
  wrapGenerate?: (args: { doGenerate: () => Promise<unknown>; doStream?: () => Promise<unknown>; params: unknown; model: unknown }) => Promise<unknown>;
  wrapStream?: (args: { doGenerate?: () => Promise<unknown>; doStream: () => Promise<unknown>; params: unknown; model: unknown }) => Promise<unknown>;
};

export interface AgentGuardAiSdkOptions extends FrameworkAdapterOptions {
  framework?: 'vercel-ai-sdk';
}

export function agentguardAiSdkMiddleware(opts: AgentGuardAiSdkOptions): AgentGuardAiSdkMiddleware {
  const guard = createFrameworkGuard({ ...opts, framework: 'vercel-ai-sdk' });
  return {
    specificationVersion: 'v3',
    transformParams: ({ params }) => params,
    wrapGenerate: async ({ doGenerate, params, model }) => {
      const preflight = await aiSdkPreflight(guard, opts, params, model, 'generate');
      const result = await doGenerate();
      await settleFrameworkCall(preflight, result);
      await recordFrameworkReceipt(guard, 'vercel-ai-sdk', {
        event: 'post_generate',
        decisionId: preflight.decision.decisionId,
        model: preflight.decision.modelResolved,
        requestShape: requestShape(params, 'generate'),
        hasUsage: Boolean(objectRecord(result)?.usage),
      });
      return result;
    },
    wrapStream: async ({ doStream, params, model }) => {
      const preflight = await aiSdkPreflight(guard, opts, params, model, 'stream');
      const result = await doStream();
      return wrapAiSdkStreamResult(guard, preflight, result, params);
    },
  };
}

export const agentguard = agentguardAiSdkMiddleware;

async function aiSdkPreflight(
  guard: SpendGuard,
  opts: AgentGuardAiSdkOptions,
  params: unknown,
  model: unknown,
  type: 'generate' | 'stream',
): Promise<FrameworkPreflight> {
  return preflightFrameworkCall(guard, opts, {
    framework: 'vercel-ai-sdk',
    model: modelFromUnknown(model, opts.defaultModel ?? 'unknown'),
    params,
    metadata: agentguardMetadata(params),
    requestShape: requestShape(params, type),
  });
}

function requestShape(params: unknown, type: 'generate' | 'stream'): Record<string, unknown> {
  const record = objectRecord(params) ?? {};
  const tools = objectRecord(record.tools) ?? {};
  return {
    type,
    maxOutputTokens: record.maxOutputTokens ?? record.maxTokens,
    toolNames: Object.keys(tools).sort(),
    providerOptionsPresent: Boolean(record.providerOptions),
    redacted: redactDataPlane({ prompt: record.prompt, messages: record.messages }),
  };
}

function agentguardMetadata(params: unknown): Record<string, unknown> {
  const record = objectRecord(params) ?? {};
  const providerOptions = objectRecord(record.providerOptions) ?? objectRecord(record.providerMetadata) ?? {};
  const metadata = objectRecord(providerOptions.agentguard) ?? {};
  return metadata;
}

async function wrapAiSdkStreamResult(
  guard: SpendGuard,
  preflight: FrameworkPreflight,
  result: unknown,
  params: unknown,
): Promise<unknown> {
  const record = objectRecord(result);
  const stream = record?.stream;
  if (!stream || typeof (stream as { pipeThrough?: unknown }).pipeThrough !== 'function') {
    await settleFrameworkCall(preflight, result, true);
    await recordFrameworkReceipt(guard, 'vercel-ai-sdk', {
      event: 'post_stream',
      decisionId: preflight.decision.decisionId,
      model: preflight.decision.modelResolved,
      requestShape: requestShape(params, 'stream'),
      partial: true,
    });
    return result;
  }
  const TransformStreamCtor = (globalThis as unknown as { TransformStream?: new (handlers: Record<string, unknown>) => unknown }).TransformStream;
  if (!TransformStreamCtor) return result;
  let finalUsage: unknown = null;
  const transform = new TransformStreamCtor({
    transform(chunk: unknown, controller: { enqueue: (chunk: unknown) => void }) {
      const usage = objectRecord(chunk)?.usage;
      if (usage) finalUsage = { usage };
      controller.enqueue(chunk);
    },
    async flush() {
      await settleFrameworkCall(preflight, finalUsage ?? result);
      await recordFrameworkReceipt(guard, 'vercel-ai-sdk', {
        event: 'post_stream',
        decisionId: preflight.decision.decisionId,
        model: preflight.decision.modelResolved,
        requestShape: requestShape(params, 'stream'),
        partial: false,
      });
    },
  });
  return { ...record, stream: (stream as { pipeThrough: (transform: unknown) => unknown }).pipeThrough(transform) };
}