import { createHash } from "node:crypto";
import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { dirname } from "node:path";
import type { Pool, ResolvedModel } from "./router-core.ts";

export interface RateLimitSnapshot {
  remaining?: number;
  limit?: number;
  resetAt?: number;
}

export type PlanStatus = "ok" | "cooldown";

export interface PlanState {
  planKey: string;
  status: PlanStatus;
  cooldownUntil?: number;
  reason?: string;
  lastSnapshot?: RateLimitSnapshot;
  updatedAt: number;
}

export interface QuotaConfig {
  enabled: boolean;
  reserveRatio: number;
  inTurnRetry: boolean;
  maxRetries: number;
  defaultCooldownMs: number;
}

export interface PlanIdentity {
  provider: string;
  baseUrl: string;
  apiKey?: string;
  headers?: Record<string, string>;
  env?: Record<string, string>;
}

export const DEFAULT_QUOTA_CONFIG: QuotaConfig = {
  enabled: true,
  reserveRatio: 0.05,
  inTurnRetry: false,
  maxRetries: 2,
  defaultCooldownMs: 300_000,
};

const QUOTA_STATE_VERSION = 2;

interface PersistedQuotaState {
  version: typeof QUOTA_STATE_VERSION;
  plans: PlanState[];
}

export class QuotaState {
  readonly config: QuotaConfig;
  private readonly plans = new Map<string, PlanState>();

  constructor(config: QuotaConfig | Partial<QuotaConfig> = DEFAULT_QUOTA_CONFIG) {
    this.config = { ...DEFAULT_QUOTA_CONFIG, ...config };
  }

  isAvailable(planKey: string, now: number): boolean {
    const state = this.plans.get(planKey);
    if (!state || state.status === "ok") return true;

    if (state.cooldownUntil != null && now >= state.cooldownUntil) {
      state.status = "ok";
      state.cooldownUntil = undefined;
      state.reason = undefined;
      state.updatedAt = now;
      return true;
    }

    return false;
  }

  recordResponse(
    planKey: string,
    status: number,
    headers: Record<string, string>,
    provider: string,
    now: number,
  ): PlanState {
    const snapshot = parseRateLimitHeaders(provider, headers, now);
    const state = this.get(planKey, now);
    state.lastSnapshot = isEmptySnapshot(snapshot) ? undefined : snapshot;
    state.updatedAt = now;

    if (status === 429) {
      return this.recordRateLimited(planKey, parseRetryAfter(headers["retry-after"], now), snapshot.resetAt, now);
    }

    if (snapshot.remaining != null && snapshot.limit != null && snapshot.limit > 0) {
      if (snapshot.remaining / snapshot.limit <= this.config.reserveRatio) {
        state.status = "cooldown";
        state.cooldownUntil = snapshot.resetAt ?? now + this.config.defaultCooldownMs;
        state.reason = "low-remaining";
        return state;
      }
    }

    state.status = "ok";
    state.cooldownUntil = undefined;
    state.reason = undefined;
    return state;
  }

  recordRateLimited(planKey: string, retryAfterMs: number | undefined, resetAt: number | undefined, now: number): PlanState {
    const state = this.get(planKey, now);
    state.status = "cooldown";
    state.cooldownUntil = resetAt ?? (retryAfterMs != null ? now + retryAfterMs : now + this.config.defaultCooldownMs);
    state.reason = "429";
    state.updatedAt = now;
    return state;
  }

  snapshot(planKey: string): PlanState | undefined {
    const state = this.plans.get(planKey);
    return state ? { ...state, lastSnapshot: state.lastSnapshot ? { ...state.lastSnapshot } : undefined } : undefined;
  }

  snapshots(): PlanState[] {
    return [...this.plans.values()].map((state) => ({
      ...state,
      lastSnapshot: state.lastSnapshot ? { ...state.lastSnapshot } : undefined,
    }));
  }

