import { randomUUID } from "node:crypto";
import path from "node:path";
import { EncryptJWT, jwtDecrypt } from "jose";
import type { AppConfig } from "../config.js";
import { deriveKey } from "../security/keys.js";
import { ReplayStore } from "../security/replay-store.js";
import { ComposeMessageSchema, type ComposeMessage } from "./schemas.js";

export type PreparedMessage = {
  message: ComposeMessage;
  jti: string;
  expiresAt: number;
};

export class ActionTokenService {
  private readonly key: Uint8Array;
  private readonly replayStore: ReplayStore;

  constructor(private readonly config: AppConfig) {
    this.key = deriveKey(config.tokenSecret, "prepared-email-action");
    this.replayStore = new ReplayStore(path.join(config.stateDir, "mail-actions.json"));
  }

  async prepare(messageInput: ComposeMessage): Promise<{ token: string; expiresAt: number }> {
    const message = ComposeMessageSchema.parse(messageInput);
    const now = Math.floor(Date.now() / 1000);
    const expiresAt = now + this.config.limits.preparedEmailTtlSeconds;
    const token = await new EncryptJWT({ message })
      .setProtectedHeader({ alg: "dir", enc: "A256GCM", typ: "mail-action+jwt" })
      .setIssuer("cpanel-mail-mcp")
      .setAudience("send-prepared-email")
      .setIssuedAt(now)
      .setExpirationTime(expiresAt)
      .setJti(randomUUID())
      .encrypt(this.key);
    return { token, expiresAt };
  }

  async decode(token: string): Promise<PreparedMessage> {
    try {
      const result = await jwtDecrypt(token, this.key, {
        issuer: "cpanel-mail-mcp",
        audience: "send-prepared-email",
        keyManagementAlgorithms: ["dir"],
        contentEncryptionAlgorithms: ["A256GCM"],
      });
      const message = ComposeMessageSchema.parse(result.payload.message);
      if (!result.payload.jti || !result.payload.exp) throw new Error("Missing token claims");
      return {
        message,
        jti: result.payload.jti,
        expiresAt: result.payload.exp,
      };
    } catch {
      throw new Error("Prepared email token is invalid or expired; prepare the email again");
    }
  }

  async reserveForSending(prepared: PreparedMessage): Promise<void> {
    const reserved = await this.replayStore.consume(
      "send-action",
      prepared.jti,
      prepared.expiresAt,
    );
    if (!reserved) throw new Error("This prepared email has already been used");
  }
}