type WindowCounter = {
  count: number
  resetAt: number
}

const ipRequestStore = new Map<string, WindowCounter>()
const sessionRequestStore = new Map<string, WindowCounter>()
const sessionOutputStore = new Map<string, WindowCounter>()

function now() {
  return Date.now()
}

function cleanup(store: Map<string, WindowCounter>) {
  const currentTime = now()

  for (const [key, value] of store.entries()) {
    if (value.resetAt <= currentTime) {
      store.delete(key)
    }
  }
}

function readOrCreate(
  store: Map<string, WindowCounter>,
  key: string,
  windowMs: number
) {
  cleanup(store)

  const currentTime = now()
  const existing = store.get(key)

  if (!existing || existing.resetAt <= currentTime) {
    const fresh = {
      count: 0,
      resetAt: currentTime + windowMs,
    }

    store.set(key, fresh)
    return fresh
  }

  return existing
}

function consume(
  store: Map<string, WindowCounter>,
  key: string,
  windowMs: number,
  amount: number
) {
  const bucket = readOrCreate(store, key, windowMs)
  bucket.count += amount
  store.set(key, bucket)
  return bucket
}

function peek(
  store: Map<string, WindowCounter>,
  key: string,
  windowMs: number
) {
  return readOrCreate(store, key, windowMs)
}

export function getChatRateLimitStatus({
  ip,
  maxIpRequests,
  maxSessionRequests,
  maxSessionOutputChars,
  requestWindowMs,
  sessionId,
  outputWindowMs,
}: {
  ip: string
  maxIpRequests: number
  maxSessionRequests: number
  maxSessionOutputChars: number
  outputWindowMs: number
  requestWindowMs: number
  sessionId: string
}) {
  const ipBucket = peek(ipRequestStore, ip, requestWindowMs)
  const sessionBucket = peek(sessionRequestStore, sessionId, requestWindowMs)
  const outputBucket = peek(sessionOutputStore, sessionId, outputWindowMs)

  const ipRemaining = Math.max(0, maxIpRequests - ipBucket.count)
  const sessionRemaining = Math.max(0, maxSessionRequests - sessionBucket.count)
  const outputCharsRemaining = Math.max(
    0,
    maxSessionOutputChars - outputBucket.count
  )

  return {
    ipRemaining,
    sessionRemaining,
    outputCharsRemaining,
    requestResetAt: new Date(
      Math.max(ipBucket.resetAt, sessionBucket.resetAt)
    ).toISOString(),
    outputResetAt: new Date(outputBucket.resetAt).toISOString(),
    blocked:
      ipRemaining <= 0 || sessionRemaining <= 0 || outputCharsRemaining <= 0,
  }
}

export function consumeChatRequest({
  ip,
  requestWindowMs,
  sessionId,
}: {
  ip: string
  requestWindowMs: number
  sessionId: string
}) {
  consume(ipRequestStore, ip, requestWindowMs, 1)
  consume(sessionRequestStore, sessionId, requestWindowMs, 1)
}

export function consumeChatOutput({
  chars,
  outputWindowMs,
  sessionId,
}: {
  chars: number
  outputWindowMs: number
  sessionId: string
}) {
  consume(sessionOutputStore, sessionId, outputWindowMs, chars)
}