  /** Returns true when no matching plan is known or at least one matching auth plan is usable. */
  isAnyAvailableMatching(planKeyPrefix: string, now: number): boolean {
    const matching = [...this.plans.keys()].filter((planKey) => planKey.startsWith(planKeyPrefix));
    return matching.length === 0 || matching.some((planKey) => this.isAvailable(planKey, now));
  }

  load(file: string): void {
    if (!existsSync(file)) return;

    try {
      const parsed = JSON.parse(readFileSync(file, "utf8")) as Partial<PersistedQuotaState>;
      if (parsed.version !== QUOTA_STATE_VERSION || !Array.isArray(parsed.plans)) return;

      this.plans.clear();
      for (const plan of parsed.plans) {
        if (!isValidPlanState(plan)) continue;
        this.plans.set(plan.planKey, { ...plan, lastSnapshot: plan.lastSnapshot ? { ...plan.lastSnapshot } : undefined });
      }
    } catch {
      return;
    }
  }

  persist(file: string): void {
    mkdirSync(dirname(file), { recursive: true });
    const data: PersistedQuotaState = { version: QUOTA_STATE_VERSION, plans: this.snapshots() };
    writeFileSync(file, `${JSON.stringify(data, null, 2)}\n`, "utf8");
  }

  private get(planKey: string, now: number): PlanState {
    const existing = this.plans.get(planKey);
    if (existing) return existing;

    const state: PlanState = { planKey, status: "ok", updatedAt: now };
    this.plans.set(planKey, state);
    return state;
  }
}

export function buildPlanKey(identity: PlanIdentity): string {
  const provider = identity.provider.trim().toLowerCase();
  const baseUrl = normalizeBaseUrl(identity.baseUrl);
  const authHash = hashStable({
    apiKey: identity.apiKey ?? null,
    headers: normalizeRecord(identity.headers),
    env: normalizeRecord(identity.env),
  });
  return `${provider}|${baseUrl}|auth:${authHash}`;
}

export function filterPoolByQuota(
  pool: Pool,
  quota: QuotaState | undefined,
  now: number,
  excluded = new Set<string>(),
  planKeyFor: (item: ResolvedModel) => string = (item) =>
    buildPlanKey({ provider: item.model.provider, baseUrl: item.model.baseUrl }),
): Pool {
  if (!quota?.config.enabled) return pool;

  const keep = (item: ResolvedModel) => {
    const planKey = planKeyFor(item);
    return !excluded.has(planKey) && quota.isAvailable(planKey, now);
  };
  const filter = (items: ResolvedModel[]) => items.filter(keep);
  const next: Pool = {
    cheapPool: filter(pool.cheapPool),
    standardPool: filter(pool.standardPool),
    strongPool: filter(pool.strongPool),
    unknownPool: filter(pool.unknownPool),
    all: filter(pool.all),
  };

  return next.all.length === 0 ? pool : next;
}

/** Filters models when every persisted auth plan for the same provider endpoint is cooling down. */
export function filterPoolByQuotaPlanPrefix(pool: Pool, quota: QuotaState | undefined, now: number): Pool {
  if (!quota?.config.enabled) return pool;

  const keep = (item: ResolvedModel) => quota.isAnyAvailableMatching(
    quotaPlanPrefix(item.model.provider, item.model.baseUrl),
    now,
  );
  const filter = (items: ResolvedModel[]) => items.filter(keep);
  const next: Pool = {
    cheapPool: filter(pool.cheapPool),
    standardPool: filter(pool.standardPool),
    strongPool: filter(pool.strongPool),
    unknownPool: filter(pool.unknownPool),
    all: filter(pool.all),
  };
  return next.all.length === 0 ? pool : next;
}

export function quotaPlanPrefix(provider: string, baseUrl: string): string {
  return `${provider.trim().toLowerCase()}|${normalizeBaseUrl(baseUrl)}|auth:`;
}

