import type { Context, Env, MiddlewareHandler } from 'hono'
import { matchedRoutes } from 'hono/route'
import { Base64 } from 'ox'

import type * as Store from './Store.js'

/**
 * Per-request edge-cache policy published by `Cache.response` for routes that
 * opted into response caching. Its presence is what marks a response as safe to
 * persist in the shared edge cache.
 */
export type Policy = {
  /** Shared max-age / s-maxage in seconds for the stored response. */
  maxAge: number
  /** Stale-while-revalidate window in seconds, if any. */
  staleWhileRevalidate?: number | undefined
}

/** Hono variables published by the edge-cache layer. */
export type Variables = {
  /** Edge-cache policy for the current route, set by `Cache.response`. */
  edgeCache?: Policy | undefined
  /** `'hit'` when the edge cache served the response without reaching the origin handlers. */
  edgeCacheStatus?: 'hit' | undefined
}

// Per-request headers that must never be persisted in a shared cache entry.
// `age`/`cf-cache-status` are CDN telemetry stamped by Cloudflare on Cache API
// hits (e.g. the inner `Cache.response` layer); storing them would freeze a
// stale `Age` into every copy served from this cache.
const volatileHeaders = [
  'age',
  'cf-cache-status',
  'content-encoding',
  'content-length',
  'payment-receipt',
  'payment-session-snapshot',
  'ratelimit-limit',
  'ratelimit-remaining',
  'ratelimit-reset',
  'ratelimit-scope',
  'server-timing',
  'set-cookie',
  'tempo-request-id',
  'www-authenticate',
]

/**
 * Serves GET responses from the edge cache store BEFORE auth and rate-limiting
 * run, so cache hits skip the per-request store round-trips that dominate
 * latency. Credential-bearing requests bypass this layer so per-key policy
 * enforcement always runs. Place this ahead of the auth middleware.
 *
 * Only 2xx responses from routes that opted into caching (via `Cache.response`,
 * which publishes the `edgeCache` policy) are stored, and the policy is only
 * published for routes whose responses are public and URL-keyed (no per-principal
 * data), so a single shared entry is safe to serve to all callers. The trade-off
 * is that cache hits bypass origin metering: they are not counted against
 * rate-limit quotas.
 *
 * Lookups are equally scoped: only requests whose matched route carries the
 * {@link markCacheable} marker (applied by `Cache.response`) pay the store
 * read, so GETs to non-cacheable routes skip the round trip entirely.
 *
 * The cache is a plain `Store.Store` — the same store the origin read cache uses
 * — holding the response serialized to a string under a URL key. It is a no-op
 * when no store is supplied.
 *
 * Standard request `Cache-Control` directives let a caller invalidate without an
 * out-of-band purge: `no-cache` forces a refresh (skip the stored hit, re-fetch,
 * overwrite the entry) and `no-store` bypasses the cache entirely. Both fall
 * through to the origin and are metered normally. The store is colo-local, so a
 * refresh only affects the data center serving that request.
 */
