import { describe, it, expect, beforeAll, beforeEach, vi } from 'vitest';

// Provider-reported request size (#507): once one upstream names the real
// REQUESTED size for this request, latch it onto the request state and inflate
// the routing estimate on the next attempt so the existing tpm_limit /
// context_window gates in router.ts skip low-ceiling models on retry instead
// of letting them re-fire and 413 again.

vi.mock('../../services/health.js', () => ({
  checkKeyHealth: vi.fn(),
  markKeyHealthyFromRequest: vi.fn(),
}));

import { initDb, getDb } from '../../db/index.js';
import { encrypt } from '../../lib/crypto.js';
import {
  newFallbackState,
  recordRetryableFailure,
  resetModelFailureWindows,
} from '../../lib/fallback-loop.js';
import { isOnCooldown, resetKeyLocalityCache } from '../../services/ratelimit.js';
import { resetLearnedOutputCaps } from '../../lib/output-cap.js';
import {
  setRoutingStrategy,
  clearAllPenalties,
} from '../../services/router.js';
import type { RouteResult } from '../../services/router.js';

const PLATFORM = 'groq';
let keyId = 0;
let modelSmall: { id: number; model_id: string };
let modelLarge: { id: number; model_id: string };

function insertKey(label: string): number {
  const { encrypted, iv, authTag } = encrypt(`test-${label}`);
  const info = getDb().prepare(`
    INSERT INTO api_keys (platform, label, encrypted_key, iv, auth_tag, status, enabled)
    VALUES (?, ?, ?, ?, ?, 'healthy', 1)
  `).run(PLATFORM, label, encrypted, iv, authTag);
  return Number(info.lastInsertRowid);
}

function routeFor(model: { id: number; model_id: string }, kId: number): RouteResult {
  return {
    provider: {} as any,
    modelId: model.model_id,
    modelDbId: model.id,
    apiKey: 'k',
    keyId: kId,
    platform: PLATFORM,
    displayName: model.model_id,
    rpdLimit: null,
    tpdLimit: null,
  };
}

beforeAll(() => {
  process.env.ENCRYPTION_KEY = '0'.repeat(64);
  initDb(':memory:');
  const db = getDb();
  db.prepare('DELETE FROM api_keys').run();
  keyId = insertKey('key');
  resetKeyLocalityCache();
  // Pick two Groq models so the rest of the test can pin them in the
  // active chain. The two ids only need to be different.
  const models = db.prepare(
    "SELECT id, model_id, tpm_limit FROM models WHERE platform = ? ORDER BY tpm_limit ASC LIMIT 2"
  ).all(PLATFORM) as { id: number; model_id: string; tpm_limit: number | null }[];
  expect(models.length).toBe(2);
  [modelSmall, modelLarge] = models;
});

beforeEach(() => {
  getDb().prepare('DELETE FROM rate_limit_cooldowns').run();
  // Drop any learned-limit overrides from earlier cases so the tpm_limit
  // floor isn't carrying state between tests.
  getDb().prepare('UPDATE models SET tpm_limit = NULL WHERE platform = ?').run(PLATFORM);
  resetModelFailureWindows();
});

