import { WebSearchError } from "@umt/adapters";
import { CapabilityError, type CapabilityErrorCode, type CapabilityRegistry } from "@umt/agent";
import { capabilitySpec } from "@umt/contracts";
import { allowScopeKey, type WebSearchService, WebSearchServiceError } from "./service.ts";

// web.search / web.fetch(read,low):错误按类别映射成能力错误;输出本身带 trust = untrusted 与不可信提示,
// 进入模型上下文前 PiAgentRuntime 还会整体过一遍 redactForModel。

const CODE: Record<string, CapabilityErrorCode> = {
  disabled: "precondition",
  not_allowed: "precondition",
  blocked: "precondition",
  unsupported: "precondition",
  bad_request: "precondition",
  auth: "precondition",
  quota: "budget",
  rate: "budget",
  upstream: "precondition",
  network: "precondition",
  timeout: "precondition",
  protocol: "precondition",
};

async function guard<T>(fn: () => Promise<T>): Promise<T> {
  try {
    return await fn();
  } catch (e) {
    if (e instanceof WebSearchError || e instanceof WebSearchServiceError)
      throw new CapabilityError(CODE[e.kind] ?? "precondition", e.message, { kind: e.kind });
    throw e;
  }
}

export function registerWebSearchCapabilities(registry: CapabilityRegistry, svc: () => WebSearchService): void {
  registry.register(capabilitySpec("web.search"), (input, c) =>
    guard(async () => {
      const { durationMs: _d, ...out } = await svc().search(input, { scopeKey: allowScopeKey(c), signal: c.signal });
      return out;
    }),
  );
  registry.register(capabilitySpec("web.fetch"), (input, c) =>
    guard(() => svc().fetch(input.url, { scopeKey: allowScopeKey(c), signal: c.signal })),
  );
}