import { z } from 'zod';

import type { AgentCredentials } from './agent-credentials.js';
import { agentRequest, agentResponse } from './agent-http.js';

const resourceSchema = z.object({ authorization_servers: z.array(z.url()).min(1) });
const discoverySchema = z.object({
  issuer: z.url(),
  token_endpoint: z.url(),
  agent_auth: z.object({
    identity_endpoint: z.url(),
    claim_endpoint: z.url(),
    identity_types_supported: z.array(z.string()),
  }),
});
export const agentIdentityAssertionSchema = z.object({ assertion: z.string().min(1) });
const registrationSchema = z.object({
  id: z.string().min(1),
  identity: agentIdentityAssertionSchema,
  claim: z.object({ token: z.string().min(1) }),
});
const tokenSchema = z.object({
  access_token: z.string().min(1),
  expires_in: z.number().positive(),
  token_type: z.string().regex(/^bearer$/i),
});
const identitySchema = z.discriminatedUnion('identityType', [
  z.object({ registrationId: z.string(), identityType: z.literal('anonymous') }),
  z.object({
    registrationId: z.string(),
    identityType: z.literal('user_bound'),
    userId: z.string(),
    orgId: z.string(),
  }),
]);

export type AgentDiscovery = z.infer<typeof discoverySchema>;

export async function discoverAgentAuth(apiUrl: string): Promise<AgentDiscovery> {
  const resource = await agentResponse(
    await agentRequest(`${apiUrl}/api/agent/.well-known/oauth-protected-resource`),
    resourceSchema
  );
  const issuer = resource.authorization_servers[0];
  if (!issuer || new URL(issuer).protocol !== 'https:')
    throw new Error('Agent authentication requires an HTTPS issuer.');
  const discovery = await agentResponse(
    await agentRequest(`${issuer.replace(/\/$/, '')}/.well-known/oauth-authorization-server`),
    discoverySchema
  );
  const origin = new URL(issuer).origin;
  if (
    [
      discovery.issuer,
      discovery.token_endpoint,
      discovery.agent_auth.identity_endpoint,
      discovery.agent_auth.claim_endpoint,
    ].some((url) => new URL(url).origin !== origin)
  ) {
    throw new Error('Agent authentication endpoints do not match the configured issuer.');
  }
  return discovery;
}

export async function registerAgent(discovery: AgentDiscovery): Promise<AgentCredentials> {
  if (!discovery.agent_auth.identity_types_supported.includes('anonymous')) {
    throw new Error('Anonymous registration is not enabled for this environment.');
  }
  const result = await agentResponse(
    await agentRequest(discovery.agent_auth.identity_endpoint, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ type: 'anonymous' }),
    }),
    registrationSchema
  );
  return {
    registrationId: result.id,
    issuer: discovery.issuer,
    assertion: result.identity.assertion,
    claimToken: result.claim.token,
  };
}

export async function exchangeAgentToken(
  discovery: AgentDiscovery,
  credentials: AgentCredentials
): Promise<AgentCredentials & { accessToken: string }> {
  if (new URL(credentials.issuer).origin !== new URL(discovery.issuer).origin) {
    throw new Error('Saved agent credentials belong to another issuer.');
  }
  if (credentials.accessToken && (credentials.accessTokenExpiresAt ?? 0) > Date.now() + 30_000) {
    return { ...credentials, accessToken: credentials.accessToken };
  }
  const token = await agentResponse(
    await agentRequest(discovery.token_endpoint, {
      method: 'POST',
      body: new URLSearchParams({
        grant_type: 'urn:ietf:params:oauth:grant-type:jwt-bearer',
        assertion: credentials.assertion,
      }),
    }),
    tokenSchema
  );
  return {
    ...credentials,
    accessToken: token.access_token,
    accessTokenExpiresAt: Date.now() + token.expires_in * 1000,
  };
}

export async function getAgentIdentity(apiUrl: string, accessToken: string) {
  return agentResponse(
    await agentRequest(`${apiUrl}/api/agent/identity`, {
      headers: { Authorization: `Bearer ${accessToken}` },
    }),
    identitySchema
  );
}