/**
 * OIDCKeyService — manages OIDC signing key pairs for ID token issuance.
 * Keys are RSA-2048 (RS256) or EC P-256 (ES256) depending on OIDC_KEY_ALGORITHM env var.
 * Private keys are stored in Vault KV v2; public keys in the oidc_keys table.
 * Only one key is `is_current = true` at a time.
 *
 * Key storage strategy (mirrors DIDService pattern):
 * - When VAULT_ADDR + VAULT_TOKEN are set: private key is stored in Vault KV v2 at
 *   `{mount}/data/agentidp/oidc/keys/{kid}`.
 * - When Vault is not configured (dev mode): `vault_key_path` column stores the marker
 *   `dev:no-vault` and the private key is held in-memory only (ephemeral).
 */

import { Pool, QueryResult } from 'pg';
import { generateKeyPairSync, KeyObject } from 'crypto';
import nodeVault from 'node-vault';
import { RedisClientType } from 'redis';

import { IOIDCKey, IJWKSKey, IJWKSResponse } from '../types/oidc.js';
import { SentryAgentError } from '../utils/errors.js';

/** Raw database row for oidc_keys. */
interface OIDCKeyRow {
  id: string;
  kid: string;
  algorithm: string;
  public_key_jwk: IJWKSKey;
  vault_key_path: string;
  is_current: boolean;
  created_at: Date;
  expires_at: Date;
}

/** Dev-mode in-memory store for private keys (not persisted). */
const devPrivateKeys = new Map<string, string>();

/**
 * 500 — No current OIDC signing key is available.
 */
class OIDCKeyNotFoundError extends SentryAgentError {
  constructor() {
    super(
      'No current OIDC signing key is configured. Call ensureCurrentKey() first.',
      'OIDC_KEY_NOT_FOUND',
      500,
    );
  }
}

/**
 * Maps a raw database row to the IOIDCKey domain model.
 *
 * @param row - Raw row from the oidc_keys table.
 * @returns Typed IOIDCKey object.
 */
function mapRowToOIDCKey(row: OIDCKeyRow): IOIDCKey {
  return {
    id: row.id,
    kid: row.kid,
    algorithm: row.algorithm,
    public_key_jwk: row.public_key_jwk,
    vault_key_path: row.vault_key_path,
    is_current: row.is_current,
    created_at: row.created_at,
    expires_at: row.expires_at,
  };
}

/**
 * Service that manages RSA-2048 or EC P-256 signing key pairs for OIDC ID token issuance.
 * Integrates with Vault for private key storage and Redis for JWKS caching.
 */
export class OIDCKeyService {
  /**
   * @param pool - PostgreSQL connection pool.
   * @param redis - Redis client for JWKS caching.
   */
  constructor(
    private readonly pool: Pool,
    private readonly redis: RedisClientType,
  ) {}

  // ─────────────────────────────────────────────────────────────────────────────
  // Public API
  // ─────────────────────────────────────────────────────────────────────────────

  /**
   * Ensures a current signing key exists. If none exists, generates one.
   * Idempotent — safe to call on every application startup.
   *
   * @returns Promise that resolves when a current key is guaranteed to exist.
   */
  async ensureCurrentKey(): Promise<void> {
    const result: QueryResult<{ count: string }> = await this.pool.query(
      `SELECT COUNT(*) AS count FROM oidc_keys WHERE is_current = TRUE`,
    );
    const count = parseInt(result.rows[0].count, 10);
    if (count === 0) {
      await this.generateSigningKeyPair();
    }
  }

  /**
   * Generates a new RSA-2048 (RS256) or EC P-256 (ES256) signing key pair.
   * Stores the private key in Vault (or in-memory for dev mode).
   * Stores the public key in the oidc_keys table and sets it as the current key.
   * The previously current key is demoted to is_current = false.
   *
   * Algorithm is controlled by the OIDC_KEY_ALGORITHM environment variable (default: RS256).
   *
   * @returns The newly created IOIDCKey record.
   */
  async generateSigningKeyPair(): Promise<IOIDCKey> {
    const algorithm = this.getAlgorithm();
    const kid = this.buildKid();
    const ttlSeconds = this.getTokenTTLSeconds();

    const { publicKey, privateKey } = this.generateKeyPair(algorithm);

    const publicKeyJwk = this.exportPublicJWK(publicKey, algorithm, kid);
    const privateKeyPem = privateKey.export({ format: 'pem', type: 'pkcs8' }) as string;

    const vaultKeyPath = await this.storePrivateKey(kid, privateKeyPem);

    // Demote the previous current key
    await this.pool.query(
      `UPDATE oidc_keys SET is_current = FALSE WHERE is_current = TRUE`,
    );

    // Insert the new key
    const expiresAt = new Date(Date.now() + ttlSeconds * 1000);
    const insertResult: QueryResult<OIDCKeyRow> = await this.pool.query(
      `INSERT INTO oidc_keys (kid, algorithm, public_key_jwk, vault_key_path, is_current, expires_at)
       VALUES ($1, $2, $3, $4, TRUE, $5)
       RETURNING *`,
      [kid, algorithm, JSON.stringify(publicKeyJwk), vaultKeyPath, expiresAt],
    );

    // Invalidate JWKS cache
    await this.invalidateJWKSCache();

    return mapRowToOIDCKey(insertResult.rows[0]);
  }

