import { randomUUID } from "node:crypto";
import path from "node:path";
import type { Request, Response } from "express";
import {
  AccessDeniedError,
  InvalidGrantError,
  InvalidScopeError,
  InvalidTargetError,
  InvalidTokenError,
} from "@modelcontextprotocol/sdk/server/auth/errors.js";
import type {
  AuthorizationParams,
  OAuthServerProvider,
} from "@modelcontextprotocol/sdk/server/auth/provider.js";
import type { AuthInfo } from "@modelcontextprotocol/sdk/server/auth/types.js";
import type {
  OAuthClientInformationFull,
  OAuthTokenRevocationRequest,
  OAuthTokens,
} from "@modelcontextprotocol/sdk/shared/auth.js";
import type { OAuthRegisteredClientsStore } from "@modelcontextprotocol/sdk/server/auth/clients.js";
import { EncryptJWT, SignJWT, jwtDecrypt, jwtVerify, type JWTPayload } from "jose";
import type { AppConfig } from "../config.js";
import { deriveKey } from "../security/keys.js";
import { verifyPassword } from "../security/password.js";
import { ReplayStore } from "../security/replay-store.js";
import { PersistentClientStore } from "./client-store.js";

const ALL_SCOPES = ["mail.read", "mail.draft", "mail.send"] as const;
const CODE_TTL_SECONDS = 90;
const AUTH_REQUEST_TTL_SECONDS = 600;

type OAuthClaims = JWTPayload & {
  client_id: string;
  redirect_uri?: string;
  code_challenge?: string;
  resource: string;
  scope: string;
  name?: string;
  email?: string;
};

function escapeHtml(value: string): string {
  return value
    .replaceAll("&", "&")
    .replaceAll("<", "&lt;")
    .replaceAll(">", "&gt;")
    .replaceAll('"', "&quot;")
    .replaceAll("'", "&#039;");
}

function asString(payload: JWTPayload, key: string): string {
  const value = payload[key];
  if (typeof value !== "string" || !value) throw new InvalidGrantError("Invalid token claims");
  return value;
}

export class MailOAuthProvider implements OAuthServerProvider {
  readonly clientsStore: OAuthRegisteredClientsStore;
  private readonly accessKey: Uint8Array;
  private readonly encryptedTokenKey: Uint8Array;
  private readonly approvalKey: Uint8Array;
  private readonly replayStore: ReplayStore;
  private readonly issuer: string;
  private readonly resource: string;
  private readonly enabledScopes: string[];

  constructor(private readonly config: AppConfig) {
    this.issuer = config.baseUrl.href;
    this.resource = config.mcpUrl.href;
    this.accessKey = deriveKey(config.tokenSecret, "oauth-access-token");
    this.encryptedTokenKey = deriveKey(config.tokenSecret, "oauth-encrypted-token");
    this.approvalKey = deriveKey(config.tokenSecret, "oauth-approval-request");
    this.replayStore = new ReplayStore(path.join(config.stateDir, "oauth-replay.json"));
    this.enabledScopes = [
      "mail.read",
      ...(config.features.drafts ? ["mail.draft"] : []),
      ...(config.features.send ? ["mail.send"] : []),
    ];

    const persistentClients = new PersistentClientStore(
      path.join(config.stateDir, "oauth-clients.json"),
      config.oauth.allowedRedirectHosts,
      config.oauth.allowDynamicRegistration,
      config.oauth.staticClient,
      config.nodeEnv !== "production",
    );
    this.clientsStore = config.oauth.allowDynamicRegistration
      ? persistentClients
      : { getClient: (clientId) => persistentClients.getClient(clientId) };
  }

  get scopesSupported(): string[] {
    return [...this.enabledScopes];
  }

  async authorize(
    client: OAuthClientInformationFull,
    params: AuthorizationParams,
    res: Response,
  ): Promise<void> {
    this.assertResource(params.resource);
    const scopes = this.normalizeScopes(params.scopes);
    const now = Math.floor(Date.now() / 1000);
    const transaction = await new EncryptJWT({
      client_id: client.client_id,
      redirect_uri: params.redirectUri,
      code_challenge: params.codeChallenge,
      resource: this.resource,
      scope: scopes.join(" "),
      ...(params.state ? { state: params.state } : {}),
    })
      .setProtectedHeader({ alg: "dir", enc: "A256GCM", typ: "oauth-approval+jwt" })
      .setIssuer(this.issuer)
      .setAudience("oauth-approval")
      .setIssuedAt(now)
      .setExpirationTime(now + AUTH_REQUEST_TTL_SECONDS)
      .setJti(randomUUID())
      .encrypt(this.approvalKey);

    res.status(200).type("html").send(this.renderLogin(transaction, scopes));
  }