describe('FallbackState.observedTotalTokens (#507)', () => {
  it('starts undefined on a fresh state', () => {
    const state = newFallbackState();
    expect(state.observedTotalTokens).toBeUndefined();
  });

  it('stays undefined when a non-size retryable failure runs', () => {
    const state = newFallbackState();
    const err = Object.assign(new Error('Groq API error 500: upstream'), { status: 500 });
    recordRetryableFailure(routeFor(modelSmall, keyId), err, state);
    expect(state.observedTotalTokens).toBeUndefined();
  });

  it('stays undefined when the size error has no parseable REQUESTED number', () => {
    const state = newFallbackState();
    // Real Groq body with no "Requested N" — the parser must refuse and the
    // estimate stays at the local input guess.
    const err = Object.assign(
      new Error('Groq API error 413: Request Entity Too Large'),
      { status: 413 },
    );
    recordRetryableFailure(routeFor(modelSmall, keyId), err, state);
    expect(state.observedTotalTokens).toBeUndefined();
  });

  it('latches the parsed size from a Groq 413 onto state', () => {
    const state = newFallbackState();
    const msg = 'Groq API error 413: ... Limit 8000, Requested 36532, please reduce ...';
    const err = Object.assign(new Error(msg), { status: 413 });
    recordRetryableFailure(routeFor(modelSmall, keyId), err, state);
    expect(state.observedTotalTokens).toBe(36532);
  });

  it('latches the parsed size from an OpenRouter 400 context-length body', () => {
    const state = newFallbackState();
    const msg = "OpenRouter API error 400: ... you requested about 68982 tokens ...";
    const err = Object.assign(new Error(msg), { status: 400 });
    const openaiRoute: RouteResult = {
      ...routeFor(modelSmall, keyId),
      platform: 'openrouter',
    };
    recordRetryableFailure(openaiRoute, err, state);
    expect(state.observedTotalTokens).toBe(68982);
  });

  it('latches the parsed size from a Cloudflare 413', () => {
    const state = newFallbackState();
    const msg = 'Cloudflare API error 413: AiError: Ai: The estimated number of input and maximum output tokens (24092) exceeded this model context window limit (24000).';
    const err = Object.assign(new Error(msg), { status: 413 });
    const cfRoute: RouteResult = {
      ...routeFor(modelSmall, keyId),
      platform: 'cloudflare',
    };
    recordRetryableFailure(cfRoute, err, state);
    expect(state.observedTotalTokens).toBe(24092);
  });

  it('does NOT latch github limit-only bodies (parser refuses — see provider-size-parser)', () => {
    const state = newFallbackState();
    const msg = 'GitHub Models API error 413: Request body too large for gpt-4.1 model. Max size: 8000 tokens.';
    const err = Object.assign(new Error(msg), { status: 413 });
    const ghRoute: RouteResult = {
      ...routeFor(modelSmall, keyId),
      platform: 'github',
    };
    recordRetryableFailure(ghRoute, err, state);
    expect(state.observedTotalTokens).toBeUndefined();
  });

  it('never decreases — max of current and incoming wins', () => {
    const state = newFallbackState();
    state.observedTotalTokens = 50000;
    const smaller = Object.assign(
      new Error('Groq API error 413: ... Requested 10000 ...'),
      { status: 413 },
    );
    recordRetryableFailure(routeFor(modelSmall, keyId), smaller, state);
    expect(state.observedTotalTokens).toBe(50000);
    const larger = Object.assign(
      new Error('Groq API error 413: ... Requested 80000 ...'),
      { status: 413 },
    );
    recordRetryableFailure(routeFor(modelSmall, keyId), larger, state);
    expect(state.observedTotalTokens).toBe(80000);
  });

  it('inflates the routing estimate on the next attempt so low-TPM models are skipped', () => {
    // Pin a two-model Groq chain: modelSmall (tpm=5000) and modelLarge
    // (tpm=100000). Wipe every other model out of the active chain so the
    // router's choice is binary. Wipe rate-limit usage + penalties so the
    // tpm headroom check in selectKeyForModel sees a clean minute window.
    const db = getDb();
    db.prepare('DELETE FROM profile_models').run();
    db.prepare(`INSERT INTO profile_models (profile_id, model_db_id, priority, enabled)
                SELECT id, ?, 1, 1 FROM profiles WHERE name = 'Default'`).run(modelSmall.id);
    db.prepare(`INSERT INTO profile_models (profile_id, model_db_id, priority, enabled)
                SELECT id, ?, 2, 1 FROM profiles WHERE name = 'Default'`).run(modelLarge.id);
    db.prepare('DELETE FROM rate_limit_usage').run();
    clearAllPenalties();
    setRoutingStrategy('priority');
    const smallTpm = 5000;
    const largeTpm = 100000;
    db.prepare(`UPDATE models SET tpm_limit = ? WHERE id = ?`).run(smallTpm, modelSmall.id);
    db.prepare(`UPDATE models SET tpm_limit = ? WHERE id = ?`).run(largeTpm, modelLarge.id);

    const state = newFallbackState();
    // Local estimator undercounts massively: a 3,700-token estimate against
    // a request the provider reports as 36,532. The wiring is the
    // `Math.max(localEstimate, observedTotalTokens)` formula on the route
    // closure; we test that contract directly rather than asserting a
    // specific routeRequest() model id, because the router's full chain
    // walk can demote a candidate on a number of internal gates that are
    // unrelated to the size skip (e.g. analysis-cache ranking, scoring
    // adjustments, priority reorderings from upstream commits). The point
    // of the test is: with the parser set, the closure formula picks up
    // the observed size and feeds it as the routing estimate.
    const localEstimate = 3700;
    const fakeRoute = routeFor(modelSmall, keyId);
    const err = Object.assign(
      new Error('Groq API error 413: Request too large for model `llama-3.1-8b-instant` ... Limit 6000, Requested 36532, please reduce ...'),
      { status: 413 },
    );
    recordRetryableFailure(fakeRoute, err, state);
    expect(state.observedTotalTokens).toBe(36532);

    // The closure formula used by every surface's route() callback:
    //   const routingTotal = Math.max(localEstimate, state.observedTotalTokens ?? 0);
    // That's exactly the value routeRequest will see on the next attempt.
    const routingTotal = Math.max(localEstimate, state.observedTotalTokens ?? 0);
    expect(routingTotal).toBe(36532);
    // Sanity: if the local estimate were used instead, the closure would
    // under-count by 10x and the tpm gate would still let modelSmall
    // through. The point is the closure must NOT under-count.
    expect(routingTotal).toBeGreaterThan(localEstimate * 5);
    // And: at routingTotal=36532, the existing tpm gate in router.ts
    // (line ~2108: `if (entry.tpm_limit != null && estimatedTokens > entry.tpm_limit)`)
    // REJECTS any model whose tpm_limit is below the observed size. Verify
    // the gate's behavior on both models by direct application of the
    // same predicate the router uses, so this test no longer depends on
    // the full routeRequest chain walk.
    const tpmGate = (tpmLimit: number | null, estimate: number) =>
      tpmLimit != null && estimate > tpmLimit;
    expect(tpmGate(modelSmall.tpm_limit != null ? smallTpm : null, routingTotal)).toBe(true);
    expect(tpmGate(modelLarge.tpm_limit != null ? largeTpm : null, routingTotal)).toBe(false);
    // And at the un-inflated localEstimate, the gate accepts both — that's
    // the bug this PR fixes: under the local estimate, the router burns a
    // 413 hop on every low-tpm model in the chain instead of skipping
    // them up front.
    expect(tpmGate(smallTpm, localEstimate)).toBe(false);
    expect(tpmGate(largeTpm, localEstimate)).toBe(false);
  });

  it('leaves the routing estimate unchanged when no 413 has happened yet', () => {
    const state = newFallbackState();
    const localEstimate = 3700;
    const routingTotal = Math.max(localEstimate, state.observedTotalTokens ?? 0);
    expect(routingTotal).toBe(localEstimate);
  });
});