  /**
   * Returns the current signing key from the database.
   *
   * @returns The current IOIDCKey.
   * @throws OIDCKeyNotFoundError if no current key is configured.
   */
  async getCurrentKey(): Promise<IOIDCKey> {
    const result: QueryResult<OIDCKeyRow> = await this.pool.query(
      `SELECT * FROM oidc_keys WHERE is_current = TRUE LIMIT 1`,
    );
    if (result.rows.length === 0) {
      throw new OIDCKeyNotFoundError();
    }
    return mapRowToOIDCKey(result.rows[0]);
  }

  /**
   * Returns the private key PEM for a given key ID.
   * Fetches from Vault if configured; falls back to dev in-memory store.
   *
   * @param kid - The key ID.
   * @param vaultKeyPath - The Vault path recorded in the database.
   * @returns The PEM-encoded private key string.
   * @throws Error if the private key cannot be retrieved.
   */
  async getPrivateKeyPem(kid: string, vaultKeyPath: string): Promise<string> {
    if (vaultKeyPath === 'dev:no-vault') {
      const pem = devPrivateKeys.get(kid);
      if (!pem) {
        throw new Error(
          `Dev mode: private key for kid "${kid}" not found in memory. Was the process restarted?`,
        );
      }
      return pem;
    }

    const vaultAddr = process.env['VAULT_ADDR'];
    const vaultToken = process.env['VAULT_TOKEN'];
    if (!vaultAddr || !vaultToken) {
      throw new Error('VAULT_ADDR and VAULT_TOKEN are required to read private keys from Vault.');
    }

    const vault = nodeVault({ endpoint: vaultAddr, token: vaultToken });
    const response = await vault.read(vaultKeyPath) as { data: { data: { privateKeyPem: string } } };
    return response.data.data.privateKeyPem;
  }

  /**
   * Returns the JWKS (JSON Web Key Set) for all non-expired keys.
   * Includes older keys that are past is_current=false but within their expires_at window,
   * so consumers can still verify tokens issued before the last rotation.
   * Result is cached in Redis with TTL from OIDC_JWKS_CACHE_TTL_SECONDS (default: 3600).
   *
   * @returns IJWKSResponse containing all valid public keys.
   */
  async getPublicJWKS(): Promise<IJWKSResponse> {
    const cacheKey = 'oidc:jwks';
    const cached = await this.redis.get(cacheKey);
    if (cached !== null) {
      return JSON.parse(cached) as IJWKSResponse;
    }

    const result: QueryResult<OIDCKeyRow> = await this.pool.query(
      `SELECT * FROM oidc_keys WHERE expires_at > NOW() ORDER BY created_at DESC`,
    );

    const jwks: IJWKSResponse = {
      keys: result.rows.map((row) => row.public_key_jwk),
    };

    const ttl = parseInt(process.env['OIDC_JWKS_CACHE_TTL_SECONDS'] ?? '3600', 10);
    await this.redis.set(cacheKey, JSON.stringify(jwks), { EX: ttl });

    return jwks;
  }

  /**
   * Generates a new key pair and promotes it to current.
   * The old current key is demoted (is_current = false) but remains in the oidc_keys table
   * until its expires_at passes — consumers can still verify tokens signed with it.
   *
   * @returns The new current IOIDCKey.
   */
  async rotateKey(): Promise<IOIDCKey> {
    return this.generateSigningKeyPair();
  }

