import { NextResponse } from "next/server";
import type { NextRequest } from "next/server";

/**
 * Rate limiting for the public API surface.
 *
 * `/api/chat` is reachable by anyone on the internet, takes no credentials,
 * and spends AWS Bedrock credit on every call (understanding + generation +
 * entailment, several model calls per turn) against a small personal budget.
 * The architecture assigns rate limiting to `api-gateway`, which is not built
 * yet, so until it exists this is the only place the limit can live.
 *
 * It runs here rather than inside the route handlers because the edge
 * middleware rejects an abusive request before any handler work — and because
 * it covers every current and future `/api/*` route by default rather than
 * one endpoint at a time.
 *
 * Known limitations, stated rather than hidden:
 * - Counters are per process and in memory. Production runs a single `web`
 *   container, so this is a real limit today; the moment that scales to more
 *   than one replica each replica gets its own allowance, and this needs to
 *   move to Redis (already reserved for exactly this in `docs/architecture.md`)
 *   or to the gateway.
 * - It is keyed by client IP, so it throttles a shared NAT as one caller. That
 *   is the correct trade for a cost guard with no authentication; per-user
 *   limits need auth, which does not exist yet.
 * - It is a cost and abuse guard, not a security control. It does not
 *   authenticate anyone and must not be described as if it does.
 */

interface Bucket {
  hits: number[];
}

interface Rule {
  windowMs: number;
  max: number;
}

// Chat is the expensive path: several Bedrock calls per request, and a single
// turn was measured taking up to ~45s of model time. Autocomplete is a local
// catalog lookup with no model call, so it can be far more generous without
// costing anything.
const RULES: Array<{ prefix: string; rules: Rule[] }> = [
  {
    prefix: "/api/chat",
    rules: [
      { windowMs: 60_000, max: 12 },
      { windowMs: 3_600_000, max: 120 },
    ],
  },
  {
    prefix: "/api/suggest",
    rules: [{ windowMs: 60_000, max: 120 }],
  },
  // `/api/pdf` streams the whole 37MB source PDF from disk on every request,
  // with no range support and no caching headers. It is unauthenticated, so
  // repeated fetches are a bandwidth and memory cost on a single small EC2
  // host. Sized against real UI behaviour rather than guessed: `tra-cuu`
  // uses it as an iframe `src` whose `#page=` fragment changes per citation
  // click, and `CitationCard` links to it, so a clinician working through a
  // long evidence list can legitimately fetch it repeatedly. 30/min is far
  // above that and still bounds an automated puller.
  {
    prefix: "/api/pdf",
    rules: [{ windowMs: 60_000, max: 30 }],
  },
  // `/api/feedback` writes one row per answer (upsert keyed on trace_id), so
  // the ceiling only needs to exceed how fast a human can rate answers. It
  // reaches PostgreSQL on every call, which is why it is bounded at all.
  {
    prefix: "/api/feedback",
    rules: [{ windowMs: 60_000, max: 60 }],
  },
];

const buckets = new Map<string, Bucket>();
const LONGEST_WINDOW_MS = 3_600_000;
let lastSweep = 0;

/** Drop entries no rule can still be counting, so the map cannot grow without bound. */
function sweep(now: number) {
  if (now - lastSweep < 60_000) return;
  lastSweep = now;
  for (const [key, bucket] of buckets) {
    const live = bucket.hits.filter((t) => now - t < LONGEST_WINDOW_MS);
    if (live.length === 0) buckets.delete(key);
    else bucket.hits = live;
  }
}

/**
 * Caddy sits in front and sets `X-Forwarded-For`; the left-most entry is the
 * original client. Falling back to a shared key rather than to "unlimited"
 * matters: an unknown IP must not become a way to opt out of the limit.
 */
function clientKey(request: NextRequest): string {
  const forwarded = request.headers.get("x-forwarded-for");
  if (forwarded) {
    const first = forwarded.split(",")[0]?.trim();
    if (first) return first;
  }
  return request.headers.get("x-real-ip")?.trim() || "unknown";
}

function matchRules(pathname: string) {
  return RULES.find((entry) => pathname.startsWith(entry.prefix))?.rules;
}

export function middleware(request: NextRequest) {
  const rules = matchRules(request.nextUrl.pathname);
  if (!rules) return NextResponse.next();

  const now = Date.now();
  sweep(now);

  const key = `${clientKey(request)}:${request.nextUrl.pathname}`;
  const bucket = buckets.get(key) ?? { hits: [] };
  bucket.hits = bucket.hits.filter((t) => now - t < LONGEST_WINDOW_MS);

  for (const rule of rules) {
    const inWindow = bucket.hits.filter((t) => now - t < rule.windowMs);
    if (inWindow.length >= rule.max) {
      const oldest = Math.min(...inWindow);
      const retryAfterSec = Math.max(1, Math.ceil((rule.windowMs - (now - oldest)) / 1000));
      // Record nothing for a rejected request: a client hammering the endpoint
      // should not keep pushing its own window forward and lock itself out for
      // longer than the rule says.
      buckets.set(key, bucket);
      return NextResponse.json(
        {
          error: "rate_limited",
          message:
            "Bạn đang gửi quá nhiều yêu cầu trong thời gian ngắn. Vui lòng đợi một lát rồi thử lại.",
        },
        {
          status: 429,
          headers: {
            "Retry-After": String(retryAfterSec),
            "X-RateLimit-Limit": String(rule.max),
            "X-RateLimit-Remaining": "0",
          },
        }
      );
    }
  }

  bucket.hits.push(now);
  buckets.set(key, bucket);

  const tightest = rules[0];
  const used = bucket.hits.filter((t) => now - t < tightest.windowMs).length;
  const response = NextResponse.next();
  response.headers.set("X-RateLimit-Limit", String(tightest.max));
  response.headers.set("X-RateLimit-Remaining", String(Math.max(0, tightest.max - used)));
  return response;
}

export const config = {
  matcher: ["/api/:path*"],
};