/**
 * Token Repository for SentryAgent.ai AgentIdP.
 * Manages token revocation tracking (Redis primary, PostgreSQL fallback)
 * and monthly token count tracking (Redis with monthly TTL).
 */

import { Pool, QueryResult } from 'pg';
import { RedisClientType } from 'redis';

/** Raw database row for a token revocation record. */
interface TokenRevocationRow {
  jti: string;
  expires_at: Date;
  revoked_at: Date;
}

/**
 * Repository for token revocation and monthly usage tracking.
 * Receives a pg.Pool and RedisClientType via constructor injection.
 */
export class TokenRepository {
  /**
   * @param pool - The PostgreSQL connection pool.
   * @param redis - The Redis client.
   */
  constructor(
    private readonly pool: Pool,
    private readonly redis: RedisClientType,
  ) {}

  /**
   * Adds a token JTI to both the Redis revocation list and the PostgreSQL
   * token_revocations table. Redis TTL is set to the token's remaining lifetime.
   *
   * @param jti - The JWT ID to revoke.
   * @param expiresAt - The token expiry date (used to calculate Redis TTL).
   */
  async addToRevocationList(jti: string, expiresAt: Date): Promise<void> {
    const nowSeconds = Math.floor(Date.now() / 1000);
    const expirySeconds = Math.floor(expiresAt.getTime() / 1000);
    const ttl = Math.max(1, expirySeconds - nowSeconds);

    const redisKey = `revoked:${jti}`;
    await this.redis.set(redisKey, '1', { EX: ttl });

    await this.pool.query(
      `INSERT INTO token_revocations (jti, expires_at)
       VALUES ($1, $2)
       ON CONFLICT (jti) DO NOTHING`,
      [jti, expiresAt],
    );
  }

  /**
   * Checks whether a JTI has been revoked.
   * Checks Redis first for performance; falls back to PostgreSQL.
   *
   * @param jti - The JWT ID to check.
   * @returns True if the token has been revoked, false otherwise.
   */
  async isRevoked(jti: string): Promise<boolean> {
    const redisKey = `revoked:${jti}`;
    const cached = await this.redis.get(redisKey);
    if (cached !== null) return true;

    const result: QueryResult<TokenRevocationRow> = await this.pool.query(
      'SELECT jti FROM token_revocations WHERE jti = $1',
      [jti],
    );
    return result.rows.length > 0;
  }

  /**
   * Increments the monthly token count for a client in Redis.
   * Sets a TTL to the end of the current month if the key is new.
   *
   * @param clientId - The agent/client ID.
   * @returns The new count after incrementing.
   */
  async incrementMonthlyCount(clientId: string): Promise<number> {
    const now = new Date();
    const year = now.getUTCFullYear();
    const month = String(now.getUTCMonth() + 1).padStart(2, '0');
    const key = `monthly:tokens:${clientId}:${year}-${month}`;

    const count = await this.redis.incr(key);

    if (count === 1) {
      // Set TTL to end of month
      const endOfMonth = new Date(Date.UTC(year, now.getUTCMonth() + 1, 1, 0, 0, 0));
      const ttlSeconds = Math.floor((endOfMonth.getTime() - Date.now()) / 1000);
      await this.redis.expire(key, Math.max(1, ttlSeconds));
    }

    return count;
  }

  /**
   * Gets the current monthly token count for a client.
   *
   * @param clientId - The agent/client ID.
   * @returns The current count, or 0 if no tokens have been issued this month.
   */
  async getMonthlyCount(clientId: string): Promise<number> {
    const now = new Date();
    const year = now.getUTCFullYear();
    const month = String(now.getUTCMonth() + 1).padStart(2, '0');
    const key = `monthly:tokens:${clientId}:${year}-${month}`;

    const value = await this.redis.get(key);
    return value !== null ? parseInt(value, 10) : 0;
  }
}
