/**
 * A2aLimiter: process-local rate / budget / concurrency counters for A2A grant enforcement.
 *
 * IMPORTANT — process-local, resets on restart:
 *   All counters live in-process memory and are lost when the server restarts.
 *   This is intentional and acceptable because:
 *   - Rate / skill-budget counters: brief gaps from restarts are fail-open by design.
 *   - Concurrency reservation: the durable backstop is the `dbCount` argument passed to
 *     tryReserveConcurrency(). The in-memory counter is an additional synchronous guard
 *     against TOCTOU races within a single process instance.
 *
 * Keyed by grantId (UUID from the a2a_delegations table).
 */

import { logger } from '../../logger.js';

// ---------------------------------------------------------------------------
// Public types
// ---------------------------------------------------------------------------

export interface A2aLimits {
  ratePerMinute: number;
  maxConcurrentTasks: number;
  maxPayloadBytes: number;
  maxStreamSeconds: number;
  skillBudgetPerHour: number;
  maxConcurrentResubscribe: number;
}

export const DEFAULT_A2A_LIMITS: A2aLimits = {
  ratePerMinute: 30,
  maxConcurrentTasks: 5,
  maxPayloadBytes: 65536,
  maxStreamSeconds: 600,
  skillBudgetPerHour: 100,
  maxConcurrentResubscribe: 5,
};

// ---------------------------------------------------------------------------
// resolveA2aLimits — strict integer≥1 validation
// ---------------------------------------------------------------------------

function posInt(v: unknown, fallback: number): number {
  return typeof v === 'number' && Number.isInteger(v) && v >= 1 ? v : fallback;
}

export function resolveA2aLimits(p?: Partial<Record<keyof A2aLimits, unknown>>): A2aLimits {
  const src = p ?? {};
  const out = {} as A2aLimits;
  for (const k of Object.keys(DEFAULT_A2A_LIMITS) as (keyof A2aLimits)[]) {
    const resolved = posInt(src[k], DEFAULT_A2A_LIMITS[k]);
    if (src[k] !== undefined && resolved !== src[k]) {
      logger.warn(`[a2a] invalid limit ${k}=${String(src[k])}, using default ${resolved}`);
    }
    out[k] = resolved;
  }
  return out;
}

// ---------------------------------------------------------------------------
// Internal state shapes
// ---------------------------------------------------------------------------

interface RateBucket {
  tokens: number;
  lastRefillMs: number;
}

const ONE_HOUR_MS = 3_600_000;

// ---------------------------------------------------------------------------
// A2aLimiter
// ---------------------------------------------------------------------------

export class A2aLimiter {
  readonly limits: A2aLimits;
  private readonly _now: () => number;

  /** Token buckets for rate limiting: grantId → { tokens, lastRefillMs } */
  private readonly _rateBuckets = new Map<string, RateBucket>();
  /** Rolling-hour timestamps for skill budget: grantId → timestamp[] */
  private readonly _skillHits = new Map<string, number[]>();
  /** In-memory concurrency reservations: grantId → count */
  private readonly _reservations = new Map<string, number>();
  /** Resubscribe slot counters: grantId → count */
  private readonly _resubscribes = new Map<string, number>();

  constructor(limits: A2aLimits, now: () => number = () => Date.now()) {
    this.limits = limits;
    this._now = now;
  }

  // -------------------------------------------------------------------------
  // Rate limiting (token bucket)
  // -------------------------------------------------------------------------

  tryConsumeRate(grantId: string): boolean {
    const nowMs = this._now();
    const capacity = this.limits.ratePerMinute;
    let bucket = this._rateBuckets.get(grantId);

    if (bucket) {
      const elapsed = Math.max(0, nowMs - bucket.lastRefillMs);
      bucket.tokens = Math.min(capacity, bucket.tokens + elapsed * capacity / 60_000);

      // Opportunistic eviction: full bucket idle for >1h → drop stale entry.
      // The next access will start fresh (same observable effect as a full refill).
      if (bucket.tokens >= capacity && elapsed >= ONE_HOUR_MS) {
        this._rateBuckets.delete(grantId);
        bucket = undefined;
      } else {
        bucket.lastRefillMs = nowMs;
      }
    }

    if (!bucket) {
      bucket = { tokens: capacity, lastRefillMs: nowMs };
      this._rateBuckets.set(grantId, bucket);
    }

    if (bucket.tokens >= 1) {
      bucket.tokens -= 1;
      return true;
    }
    return false;
  }

  // -------------------------------------------------------------------------
  // Skill budget (rolling-hour window)
  // -------------------------------------------------------------------------

  tryConsumeSkillBudget(grantId: string): boolean {
    const nowMs = this._now();
    const cutoff = nowMs - ONE_HOUR_MS;
    const existing = this._skillHits.get(grantId) ?? [];

    // Prune timestamps older than 1h (keep only recent ones)
    const hits = existing.filter(ts => ts > cutoff);

    if (hits.length < this.limits.skillBudgetPerHour) {
      hits.push(nowMs);
      this._skillHits.set(grantId, hits);
      return true;
    }

    // Budget exhausted. Update the pruned array; evict if empty (defensive).
    if (hits.length > 0) {
      this._skillHits.set(grantId, hits);
    } else {
      this._skillHits.delete(grantId);
    }
    return false;
  }

  // -------------------------------------------------------------------------
  // Concurrency reservation (synchronous TOCTOU guard)
  // -------------------------------------------------------------------------

  /**
   * Try to reserve one concurrency slot for grantId.
   * @param dbCount  Current running-job count from the DB (durable backstop).
   * Returns true and increments the in-memory counter if
   * Math.max(dbCount, reservedCount) < maxConcurrentTasks.
   * This is a plain synchronous method — no awaits.
   */
  tryReserveConcurrency(grantId: string, dbCount: number): boolean {
    const reserved = this._reservations.get(grantId) ?? 0;
    if (Math.max(dbCount, reserved) >= this.limits.maxConcurrentTasks) return false;
    this._reservations.set(grantId, reserved + 1);
    return true;
  }

  /** Decrement the concurrency reservation for grantId, flooring at 0. */
  releaseConcurrency(grantId: string): void {
    const n = this._reservations.get(grantId) ?? 0;
    if (n <= 1) {
      this._reservations.delete(grantId);
    } else {
      this._reservations.set(grantId, n - 1);
    }
  }

  // -------------------------------------------------------------------------
  // Resubscribe slot tracking
  // -------------------------------------------------------------------------

  tryAcquireResubscribe(grantId: string): boolean {
    const current = this._resubscribes.get(grantId) ?? 0;
    if (current >= this.limits.maxConcurrentResubscribe) return false;
    this._resubscribes.set(grantId, current + 1);
    return true;
  }

  releaseResubscribe(grantId: string): void {
    const n = this._resubscribes.get(grantId) ?? 0;
    if (n <= 1) {
      this._resubscribes.delete(grantId);
    } else {
      this._resubscribes.set(grantId, n - 1);
    }
  }

  // -------------------------------------------------------------------------
  // Test hooks (Map size inspection for eviction tests)
  // -------------------------------------------------------------------------

  /** Number of live rate-bucket entries (test hook). */
  rateBucketSize(): number {
    return this._rateBuckets.size;
  }

  /** Number of live skill-hits entries (test hook). */
  skillHitsSize(): number {
    return this._skillHits.size;
  }
}