export function middleware<environment extends Env>(
  options: middleware.Options = {},
): MiddlewareHandler<environment> {
  return async (c, next) => {
    const { store } = options
    // Pay the store lookup only when a matched handler is marked cacheable
    // (mounted `Cache.response`); other GETs (authed management routes, docs,
    // 404s) can never hit, so the read would be a wasted round trip.
    const eligible =
      store !== undefined &&
      c.req.method === 'GET' &&
      !c.req.raw.headers.has('range') &&
      // Credential-bearing requests must reach auth for per-key policy enforcement.
      !hasCredential(c.req.raw) &&
      isCacheableRoute(c)
    const key = eligible
      ? await routeCacheKey(c).catch(() => {
          // Key generators may parse request input before route validation runs.
          // Bypass caching so the validation middleware can return its response.
          return undefined
        })
      : undefined

    // Honor the standard request `Cache-Control` directives so a caller can
    // refresh a stale entry on the colo serving them without any out-of-band
    // purge: `no-cache` skips the stored hit but still re-stores the fresh
    // response (revalidate), and `no-store` bypasses the cache on both sides.
    // `Pragma: no-cache` is the HTTP/1.0 fallback, honored only when
    // `Cache-Control` is absent (RFC 9111 §5.4). A forced miss falls through
    // to `next()` and is metered like any other origin request, so this is
    // not an unmetered-access vector.
    const directives = (c.req.header('cache-control') ?? c.req.header('pragma') ?? '').toLowerCase()
    const noStore = directives.includes('no-store')
    const noCache = noStore || directives.includes('no-cache')

    if (store && key && !noCache) {
      const hit = await store.get(key)
      if (hit) {
        // Mark the hit for the canonical request log line; routing never runs
        // for hits, so this is the only signal that the cache answered.
        c.set('edgeCacheStatus' as never, 'hit' as never)
        return serveHit(c.req.header('if-none-match'), hit)
      }
    }

    await next()

    if (!store || !key || noStore) return
    const policy = c.get('edgeCache' as never) as Policy | undefined
    if (!policy) return
    if (c.res.status < 200 || c.res.status >= 300) return
    if (c.res.headers.has('set-cookie')) return
    if (c.res.headers.get('content-type')?.includes('text/event-stream')) return
    const responseDirectives = (c.res.headers.get('cache-control') ?? '').toLowerCase()
    if (responseDirectives.includes('no-cache') || responseDirectives.includes('no-store')) return

    // Store a public, sanitized copy; the client still receives the original
    // (private, metered) response for this miss. Clone synchronously — the
    // original body locks once it streams to the client — then serialize and
    // write behind `waitUntil` so the miss response is not held open on the
    // body read.
    const copy = c.res.clone()
    const persist = serialize(copy, policy).then((value) =>
      store.put(key, value, { ttl: policy.maxAge * 1_000 }),
    )
    const waitUntil = getWaitUntil(c)
    if (waitUntil) waitUntil(persist)
    else await persist.catch(() => {})
    return
  }
}

function hasCredential(request: Request): boolean {
  return (
    request.headers.has('tempo-api-key') ||
    request.headers.has('x-api-key') ||
    request.headers.has('authorization') ||
    new URL(request.url).searchParams.has('key')
  )
}

// Marks handlers whose route publishes an edge-cache policy; the middleware
// pays the store lookup only when a matched handler carries it.
const cacheableSymbol = Symbol('tempoApi.edgeCacheable')
const eligibility = new WeakMap<object, (c: Context) => boolean>()
const keys = new WeakMap<object, (c: Context) => Promise<string> | string>()

/**
 * Marks a route handler/middleware as publishing an edge-cache policy, so
 * {@link middleware} looks its route up in the edge cache. `Cache.response`
 * applies it; routes composed any other way never enter the shared cache.
 */
export function markCacheable<const handler extends (...args: never[]) => unknown>(
  handler: handler,
  options: markCacheable.Options = {},
): handler {
  if (options.key) keys.set(handler, options.key)
  return Object.assign(handler, { [cacheableSymbol]: true })
}

export declare namespace markCacheable {
  /** Options for marking a route as edge-cacheable. */
  type Options = {
    /** Cache key generator. Defaults to the request URL. */
    key?: ((c: Context) => Promise<string> | string) | undefined
  }
}

/** Adds a request eligibility check; every check must permit caching. */
export function setEligibility<const handler extends (...args: never[]) => unknown>(
  handler: handler,
  predicate: (c: Context) => boolean,
): handler {
  const existing = eligibility.get(handler)
  if (cacheableSymbol in handler)
    eligibility.set(handler, existing ? (c) => existing(c) && predicate(c) : predicate)
  return handler
}

// True when a marked handler matches and every marked handler permits caching.
function isCacheableRoute(c: Context): boolean {
  let cacheable = false
  for (const route of matchedRoutes(c)) {
    if (!(cacheableSymbol in route.handler)) continue
    cacheable = true
    if (eligibility.get(route.handler)?.(c) === false) return false
  }
  return cacheable
}

async function routeCacheKey(c: Context) {
  let requestUrl = c.req.url
  for (const route of matchedRoutes(c)) {
    if (!(cacheableSymbol in route.handler)) continue
    const key = keys.get(route.handler)
    if (key) requestUrl = await key(c)
  }
  return cacheKey(requestUrl)
}

export declare namespace middleware {
  /** Options for the edge-cache middleware. */
  type Options = {
    /** Store backing the edge cache. The layer is a no-op when omitted. */
    store?: Store.Store | undefined
  }
}