  approvalHandler = async (req: Request, res: Response): Promise<void> => {
    const transaction = typeof req.body?.transaction === "string" ? req.body.transaction : "";
    if (!transaction || transaction.length > 16_384) {
      res.status(400).send("Invalid authorization request.");
      return;
    }

    let claims: OAuthClaims;
    try {
      const decrypted = await jwtDecrypt(transaction, this.approvalKey, {
        issuer: this.issuer,
        audience: "oauth-approval",
        keyManagementAlgorithms: ["dir"],
        contentEncryptionAlgorithms: ["A256GCM"],
      });
      claims = decrypted.payload as OAuthClaims;
      asString(claims, "client_id");
      asString(claims, "redirect_uri");
      asString(claims, "code_challenge");
      asString(claims, "resource");
      asString(claims, "scope");
    } catch {
      res.status(400).send("Authorization request expired or invalid.");
      return;
    }

    const redirectUri = asString(claims, "redirect_uri");
    const state = typeof claims.state === "string" ? claims.state : undefined;
    if (req.body?.decision === "deny") {
      this.redirectWithError(res, redirectUri, state, new AccessDeniedError("Access denied"));
      return;
    }

    const password = typeof req.body?.password === "string" ? req.body.password : "";
    if (password.length > 1_024 || !(await verifyPassword(password, this.config.adminPasswordHash))) {
      const scopes = asString(claims, "scope").split(" ").filter(Boolean);
      res.status(401).type("html").send(this.renderLogin(transaction, scopes, "Password incorrect."));
      return;
    }

    const clientId = asString(claims, "client_id");
    const client = await this.clientsStore.getClient(clientId);
    if (!client || !client.redirect_uris.includes(redirectUri)) {
      res.status(400).send("OAuth client is no longer valid.");
      return;
    }

    const approvalConsumed = await this.replayStore.consume(
      "oauth-approval",
      transaction,
      claims.exp ?? Math.floor(Date.now() / 1000) + AUTH_REQUEST_TTL_SECONDS,
    );
    if (!approvalConsumed) {
      res.status(400).send("This authorization request has already been used.");
      return;
    }

    const now = Math.floor(Date.now() / 1000);
    const code = await new EncryptJWT({
      client_id: clientId,
      redirect_uri: redirectUri,
      code_challenge: asString(claims, "code_challenge"),
      resource: this.resource,
      scope: asString(claims, "scope"),
    })
      .setProtectedHeader({ alg: "dir", enc: "A256GCM", typ: "oauth-code+jwt" })
      .setIssuer(this.issuer)
      .setAudience("oauth-code")
      .setSubject(this.config.profile.id)
      .setIssuedAt(now)
      .setExpirationTime(now + CODE_TTL_SECONDS)
      .setJti(randomUUID())
      .encrypt(this.encryptedTokenKey);

    const target = new URL(redirectUri);
    target.searchParams.set("code", code);
    if (state) target.searchParams.set("state", state);
    res.redirect(302, target.href);
  };

  async challengeForAuthorizationCode(
    client: OAuthClientInformationFull,
    authorizationCode: string,
  ): Promise<string> {
    const claims = await this.decryptCode(authorizationCode);
    if (asString(claims, "client_id") !== client.client_id) {
      throw new InvalidGrantError("Authorization code was issued to another client");
    }
    return asString(claims, "code_challenge");
  }

  async exchangeAuthorizationCode(
    client: OAuthClientInformationFull,
    authorizationCode: string,
    _codeVerifier?: string,
    redirectUri?: string,
    resource?: URL,
  ): Promise<OAuthTokens> {
    const claims = await this.decryptCode(authorizationCode);
    if (asString(claims, "client_id") !== client.client_id) {
      throw new InvalidGrantError("Authorization code was issued to another client");
    }
    if (!redirectUri || asString(claims, "redirect_uri") !== redirectUri) {
      throw new InvalidGrantError("redirect_uri does not match the authorization request");
    }
    this.assertResource(resource, asString(claims, "resource"));

    const consumed = await this.replayStore.consume(
      "oauth-code",
      authorizationCode,
      claims.exp ?? Math.floor(Date.now() / 1000) + CODE_TTL_SECONDS,
    );
    if (!consumed) throw new InvalidGrantError("Authorization code has already been used");

    return this.issueTokens(client.client_id, asString(claims, "scope").split(" "));
  }