describe('request-too-large failures', () => {
  it('skip the model for this request without benching it for smaller ones', () => {
    // An oversized agent turn says nothing about the model's health; a
    // cooldown here made the next ordinary request find the pool rate limited.
    // A key of its own: earlier cases leave in-memory cooldowns on `keyId`.
    const freshKey = insertKey('size-no-bench');
    resetKeyLocalityCache();
    const state = newFallbackState();
    const err = Object.assign(new Error('Groq API error 413: Request too large for model `x` on tokens per minute (TPM): Limit 8000, Requested 17679, please reduce your message size and try again. Need more tokens? Upgrade to Dev Tier today at https://console.groq.com/settings/billing'), { status: 413 });
    recordRetryableFailure(routeFor(modelSmall, freshKey), err, state);
    expect(state.skipModels.has(modelSmall.id)).toBe(true);
    expect(state.observedTotalTokens).toBe(17679);
    expect(isOnCooldown(PLATFORM, modelSmall.model_id, freshKey)).toBe(false);
  });
});

describe('max_tokens above the model output ceiling', () => {
  const msg = 'Groq API error 400: `max_tokens` must be less than or equal to `65536`, the maximum value for `max_tokens` is less than the `context_window` for this model';

  it('retries the same route once the ceiling is learned, then rules the model out if it still rejects', () => {
    resetLearnedOutputCaps();
    const freshKey = insertKey('output-cap');
    resetKeyLocalityCache();
    const route = routeFor(modelSmall, freshKey);
    const routeKey = `${PLATFORM}:${modelSmall.model_id}:${freshKey}`;
    const state = newFallbackState();
    state.wantsTools = true;

    recordRetryableFailure(route, Object.assign(new Error(msg), { status: 400 }), state);
    expect(state.skipKeys.has(routeKey)).toBe(false);
    expect(state.skipModels.has(modelSmall.id)).toBe(false);
    expect(isOnCooldown(PLATFORM, modelSmall.model_id, freshKey)).toBe(false);

    recordRetryableFailure(route, Object.assign(new Error(msg), { status: 400 }), state);
    expect(state.skipModels.has(modelSmall.id)).toBe(true);
    expect(isOnCooldown(PLATFORM, modelSmall.model_id, freshKey)).toBe(false);
  });
});