/** A response serialized for storage in the string-keyed edge cache. */
type Stored = {
  /** Response body. UTF-8 text, or base64 when `encoding` is `'base64'`. */
  body: string
  /** Set to `'base64'` when `body` holds base64-encoded binary bytes. */
  encoding?: 'base64' | undefined
  /** Response headers, already sanitized for shared caching. */
  headers: [string, string][]
  /** Response status code. */
  status: number
}

// Consumes `response` (callers pass a dedicated clone; the client streams the
// original).
async function serialize(response: Response, policy: Policy): Promise<string> {
  const headers = new Headers(response.headers)
  for (const name of volatileHeaders) headers.delete(name)
  headers.set('Cache-Control', cacheControl(policy))
  headers.set('Vary', 'Accept-Encoding')
  // The store is text-only, so binary bodies (e.g. PNG token logos) must be
  // base64-encoded or the round trip through `.text()` corrupts them. Text
  // bodies (JSON, SVG, HTML) stay verbatim to avoid the 33% base64 overhead.
  const textual = isTextual(headers.get('content-type'))
  const stored: Stored = {
    body: textual
      ? await response.text()
      : Base64.fromBytes(new Uint8Array(await response.arrayBuffer())),
    ...(textual ? {} : { encoding: 'base64' as const }),
    headers: [...headers],
    status: response.status,
  }
  return JSON.stringify(stored)
}

function serveHit(ifNoneMatch: string | undefined, value: string): Response {
  const { body, encoding, headers, status } = JSON.parse(value) as Stored
  const responseHeaders = new Headers(headers)
  const etag = responseHeaders.get('etag')
  if (Etag.matches(etag, ifNoneMatch)) {
    const notModified = new Headers()
    for (const name of ['Cache-Control', 'ETag', 'Vary']) {
      const header = responseHeaders.get(name)
      if (header) notModified.set(name, header)
    }
    return new Response(null, { headers: notModified, status: 304, statusText: 'Not Modified' })
  }
  if (encoding === 'base64') {
    // Copy the decoded bytes into an exact `ArrayBuffer`: `Base64.toBytes`
    // returns a `Uint8Array` view, which this TS lib does not accept as a
    // `BodyInit`.
    const bytes = Base64.toBytes(body)
    const buffer = new ArrayBuffer(bytes.byteLength)
    new Uint8Array(buffer).set(bytes)
    return new Response(buffer, { headers: responseHeaders, status })
  }
  return new Response(body, { headers: responseHeaders, status })
}

// A body is safe to store as verbatim text when its content type is textual;
// everything else (images, octet-stream) is treated as binary and base64-coded.
function isTextual(contentType: string | null): boolean {
  if (!contentType) return false
  return /^(?:text\/|application\/(?:json|xml|javascript|[\w.-]+\+(?:json|xml))|image\/svg\+xml)/i.test(
    contentType,
  )
}

// Sort query params so semantically identical URLs share one cache entry.
function cacheKey(requestUrl: string) {
  const url = new URL(requestUrl)
  url.searchParams.sort()
  // Version the key so entries written before a sanitizer change cannot
  // survive the deployment that changes the shared-cache boundary.
  return `edge:v2:${url.toString()}`
}

function getWaitUntil(c: { executionCtx?: { waitUntil(promise: Promise<unknown>): void } }) {
  // `c.executionCtx` throws on runtimes without one (e.g. Node), so probe it
  // defensively and fall back to awaiting the write inline.
  try {
    const ctx = c.executionCtx
    return ctx ? ctx.waitUntil.bind(ctx) : undefined
  } catch {
    return undefined
  }
}

namespace Etag {
  export function matches(etag: string | null, ifNoneMatch: string | undefined) {
    if (!etag || !ifNoneMatch) return false
    const strip = (tag: string) => tag.trim().replace(/^W\//, '')
    return ifNoneMatch.split(',').some((tag) => strip(tag) === strip(etag))
  }
}

/**
 * Builds the shared-cache `Cache-Control` value for a stored copy: `public`
 * with `max-age`/`s-maxage` bound to the policy plus an optional
 * `stale-while-revalidate` window. Used for this layer's stored copies and by
 * `Cache.response` to bound the inner response cache's entry lifetime.
 */
export function cacheControl(policy: Policy) {
  const directives = ['public', `max-age=${policy.maxAge}`, `s-maxage=${policy.maxAge}`]
  if (policy.staleWhileRevalidate)
    directives.push(`stale-while-revalidate=${policy.staleWhileRevalidate}`)
  return directives.join(', ')
}