  async exchangeRefreshToken(
    client: OAuthClientInformationFull,
    refreshToken: string,
    requestedScopes?: string[],
    resource?: URL,
  ): Promise<OAuthTokens> {
    let payload: JWTPayload;
    try {
      const result = await jwtDecrypt(refreshToken, this.encryptedTokenKey, {
        issuer: this.issuer,
        audience: "oauth-refresh",
        keyManagementAlgorithms: ["dir"],
        contentEncryptionAlgorithms: ["A256GCM"],
      });
      payload = result.payload;
    } catch {
      throw new InvalidGrantError("Invalid or expired refresh token");
    }

    if (asString(payload, "client_id") !== client.client_id) {
      throw new InvalidGrantError("Refresh token was issued to another client");
    }
    this.assertResource(resource, asString(payload, "resource"));
    if (await this.replayStore.contains("revoked-token", refreshToken)) {
      throw new InvalidGrantError("Refresh token has been revoked");
    }

    const originalScopes = asString(payload, "scope").split(" ").filter(Boolean);
    const scopes = requestedScopes?.length ? requestedScopes : originalScopes;
    if (scopes.some((scope) => !originalScopes.includes(scope))) {
      throw new InvalidScopeError("Refresh token cannot be upgraded to additional scopes");
    }

    const rotated = await this.replayStore.consume(
      "used-refresh-token",
      refreshToken,
      payload.exp ?? Math.floor(Date.now() / 1000) + this.config.limits.refreshTokenTtlSeconds,
    );
    if (!rotated) throw new InvalidGrantError("Refresh token has already been used");
    return this.issueTokens(client.client_id, scopes);
  }

  async verifyAccessToken(token: string): Promise<AuthInfo> {
    let payload: JWTPayload;
    try {
      const result = await jwtVerify(token, this.accessKey, {
        issuer: this.issuer,
        audience: this.resource,
        algorithms: ["HS256"],
      });
      payload = result.payload;
    } catch {
      throw new InvalidTokenError("Invalid or expired access token");
    }

    if (await this.replayStore.contains("revoked-token", token)) {
      throw new InvalidTokenError("Access token has been revoked");
    }

    const clientId = asString(payload, "client_id");
    const scopes = asString(payload, "scope").split(" ").filter(Boolean);
    if (!payload.exp || !payload.sub) throw new InvalidTokenError("Invalid access token claims");

    return {
      token,
      clientId,
      scopes,
      expiresAt: payload.exp,
      resource: new URL(this.resource),
      extra: {
        sub: payload.sub,
        ...(typeof payload.name === "string" ? { name: payload.name } : {}),
        ...(typeof payload.email === "string" ? { email: payload.email } : {}),
      },
    };
  }

  async revokeToken(
    _client: OAuthClientInformationFull,
    request: OAuthTokenRevocationRequest,
  ): Promise<void> {
    let expiresAt = Math.floor(Date.now() / 1000) + this.config.limits.refreshTokenTtlSeconds;
    try {
      const access = await jwtVerify(request.token, this.accessKey, {
        issuer: this.issuer,
        audience: this.resource,
        algorithms: ["HS256"],
      });
      if (access.payload.exp) expiresAt = access.payload.exp;
    } catch {
      try {
        const refresh = await jwtDecrypt(request.token, this.encryptedTokenKey, {
          issuer: this.issuer,
          audience: "oauth-refresh",
          keyManagementAlgorithms: ["dir"],
          contentEncryptionAlgorithms: ["A256GCM"],
        });
        if (refresh.payload.exp) expiresAt = refresh.payload.exp;
      } catch {
        return;
      }
    }
    await this.replayStore.consume("revoked-token", request.token, expiresAt);
  }

