import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
import { mkdtempSync, rmSync } from 'fs';
import { tmpdir } from 'os';
import { join } from 'path';
import { Repository } from './db/repository.js';
import { maybeEnqueueReflection } from './worker.js';

describe('maybeEnqueueReflection', () => {
  let dir: string;
  let repo: Repository;

  beforeEach(() => {
    dir = mkdtempSync(join(tmpdir(), 'reflect-enq-'));
    repo = new Repository(join(dir, 'db.sqlite'));
  });
  afterEach(() => {
    repo.close?.();
    rmSync(dir, { recursive: true, force: true });
    vi.useRealTimers();
  });

  it('creates a reflection job after an agent job completes', async () => {
    const job = await repo.createJob({
      repo: 'local/task-1', issueNumber: 1,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    // Simulate the job having completed (moved out of queued) before reflection enqueue
    await repo.updateJob(job.id, { status: 'succeeded' });
    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: 0,
    });
    const queued = await repo.getJobsByStatus('queued');
    expect(queued).toHaveLength(1);
    expect(queued[0].taskKind).toBe('reflection');
    expect(queued[0].pieceName).toBe('reflection');
    const payload = JSON.parse(queued[0].payload!);
    expect(payload.originalJobId).toBe(job.id);
    expect(payload.userId).toBe('u-1');
    expect(payload.outcome).toBe('succeeded');
  });

  it('does NOT recurse on reflection jobs', async () => {
    const job = await repo.createJob({
      repo: 'local/reflection-x', issueNumber: 0,
      instruction: '', pieceName: 'reflection',
      role: 'reflection', taskKind: 'reflection',
      payload: '{}', ownerId: 'u-1',
    } as any);
    // Simulate the reflection job having completed
    await repo.updateJob(job.id, { status: 'succeeded' });
    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: 0,
    });
    expect(await repo.getJobsByStatus('queued')).toHaveLength(0);
  });

  it('does nothing when enabled=false', async () => {
    const job = await repo.createJob({
      repo: 'local/task-1', issueNumber: 1,
      instruction: 'x', pieceName: 'chat', ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });
    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: false, workerRequired: false, perUserDailyBudgetTokens: 0,
    });
    expect(await repo.getJobsByStatus('queued')).toHaveLength(0);
  });

  it('does nothing when ownerId is null (no user to learn for)', async () => {
    const job = await repo.createJob({
      repo: 'local/task-1', issueNumber: 1,
      instruction: 'x', pieceName: 'chat',
    } as any);   // no ownerId
    await repo.updateJob(job.id, { status: 'succeeded' });
    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: 0,
    });
    expect(await repo.getJobsByStatus('queued')).toHaveLength(0);
  });

  // ── Budget enforcement tests ─────────────────────────────────────────────

  /**
   * Helper: insert a reflection_metrics row for userId with given token counts
   * and a created_at timestamp (epoch ms).
   */
  function seedMetric(
    userId: string,
    tokensIn: number,
    tokensOut: number,
    createdAtMs: number,
  ): void {
    // We insert via raw SQL so we can control the timestamp. Use the public
    // getDb() getter rather than reaching into a private field with `as any`.
    const db = repo.getDb();
    db.prepare(`
      INSERT INTO reflection_metrics
        (reflection_job_id, original_job_id, user_id, piece_name, outcome,
         memory_changes, piece_edited, tokens_in, tokens_out, duration_ms, created_at)
      VALUES
        (?, NULL, ?, 'chat', 'applied', 0, 0, ?, ?, 100, ?)
    `).run(`rjob-${Math.random()}`, userId, tokensIn, tokensOut, createdAtMs);
  }

  it('skips enqueue when user is at or above the daily budget cap', async () => {
    const job = await repo.createJob({
      repo: 'local/task-2', issueNumber: 2,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    // Seed today's usage equal to the cap (2 000 000 tokens total)
    const cap = 2_000_000;
    const todayStartMs = Date.UTC(
      new Date().getUTCFullYear(),
      new Date().getUTCMonth(),
      new Date().getUTCDate(),
    );
    seedMetric('u-1', 1_200_000, 800_000, todayStartMs + 1_000); // total = 2 000 000

    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: cap,
    });

    // No reflection job should have been created.
    expect(await repo.getJobsByStatus('queued')).toHaveLength(0);
  });

  it('enqueues when user is just under the daily budget cap', async () => {
    const job = await repo.createJob({
      repo: 'local/task-3', issueNumber: 3,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    const cap = 2_000_000;
    const todayStartMs = Date.UTC(
      new Date().getUTCFullYear(),
      new Date().getUTCMonth(),
      new Date().getUTCDate(),
    );
    // Seed 1 token below the cap
    seedMetric('u-1', 1_000_000, 999_999, todayStartMs + 1_000); // total = 1 999 999

    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: cap,
    });

    const queued = await repo.getJobsByStatus('queued');
    expect(queued).toHaveLength(1);
    expect(queued[0].taskKind).toBe('reflection');
  });

  it('cap=0 means "no limit" — enqueue happens even with massive token usage', async () => {
    const job = await repo.createJob({
      repo: 'local/task-4', issueNumber: 4,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    const todayStartMs = Date.UTC(
      new Date().getUTCFullYear(),
      new Date().getUTCMonth(),
      new Date().getUTCDate(),
    );
    // Seed an enormous amount of token usage — shouldn't matter when cap=0
    seedMetric('u-1', 50_000_000, 50_000_000, todayStartMs + 1_000);

    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: 0,
    });

    const queued = await repo.getJobsByStatus('queued');
    expect(queued).toHaveLength(1);
    expect(queued[0].taskKind).toBe('reflection');
  });

  it('budget window is UTC-day: row from 25h ago does NOT count toward today', async () => {
    const job = await repo.createJob({
      repo: 'local/task-5', issueNumber: 5,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    const cap = 1_000_000;
    const todayStartMs = Date.UTC(
      new Date().getUTCFullYear(),
      new Date().getUTCMonth(),
      new Date().getUTCDate(),
    );
    // Insert a row from 25 hours ago (before today's UTC window)
    const twentyFiveHoursAgoMs = Date.now() - 25 * 60 * 60 * 1000;
    // Only insert if the row would actually fall before today's start
    // (i.e., 25h ago is always yesterday or earlier for any UTC midnight)
    seedMetric('u-1', 900_000, 200_000, twentyFiveHoursAgoMs); // total = 1 100 000 > cap, but YESTERDAY

    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: cap,
    });

    // The 25h-ago row is before today's UTC start, so it must not count.
    // Today's window has 0 tokens → should enqueue.
    const queued = await repo.getJobsByStatus('queued');
    expect(queued).toHaveLength(1);
    expect(queued[0].taskKind).toBe('reflection');
  });

  // ── worker_required enforcement tests ───────────────────────────────────────

  it('workerRequired=true and no reflection worker → skips enqueue', async () => {
    const job = await repo.createJob({
      repo: 'local/task-10', issueNumber: 10,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: true, perUserDailyBudgetTokens: 0,
    }, [
      // workers without 'reflection' in roles
      { id: 'w1', endpoint: 'http://localhost:11434/v1', roles: ['auto', 'fast'] },
      { id: 'w2', endpoint: 'http://localhost:11435/v1', roles: ['quality'] },
    ]);

    expect(await repo.getJobsByStatus('queued')).toHaveLength(0);
  });

  it('workerRequired=true and a reflection worker exists → enqueues', async () => {
    const job = await repo.createJob({
      repo: 'local/task-11', issueNumber: 11,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: true, perUserDailyBudgetTokens: 0,
    }, [
      { id: 'w1', endpoint: 'http://localhost:11434/v1', roles: ['auto', 'fast'] },
      { id: 'w-reflect', endpoint: 'http://localhost:11436/v1', roles: ['reflection'] },
    ]);

    const queued = await repo.getJobsByStatus('queued');
    expect(queued).toHaveLength(1);
    expect(queued[0].taskKind).toBe('reflection');
  });

  it('workerRequired=false → enqueues regardless of workers list', async () => {
    const job = await repo.createJob({
      repo: 'local/task-12', issueNumber: 12,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    await repo.updateJob(job.id, { status: 'succeeded' });

    // Empty workers list — but workerRequired=false so the check is skipped
    await maybeEnqueueReflection(repo, job, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: 0,
    }, []);

    const queued = await repo.getJobsByStatus('queued');
    expect(queued).toHaveLength(1);
    expect(queued[0].taskKind).toBe('reflection');
  });

  it('different users have independent budgets', async () => {
    const jobU1 = await repo.createJob({
      repo: 'local/task-6', issueNumber: 6,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-1',
    } as any);
    const jobU2 = await repo.createJob({
      repo: 'local/task-7', issueNumber: 7,
      instruction: 'do thing', pieceName: 'chat',
      ownerId: 'u-2',
    } as any);
    await repo.updateJob(jobU1.id, { status: 'succeeded' });
    await repo.updateJob(jobU2.id, { status: 'succeeded' });

    const cap = 1_000_000;
    const todayStartMs = Date.UTC(
      new Date().getUTCFullYear(),
      new Date().getUTCMonth(),
      new Date().getUTCDate(),
    );
    // u-1 is over cap; u-2 has nothing
    seedMetric('u-1', 600_000, 500_000, todayStartMs + 1_000); // total = 1 100 000 > cap

    await maybeEnqueueReflection(repo, jobU1, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: cap,
    });
    await maybeEnqueueReflection(repo, jobU2, 'succeeded', {
      enabled: true, workerRequired: false, perUserDailyBudgetTokens: cap,
    });

    const queued = await repo.getJobsByStatus('queued');
    // Only u-2's reflection should be queued
    expect(queued).toHaveLength(1);
    expect(JSON.parse(queued[0].payload!).userId).toBe('u-2');
  });
});
