import z from "zod"
import { Log } from "../util/log"

export namespace RateLimit {
  const log = Log.create({ service: "rate-limit" })

  export type RateLimitType = "short-term" | "long-term"

  export interface ParsedRateLimit {
    type: RateLimitType
    waitTime: number
    retryAt?: Date
  }

  export const SHORT_TERM_THRESHOLD_MS = 5 * 60 * 1000 // 5 minutes

  export function parseHeaders(headers: Headers): ParsedRateLimit | null {
    const retryAfter = headers.get("retry-after")
    const rateLimitReset = headers.get("x-ratelimit-reset")
    const rateLimitRemaining = headers.get("x-ratelimit-remaining")

    if (!retryAfter && !rateLimitReset) {
      return null
    }

    let waitTime = 0

    if (retryAfter) {
      const seconds = parseInt(retryAfter, 10)
      if (!isNaN(seconds)) {
        waitTime = seconds * 1000
      }
    }

    if (rateLimitReset && waitTime === 0) {
      const resetTime = parseInt(rateLimitReset, 10)
      if (!isNaN(resetTime)) {
        waitTime = (resetTime * 1000) - Date.now()
      }
    }

    if (waitTime === 0) {
      log.warn("Rate limit detected but could not parse wait time", { headers: Object.fromEntries(headers) })
      return null
    }

    log.info("Rate limit detected", { waitTime, retryAfter, rateLimitReset, rateLimitRemaining })

    return {
      type: waitTime < SHORT_TERM_THRESHOLD_MS ? "short-term" : "long-term",
      waitTime,
      retryAt: new Date(Date.now() + waitTime),
    }
  }

  export async function wait(parsed: ParsedRateLimit, signal?: AbortSignal): Promise<void> {
    const { waitTime, retryAt } = parsed

    log.info("Waiting for rate limit reset", { waitTime, retryAt })

    if (signal?.aborted) {
      throw new Error("Rate limit wait aborted")
    }

    if (waitTime <= 0 || !retryAt) {
      return
    }

    const now = Date.now()
    const waitMs = Math.max(0, retryAt.getTime() - now)

    await new Promise<void>((resolve, reject) => {
      const timeout = setTimeout(() => resolve(), waitMs)

      signal?.addEventListener("abort", () => {
        clearTimeout(timeout)
        reject(new Error("Rate limit wait aborted"))
      })
    })

    log.info("Rate limit wait completed")
  }

  export function isRateLimitError(error: any): boolean {
    if (error?.cause?.statusCode === 429) {
      return true
    }
    if (error?.statusCode === 429) {
      return true
    }
    if (error?.status === 429) {
      return true
    }
    const message = error?.message?.toLowerCase() || ""
    return message.includes("rate limit") || message.includes("too many requests")
  }

  const RateLimitConfigSchema = z.object({
    enabled: z.boolean().default(true),
    maxRetries: z.number().int().min(0).default(3),
    maxWaitTime: z.number().int().min(0).default(300_000), // 5 minutes
    backoffStrategy: z.enum(["linear", "exponential"]).default("exponential"),
    enableProviderSwitch: z.boolean().default(true),
    switchThreshold: z.number().int().min(0).default(300_000), // 5 minutes
  })

  export type Config = z.infer<typeof RateLimitConfigSchema>

  export const defaultConfig: Config = {
    enabled: true,
    maxRetries: 3,
    maxWaitTime: 300_000,
    backoffStrategy: "exponential",
    enableProviderSwitch: true,
    switchThreshold: 300_000,
  }

  export function parseConfig(config: any): Config {
    try {
      return RateLimitConfigSchema.parse(config)
    } catch (error) {
      log.warn("Invalid rate limit config, using defaults", { error, config })
      return defaultConfig
    }
  }

  export function calculateBackoff(
    attempt: number,
    strategy: Config["backoffStrategy"],
    initialDelay: number = 1000,
  ): number {
    switch (strategy) {
      case "linear":
        return initialDelay * (attempt + 1)
      case "exponential":
        return initialDelay * Math.pow(2, attempt)
      default:
        return initialDelay * Math.pow(2, attempt)
    }
  }
}