  private async issueTokens(clientId: string, scopesInput: string[]): Promise<OAuthTokens> {
    const scopes = this.normalizeScopes(scopesInput);
    const now = Math.floor(Date.now() / 1000);
    const commonClaims = {
      client_id: clientId,
      resource: this.resource,
      scope: scopes.join(" "),
      name: this.config.profile.name,
      ...(this.config.profile.email ? { email: this.config.profile.email } : {}),
    };

    const accessToken = await new SignJWT(commonClaims)
      .setProtectedHeader({ alg: "HS256", typ: "at+jwt" })
      .setIssuer(this.issuer)
      .setAudience(this.resource)
      .setSubject(this.config.profile.id)
      .setIssuedAt(now)
      .setExpirationTime(now + this.config.limits.accessTokenTtlSeconds)
      .setJti(randomUUID())
      .sign(this.accessKey);

    const refreshToken = await new EncryptJWT(commonClaims)
      .setProtectedHeader({ alg: "dir", enc: "A256GCM", typ: "refresh+jwt" })
      .setIssuer(this.issuer)
      .setAudience("oauth-refresh")
      .setSubject(this.config.profile.id)
      .setIssuedAt(now)
      .setExpirationTime(now + this.config.limits.refreshTokenTtlSeconds)
      .setJti(randomUUID())
      .encrypt(this.encryptedTokenKey);

    return {
      access_token: accessToken,
      token_type: "Bearer",
      expires_in: this.config.limits.accessTokenTtlSeconds,
      refresh_token: refreshToken,
      scope: scopes.join(" "),
    };
  }

  private async decryptCode(code: string): Promise<JWTPayload> {
    try {
      const result = await jwtDecrypt(code, this.encryptedTokenKey, {
        issuer: this.issuer,
        audience: "oauth-code",
        keyManagementAlgorithms: ["dir"],
        contentEncryptionAlgorithms: ["A256GCM"],
      });
      return result.payload;
    } catch {
      throw new InvalidGrantError("Invalid or expired authorization code");
    }
  }

  private normalizeScopes(scopesInput?: string[]): string[] {
    const requested = (scopesInput ?? []).filter(Boolean);
    const scopes = requested.length ? [...new Set(requested)] : [...this.enabledScopes];
    const invalid = scopes.filter(
      (scope) => !ALL_SCOPES.includes(scope as (typeof ALL_SCOPES)[number]) || !this.enabledScopes.includes(scope),
    );
    if (invalid.length) throw new InvalidScopeError(`Unsupported scope: ${invalid.join(", ")}`);
    return scopes;
  }

  private assertResource(resource?: URL, encodedResource?: string): void {
    const actual = resource?.href ?? encodedResource ?? this.resource;
    if (actual !== this.resource) {
      throw new InvalidTargetError("Invalid MCP resource identifier");
    }
  }

  private renderLogin(transaction: string, scopes: string[], error?: string): string {
    const scopeLabels: Record<string, string> = {
      "mail.read": "Pesquisar e ler mensagens e anexos",
      "mail.draft": "Criar rascunhos",
      "mail.send": "Enviar mensagens após confirmação explícita",
    };
    const list = scopes.map((scope) => `<li>${escapeHtml(scopeLabels[scope] ?? scope)}</li>`).join("");
    const errorMarkup = error ? `<p role="alert">${escapeHtml(error)}</p>` : "";
    return `<!doctype html>
<html lang="pt-PT">
<head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1"><title>Autorizar Mail MCP</title></head>
<body>
<main>
  <h1>Autorizar acesso ao correio</h1>
  <p>Conta: ${escapeHtml(this.config.profile.name)}</p>
  <p>Esta ligação pede permissão para:</p>
  <ul>${list}</ul>
  ${errorMarkup}
  <form method="post" action="${escapeHtml(new URL("oauth/approve", `${this.config.baseUrl.href.replace(/\/$/, "")}/`).href)}">
    <input type="hidden" name="transaction" value="${escapeHtml(transaction)}">
    <label>Palavra-passe de administração <input type="password" name="password" autocomplete="current-password" required autofocus></label>
    <button type="submit" name="decision" value="approve">Autorizar</button>
    <button type="submit" name="decision" value="deny" formnovalidate>Cancelar</button>
  </form>
</main>
</body>
</html>`;
  }

  private redirectWithError(
    res: Response,
    redirectUri: string,
    state: string | undefined,
    error: AccessDeniedError,
  ): void {
    const target = new URL(redirectUri);
    target.searchParams.set("error", error.errorCode);
    target.searchParams.set("error_description", error.message);
    if (state) target.searchParams.set("state", state);
    res.redirect(302, target.href);
  }
}