import { FsRateLimitError, QuotaExceededError } from '../schemas/errors.js';

export interface RateLimitConfig {
  /** Cap on op-tokens per minute per sender. Default 600. */
  maxOpsPerMinutePerSender?: number;
  /** Cap on byte-tokens per hour per sender. Default 10 GiB. */
  maxBytesPerHourPerSender?: number;
  /** Per-op token cost. Default 1 per op except write=5, delete=3. */
  opCost?: Partial<Record<string, number>> & { default?: number };
}

interface Bucket {
  tokens: number;
  capacity: number;
  /** Tokens replenished per millisecond. */
  refillRate: number;
  lastRefill: number;
}

interface SenderBuckets {
  ops: Bucket;
  bytes: Bucket;
}

/**
 * Two-bucket per-sender rate limiter. Each operation consumes:
 *  - `opCost(op)` from the ops bucket.
 *  - `estimatedBytes` from the bytes bucket (post-call reconciliation possible).
 *
 * Rejection throws either `FsRateLimitError` (op-bucket) or `QuotaExceededError`
 * (bytes-bucket), each with a `retryAfterMs` hint.
 *
 * Cancellation MUST call `release()` to return reserved tokens.
 */
export class RateLimiter {
  private readonly opsCapacity: number;
  private readonly opsRefillPerMs: number;
  private readonly bytesCapacity: number;
  private readonly bytesRefillPerMs: number;
  private readonly opCost: Required<NonNullable<RateLimitConfig['opCost']>>;
  private readonly bySender = new Map<string, SenderBuckets>();

  constructor(config: RateLimitConfig = {}) {
    this.opsCapacity = config.maxOpsPerMinutePerSender ?? 600;
    this.opsRefillPerMs = this.opsCapacity / 60_000;
    this.bytesCapacity = config.maxBytesPerHourPerSender ?? 10 * 1024 * 1024 * 1024;
    this.bytesRefillPerMs = this.bytesCapacity / 3_600_000;
    this.opCost = {
      list: config.opCost?.list ?? 1,
      stat: config.opCost?.stat ?? 1,
      mkdir: config.opCost?.mkdir ?? 2,
      delete: config.opCost?.delete ?? 3,
      move: config.opCost?.move ?? 2,
      read: config.opCost?.read ?? 1,
      write: config.opCost?.write ?? 5,
      getThumbnail: config.opCost?.getThumbnail ?? 2,
      custom: config.opCost?.custom ?? 1,
      default: config.opCost?.default ?? 1,
    } as Required<NonNullable<RateLimitConfig['opCost']>>;
  }

  /**
   * Acquire tokens for an operation. Throws if either bucket is empty;
   * otherwise atomically deducts from both. `estimatedBytes` may be 0 for
   * non-I/O ops.
   */
  acquire(sender: string, op: string, estimatedBytes: number = 0): void {
    const buckets = this.getBuckets(sender);
    this.refill(buckets, Date.now());
    const cost = this.costFor(op);

    if (buckets.ops.tokens < cost) {
      const need = cost - buckets.ops.tokens;
      const retryMs = Math.ceil(need / this.opsRefillPerMs);
      throw new FsRateLimitError(`op rate limit exceeded`, retryMs);
    }
    if (estimatedBytes > 0 && buckets.bytes.tokens < estimatedBytes) {
      const need = estimatedBytes - buckets.bytes.tokens;
      const retryMs = Math.ceil(need / this.bytesRefillPerMs);
      throw new QuotaExceededError(`byte quota exceeded`, retryMs);
    }
    buckets.ops.tokens -= cost;
    if (estimatedBytes > 0) buckets.bytes.tokens -= estimatedBytes;
  }

  /**
   * Reconcile after the actual byte cost is known. `actualBytes` may be
   * higher or lower than the estimate; caps to capacity on release.
   */
  reconcile(sender: string, estimatedBytes: number, actualBytes: number): void {
    if (estimatedBytes === actualBytes) return;
    const buckets = this.getBuckets(sender);
    const delta = estimatedBytes - actualBytes;
    if (delta > 0) {
      buckets.bytes.tokens = Math.min(buckets.bytes.capacity, buckets.bytes.tokens + delta);
    } else {
      buckets.bytes.tokens = Math.max(0, buckets.bytes.tokens + delta);
    }
  }

  /** Return tokens to both buckets (cancellation cleanup). */
  release(sender: string, op: string, bytesReserved: number = 0): void {
    const buckets = this.bySender.get(sender);
    if (buckets === undefined) return;
    const cost = this.costFor(op);
    buckets.ops.tokens = Math.min(buckets.ops.capacity, buckets.ops.tokens + cost);
    if (bytesReserved > 0) {
      buckets.bytes.tokens = Math.min(buckets.bytes.capacity, buckets.bytes.tokens + bytesReserved);
    }
  }

  /** Available token snapshot — for tests/metrics. */
  snapshot(sender: string): { ops: number; bytes: number } | null {
    const buckets = this.bySender.get(sender);
    if (buckets === undefined) return null;
    this.refill(buckets, Date.now());
    return { ops: buckets.ops.tokens, bytes: buckets.bytes.tokens };
  }

  private costFor(op: string): number {
    const explicit = (this.opCost as Record<string, number>)[op];
    return explicit ?? this.opCost.default;
  }

  private getBuckets(sender: string): SenderBuckets {
    let b = this.bySender.get(sender);
    if (b !== undefined) return b;
    const now = Date.now();
    b = {
      ops: {
        tokens: this.opsCapacity,
        capacity: this.opsCapacity,
        refillRate: this.opsRefillPerMs,
        lastRefill: now,
      },
      bytes: {
        tokens: this.bytesCapacity,
        capacity: this.bytesCapacity,
        refillRate: this.bytesRefillPerMs,
        lastRefill: now,
      },
    };
    this.bySender.set(sender, b);
    return b;
  }

  private refill(buckets: SenderBuckets, now: number): void {
    const refillBucket = (b: Bucket): void => {
      const elapsed = now - b.lastRefill;
      if (elapsed <= 0) return;
      b.tokens = Math.min(b.capacity, b.tokens + elapsed * b.refillRate);
      b.lastRefill = now;
    };
    refillBucket(buckets.ops);
    refillBucket(buckets.bytes);
  }
}
