import type { Context, MiddlewareHandler } from 'hono'

import { fail } from '../http.js'

export interface RateLimitNamespace {
  get(key: string): Promise<string | null>
  put(
    key: string,
    value: string,
    options?: {
      expirationTtl?: number
    },
  ): Promise<void>
}

interface RateLimitState {
  count: number
  resetAt: number
}

interface RateLimitOptions<TBindings extends object> {
  key: (context: Context<{ Bindings: TBindings }>) => string
  limit: number
  namespace: (bindings: TBindings) => RateLimitNamespace | undefined
  now?: () => number
  prefix?: string
  windowMs: number
}

function parseState(raw: string | null): RateLimitState | null {
  if (!raw) {
    return null
  }

  try {
    const parsed = JSON.parse(raw) as Partial<RateLimitState>

    if (typeof parsed.count !== 'number' || typeof parsed.resetAt !== 'number') {
      return null
    }

    return {
      count: parsed.count,
      resetAt: parsed.resetAt,
    }
  } catch {
    return null
  }
}

export function createRateLimit<TBindings extends object>(
  options: RateLimitOptions<TBindings>,
): MiddlewareHandler<{ Bindings: TBindings }> {
  const now = options.now ?? Date.now
  const prefix = options.prefix ?? 'rate-limit'

  return async (context, next) => {
    const namespace = options.namespace(context.env)

    if (!namespace) {
      await next()
      return
    }

    const currentTime = now()
    const requestKey = `${prefix}:${options.key(context)}`
    const currentState = parseState(await namespace.get(requestKey))
    const state =
      !currentState || currentState.resetAt <= currentTime
        ? { count: 0, resetAt: currentTime + options.windowMs }
        : currentState

    if (state.count >= options.limit) {
      return fail('RATE_LIMITED', 'Rate limit exceeded', 429)
    }

    const nextState: RateLimitState = {
      count: state.count + 1,
      resetAt: state.resetAt,
    }

    await namespace.put(requestKey, JSON.stringify(nextState), {
      expirationTtl: Math.max(1, Math.ceil((nextState.resetAt - currentTime) / 1000)),
    })

    await next()
  }
}