// src/proxy/correlation.ts
import type { ToolUseId, ToolUseRef } from "../cache/types.ts";
import type { Provider, RouteMatch } from "../provider/types.ts";

export interface Resolved {
  toolName: string;
  content: string;
  input: unknown;
}

export interface Correlation {
  observeOutbound(body: unknown, route: RouteMatch): void;
  pending(id: ToolUseId): ToolUseRef | undefined;
  resolved(id: ToolUseId): Resolved | undefined;
}

export interface CorrelationOpts {
  provider: Provider;
}

export function makeCorrelation(opts: CorrelationOpts): Correlation {
  const { provider } = opts;
  const pendingMap = new Map<string, ToolUseRef>();
  const resolvedMap = new Map<string, Resolved>();

  return {
    observeOutbound(body, route): void {
      const parsed = provider.parseRequest(body, route);
      for (const tc of provider.extractToolCalls(parsed)) {
        pendingMap.set(tc.id, { id: tc.id as ToolUseId, name: tc.name, input: tc.input });
      }
      for (const tr of provider.extractToolResults(parsed)) {
        const ref = pendingMap.get(tr.toolCallId);
        if (ref) {
          resolvedMap.set(tr.toolCallId, {
            toolName: ref.name, content: tr.content, input: ref.input,
          });
        }
      }
    },
    pending(id): ToolUseRef | undefined { return pendingMap.get(id as string); },
    resolved(id): Resolved | undefined { return resolvedMap.get(id as string); },
  };
}