export function parseRateLimitHeaders(
  _provider: string,
  headers: Record<string, string>,
  now: number,
): RateLimitSnapshot {
  const h = normalizeHeaders(headers);
  const snapshot: RateLimitSnapshot = {};

  if (h["anthropic-ratelimit-tokens-remaining"] != null) {
    snapshot.remaining = num(h["anthropic-ratelimit-tokens-remaining"]);
    snapshot.limit = num(h["anthropic-ratelimit-tokens-limit"]);
    snapshot.resetAt = parseRfc3339(h["anthropic-ratelimit-tokens-reset"]);
  } else if (h["x-ratelimit-remaining-tokens"] != null) {
    snapshot.remaining = num(h["x-ratelimit-remaining-tokens"]);
    snapshot.limit = num(h["x-ratelimit-limit-tokens"]);
    const duration = parseDuration(h["x-ratelimit-reset-tokens"]);
    if (duration != null) snapshot.resetAt = now + duration;
  }

  const retryAfter = parseRetryAfter(h["retry-after"], now);
  if (retryAfter != null && snapshot.resetAt == null) snapshot.resetAt = now + retryAfter;

  return pruneSnapshot(snapshot);
}

export function parseRfc3339(value: string | undefined): number | undefined {
  if (value == null) return undefined;
  const parsed = Date.parse(value);
  return Number.isFinite(parsed) ? parsed : undefined;
}

export function parseDuration(value: string | undefined): number | undefined {
  if (value == null) return undefined;
  const trimmed = value.trim().toLowerCase();
  if (!trimmed) return undefined;
  if (/^\d+(?:\.\d+)?$/.test(trimmed)) return Number(trimmed) * 1_000;

  let total = 0;
  let matched = false;
  const pattern = /(\d+(?:\.\d+)?)(ms|h|m|s)/g;
  for (const match of trimmed.matchAll(pattern)) {
    matched = true;
    const value = Number(match[1]);
    if (!Number.isFinite(value)) return undefined;
    if (match[2] === "ms") total += value;
    if (match[2] === "s") total += value * 1_000;
    if (match[2] === "m") total += value * 60_000;
    if (match[2] === "h") total += value * 3_600_000;
  }

  return matched ? total : undefined;
}

export function parseRetryAfter(value: string | undefined, now: number): number | undefined {
  if (value == null) return undefined;
  const trimmed = value.trim();
  if (!trimmed) return undefined;
  if (/^\d+(?:\.\d+)?$/.test(trimmed)) return Number(trimmed) * 1_000;

  const parsed = Date.parse(trimmed);
  if (!Number.isFinite(parsed)) return undefined;
  return Math.max(0, parsed - now);
}

function normalizeHeaders(headers: Record<string, string>): Record<string, string> {
  return Object.fromEntries(Object.entries(headers).map(([key, value]) => [key.toLowerCase(), value]));
}

function normalizeBaseUrl(baseUrl: string): string {
  return baseUrl.trim().replace(/\/+$/, "");
}

function normalizeRecord(record: Record<string, string> | undefined): Record<string, string> | null {
  if (!record || Object.keys(record).length === 0) return null;
  return Object.fromEntries(
    Object.entries(record)
      .map(([key, value]) => [key.toLowerCase(), value] as const)
      .sort(([a], [b]) => a.localeCompare(b)),
  );
}

function hashStable(value: unknown): string {
  return createHash("sha256").update(JSON.stringify(value)).digest("hex").slice(0, 16);
}

function num(value: string | undefined): number | undefined {
  if (value == null) return undefined;
  const parsed = Number(value);
  return Number.isFinite(parsed) ? parsed : undefined;
}

function pruneSnapshot(snapshot: RateLimitSnapshot): RateLimitSnapshot {
  return Object.fromEntries(Object.entries(snapshot).filter(([, value]) => value != null)) as RateLimitSnapshot;
}

function isEmptySnapshot(snapshot: RateLimitSnapshot): boolean {
  return snapshot.remaining == null && snapshot.limit == null && snapshot.resetAt == null;
}

function isValidPlanState(value: unknown): value is PlanState {
  if (!value || typeof value !== "object") return false;
  const plan = value as Partial<PlanState>;
  return (
    typeof plan.planKey === "string" &&
    (plan.status === "ok" || plan.status === "cooldown") &&
    typeof plan.updatedAt === "number"
  );
}