import { z } from 'zod';

import {
  agentIdentityAssertionSchema,
  discoverAgentAuth,
  exchangeAgentToken,
} from './agent-auth.js';
import {
  loadAgentCredentials,
  saveAgentCredentials,
  type AgentCredentials,
} from './agent-credentials.js';
import { agentRequest, agentResponse } from './agent-http.js';
import { DASHBOARD_URL } from './constants.js';

const claimAttemptSchema = z.object({ attempt: z.object({ verification_uri: z.url() }) });

export async function createAgentClaimLink(
  apiUrl: string,
  credentials: AgentCredentials,
  email: string
): Promise<string> {
  const loginHint = z.email().parse(email);
  const discovery = await discoverAgentAuth(apiUrl);
  if (new URL(discovery.issuer).origin !== new URL(credentials.issuer).origin) {
    throw new Error('Saved agent credentials belong to another issuer.');
  }
  const claim = await agentResponse(
    await agentRequest(discovery.agent_auth.claim_endpoint, {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        type: 'service_auth',
        claim_token: credentials.claimToken,
        login_hint: loginHint,
      }),
    }),
    claimAttemptSchema
  );
  const token = z
    .string()
    .min(1)
    .parse(new URL(claim.attempt.verification_uri).searchParams.get('token'));
  const claimUrl = new URL('/agent-claim', DASHBOARD_URL);
  claimUrl.searchParams.set('token', token);
  return claimUrl.toString();
}

const completedClaimSchema = z.object({
  identity: agentIdentityAssertionSchema,
});
const signupSchema = z.object({ orgId: z.string().min(1), claimed: z.literal(true) });

export async function claimDeployment(
  options: { apiUrl: string; directory: string; code?: string; email?: string },
  warn: (message: string) => void,
  progress: (message: string) => void = () => {}
): Promise<string> {
  const { apiUrl, directory } = options;
  const code = options.code?.trim();
  const email = options.email?.trim();
  if ((!code && !email) || (code && email)) {
    throw new Error(
      'Provide either --email <email> for a claim link or --code <code> to finish the claim.'
    );
  }
  let credentials = await loadAgentCredentials(apiUrl, directory);
  if (!credentials)
    throw new Error(
      'No saved agent registration. Run from the directory used for mint deploy, or pass --directory.'
    );
  if (email && !credentials.claimVerified) {
    progress('Creating the claim link...');
    const link = await createAgentClaimLink(apiUrl, credentials, email);
    return `Claim link for the user: ${link}\nAsk the user to open it, sign in, and return the displayed code. Then run mint claim --code <code> from the same directory (or repeat --directory).`;
  }
  const discovery = await discoverAgentAuth(apiUrl);
  if (new URL(discovery.issuer).origin !== new URL(credentials.issuer).origin) {
    throw new Error('Saved agent credentials belong to another issuer.');
  }
  const save = async (value: AgentCredentials) => {
    try {
      await saveAgentCredentials(apiUrl, directory, value);
    } catch {
      warn(
        'Could not save claimed credentials to the keyring. Continuing for this run; a retry may not recover the completed claim.'
      );
    }
  };
  if (!credentials.claimVerified) {
    progress('Completing the WorkOS claim...');
    const completed = await agentResponse(
      await agentRequest(`${discovery.agent_auth.claim_endpoint.replace(/\/$/, '')}/complete`, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify({ claim_token: credentials.claimToken, user_code: code }),
      }),
      completedClaimSchema
    );
    credentials = {
      ...credentials,
      assertion: completed.identity.assertion,
      claimVerified: true,
      accessToken: undefined,
      accessTokenExpiresAt: undefined,
    };
    await save(credentials);
  }
  progress('Finishing the site claim...');
  const authenticated = await exchangeAgentToken(discovery, credentials);
  await agentResponse(
    await agentRequest(`${apiUrl}/api/agent/signup`, {
      method: 'POST',
      headers: { Authorization: `Bearer ${authenticated.accessToken}` },
    }),
    signupSchema
  );
  await save(authenticated);
  return 'Site claimed successfully. The user can now manage it in the Mintlify dashboard.';
}