import { discoverAuthorizationServerMetadata, discoverOAuthProtectedResourceMetadata, extractWWWAuthenticateParams } from "@modelcontextprotocol/client";
import { secureMcpUrl, validateDevHttpIssuer, type McpUrlPolicy } from "./mcp-url-policy.js";

export type McpHttpFetch = (input: Parameters<typeof fetch>[0], init?: Parameters<typeof fetch>[1]) => ReturnType<typeof fetch>;
type Settings = McpUrlPolicy & { resource: string; issuer: string };
type Endpoint = { tokenEndpoint: string; authMethod: "client_secret_basic" | "client_secret_post" };
const unavailable = (message: string) => Object.assign(new Error(message), { status: 503 });

function oauthChallenge(response: Response) {
  const header = response.headers.get("WWW-Authenticate") ?? "";
  const challenges: string[] = [];
  let start = 0, quoted = false, escaped = false;
  for (let i = 0; i < header.length; i++) {
    const char = header[i];
    if (escaped) { escaped = false; continue; }
    if (quoted && char === "\\") { escaped = true; continue; }
    if (char === '"') quoted = !quoted;
    if (!quoted && char === "," && /^\s*[^\s=,]+\s+(?!=)/.test(header.slice(i + 1))) {
      challenges.push(header.slice(start, i).trim()); start = i + 1;
    }
  }
  challenges.push(header.slice(start).trim());
  for (const value of challenges) {
    if (!/^(Bearer|DPoP)\s/i.test(value)) continue;
    const parsed = extractWWWAuthenticateParams(new Response(null, { headers: { "WWW-Authenticate": value } }));
    if (parsed.resourceMetadataUrl) return parsed;
  }
  return {};
}

/** Public metadata only. Configured issuer owns the pre-registered credentials. */
export class McpAuthDiscovery {
  private cached?: { endpoint: Endpoint; expires: number };
  private pending?: Promise<Endpoint>;
  private epoch = 0;

  constructor(private readonly settings: Settings, private readonly http: McpHttpFetch = fetch) {}

  invalidate() { this.epoch++; this.cached = undefined; }

  resolve(): Promise<Endpoint> {
    if (this.cached && this.cached.expires > Date.now()) return Promise.resolve(this.cached.endpoint);
    if (!this.pending) {
      const epoch = this.epoch;
      this.pending = this.discover().then((endpoint) => {
        if (epoch === this.epoch) this.cached = { endpoint, expires: Date.now() + 300_000 };
        return endpoint;
      }).finally(() => { this.pending = undefined; });
    }
    return this.pending;
  }

  private async discover(): Promise<Endpoint> {
    validateDevHttpIssuer(this.settings);
    if (this.settings.devHttpIssuer && this.settings.devHttpIssuer !== this.settings.issuer) throw unavailable("MCP development issuer must match the configured issuer");
    secureMcpUrl(this.settings.resource, this.settings);
    secureMcpUrl(this.settings.issuer, this.settings, true);
    const deadline = AbortSignal.timeout(10_000);
    const publicFetch: McpHttpFetch = async (input, init) => {
      const url = secureMcpUrl(input instanceof Request ? input.url : String(input), this.settings, true);
      const headers = new Headers(init?.headers);
      headers.delete("Authorization"); headers.delete("Cookie");
      const response = await this.http(url, { ...init, headers, credentials: "omit", redirect: "error", signal: deadline });
      // SDK parsers receive bounded public documents; no token or client secret is attached.
      const reader = response.body?.getReader();
      const chunks: Uint8Array[] = [];
      let size = 0;
      if (reader) {
        try {
          for (;;) {
            deadline.throwIfAborted();
            const part = await reader.read();
            if (part.done) break;
            size += part.value.length;
            if (size > 64 * 1024) throw unavailable("MCP OAuth metadata exceeds its size budget");
            chunks.push(part.value);
          }
        } finally { await reader.cancel().catch(() => {}); }
      }
      return new Response(size ? Buffer.concat(chunks) : null, { status: response.status, headers: response.headers });
    };
    try {
      const probe = await publicFetch(this.settings.resource, {
        method: "POST", headers: { "Content-Type": "application/json", Accept: "application/json, text/event-stream", "MCP-Protocol-Version": "2026-07-28", "Mcp-Method": "server/discover" },
        body: JSON.stringify({ jsonrpc: "2.0", id: "auth-discovery", method: "server/discover", params: { _meta: {
          "io.modelcontextprotocol/protocolVersion": "2026-07-28", "io.modelcontextprotocol/clientInfo": { name: "tjwater-agent-host", version: "1.0.0" }, "io.modelcontextprotocol/clientCapabilities": {},
        } } }),
      });
      const challenge = probe.status === 401 ? oauthChallenge(probe) : {};
      await probe.body?.cancel();
      const resource = await discoverOAuthProtectedResourceMetadata(this.settings.resource, {
        protocolVersion: "2026-07-28", resourceMetadataUrl: challenge.resourceMetadataUrl,
      }, publicFetch);
      if (new URL(resource.resource).href !== new URL(this.settings.resource).href || !resource.authorization_servers?.includes(this.settings.issuer)) {
        throw unavailable("MCP resource metadata differs from the configured resource or issuer");
      }
      const metadata = await discoverAuthorizationServerMetadata(this.settings.issuer, { fetchFn: publicFetch, protocolVersion: "2026-07-28" });
      // Enforce exact issuer equality even where an SDK permits trailing-slash normalization.
      if (!metadata || metadata.issuer !== this.settings.issuer || !metadata.token_endpoint) {
        throw unavailable("MCP OAuth issuer metadata is incompatible with its registration");
      }
      const tokenEndpoint = secureMcpUrl(metadata.token_endpoint, this.settings, true);
      if (new URL(this.settings.issuer).protocol === "https:" && tokenEndpoint.protocol !== "https:") {
        throw unavailable("MCP OAuth token endpoint cannot downgrade HTTPS");
      }
      const methods = metadata.token_endpoint_auth_methods_supported ?? ["client_secret_basic"];
      const authMethod = methods.includes("client_secret_post") ? "client_secret_post" : methods.includes("client_secret_basic") ? "client_secret_basic" : undefined;
      if (!authMethod) throw unavailable("MCP OAuth server does not support the configured client secret");
      return { tokenEndpoint: tokenEndpoint.href, authMethod };
    } catch {
      throw unavailable("MCP OAuth discovery failed or is incompatible with the configured resource and issuer");
    }
  }
}