  /**
   * Deletes all keys whose expires_at is in the past from both the database and Vault.
   * Should be called periodically (e.g. via a cron job) to prevent table bloat.
   *
   * @returns Promise that resolves when all expired keys have been pruned.
   */
  async pruneExpiredKeys(): Promise<void> {
    const result: QueryResult<OIDCKeyRow> = await this.pool.query(
      `DELETE FROM oidc_keys WHERE expires_at <= NOW() RETURNING *`,
    );

    // Remove private keys from Vault for each pruned key
    const vaultAddr = process.env['VAULT_ADDR'];
    const vaultToken = process.env['VAULT_TOKEN'];

    for (const row of result.rows) {
      if (row.vault_key_path === 'dev:no-vault') {
        devPrivateKeys.delete(row.kid);
        continue;
      }
      if (vaultAddr && vaultToken) {
        try {
          const vault = nodeVault({ endpoint: vaultAddr, token: vaultToken });
          await vault.delete(row.vault_key_path);
        } catch {
          // Best-effort cleanup — log but do not throw
          console.warn(`[OIDCKeyService] Failed to delete Vault key at ${row.vault_key_path}`);
        }
      }
    }

    // Invalidate cache if any keys were pruned
    if (result.rows.length > 0) {
      await this.invalidateJWKSCache();
    }
  }

  // ─────────────────────────────────────────────────────────────────────────────
  // Private helpers
  // ─────────────────────────────────────────────────────────────────────────────

  /**
   * Reads the OIDC key algorithm from environment.
   *
   * @returns "RS256" or "ES256".
   */
  private getAlgorithm(): string {
    const alg = process.env['OIDC_KEY_ALGORITHM'] ?? 'RS256';
    if (alg !== 'RS256' && alg !== 'ES256') {
      throw new Error(`Unsupported OIDC_KEY_ALGORITHM: "${alg}". Must be RS256 or ES256.`);
    }
    return alg;
  }

  /**
   * Reads the ID token TTL from environment.
   *
   * @returns Token TTL in seconds (default: 3600).
   */
  private getTokenTTLSeconds(): number {
    return parseInt(process.env['OIDC_ID_TOKEN_TTL_SECONDS'] ?? '3600', 10);
  }

  /**
   * Builds a deterministic key ID based on the current timestamp.
   *
   * @returns A unique kid string.
   */
  private buildKid(): string {
    const now = new Date();
    const date = now.toISOString().slice(0, 10).replace(/-/g, '');
    const ms = now.getTime().toString(36).toUpperCase();
    return `key-${date}-${ms}`;
  }

  /**
   * Generates an RSA-2048 or EC P-256 key pair depending on the algorithm.
   *
   * @param algorithm - "RS256" for RSA-2048, "ES256" for EC P-256.
   * @returns The generated public and private key objects.
   */
  private generateKeyPair(algorithm: string): { publicKey: KeyObject; privateKey: KeyObject } {
    if (algorithm === 'ES256') {
      return generateKeyPairSync('ec', { namedCurve: 'P-256' });
    }
    // For RSA, use the overload that returns KeyObject (no encoding options).
    return generateKeyPairSync('rsa', { modulusLength: 2048 });
  }

  /**
   * Exports a public key as a JWK with OIDC-specific metadata fields.
   *
   * @param publicKey - The public KeyObject to export.
   * @param algorithm - "RS256" or "ES256".
   * @param kid - The key ID to embed in the JWK.
   * @returns The JWK representation of the public key.
   */
  private exportPublicJWK(publicKey: KeyObject, algorithm: string, kid: string): IJWKSKey {
    const jwk = publicKey.export({ format: 'jwk' }) as Record<string, string>;
    return {
      kid,
      kty: jwk['kty'] ?? (algorithm === 'ES256' ? 'EC' : 'RSA'),
      use: 'sig',
      alg: algorithm,
      ...(algorithm === 'RS256'
        ? { n: jwk['n'], e: jwk['e'] }
        : { crv: jwk['crv'], x: jwk['x'], y: jwk['y'] }),
    };
  }

  /**
   * Stores the private key PEM in Vault (if configured) or the dev in-memory map.
   *
   * @param kid - The key ID (used to build the Vault path and dev store key).
   * @param privateKeyPem - The PKCS#8 PEM-encoded private key.
   * @returns The Vault path where the key was stored, or "dev:no-vault".
   */
  private async storePrivateKey(kid: string, privateKeyPem: string): Promise<string> {
    const vaultAddr = process.env['VAULT_ADDR'];
    const vaultToken = process.env['VAULT_TOKEN'];

    if (vaultAddr && vaultToken) {
      const mount = process.env['VAULT_MOUNT'] ?? 'secret';
      const vaultPath = `${mount}/data/agentidp/oidc/keys/${kid}`;
      const vault = nodeVault({ endpoint: vaultAddr, token: vaultToken });
      await vault.write(vaultPath, { data: { privateKeyPem } });
      return vaultPath;
    }

    // Dev mode: private key is held in-memory only — not persisted across restarts
    devPrivateKeys.set(kid, privateKeyPem);
    return 'dev:no-vault';
  }

  /**
   * Invalidates the JWKS Redis cache entry.
   */
  private async invalidateJWKSCache(): Promise<void> {
    await this.redis.del('oidc:jwks');
  }
}
