// src/engine/reflection/snapshot.test.ts
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import {
  mkdtempSync,
  rmSync,
  existsSync,
  readFileSync,
  mkdirSync,
  writeFileSync,
} from 'fs';
import { join } from 'path';
import { tmpdir } from 'os';
import {
  writeSnapshot,
  revertSnapshot,
  revertSnapshotForUser,
  listSnapshots,
  readSnapshot,
  type SnapshotDeps,
  type WriteSnapshotMeta,
  type FileSnapshot,
} from './snapshot.js';

// ── Fixtures ──────────────────────────────────────────────────────────────────

const USER_ID = 'u-snap-test';

function makeDeps(dataDir: string, storeLlmRaw = false): SnapshotDeps {
  return { dataDir, storeLlmRaw };
}

function makeMeta(overrides: Partial<WriteSnapshotMeta> = {}): WriteSnapshotMeta {
  return {
    originalJobId: 'j-001',
    userId: USER_ID,
    pieceName: 'chat',
    outcome: 'applied',
    reasoning: 'The user prefers short answers.',
    modelUsed: 'qwen2.5:3b',
    tokensIn: 1200,
    tokensOut: 80,
    ratingAtTime: null,
    memoryChanges: 2,
    pieceEdited: false,
    ...overrides,
  };
}

const FIXED_DATE = new Date('2026-05-11T10:23:00Z');
const EXPECTED_TS_DIR = '20260511T102300Z-j-001';
const EXPECTED_ISO_TS = '2026-05-11T10:23:00Z';

// ── Test suite ────────────────────────────────────────────────────────────────

describe('writeSnapshot', () => {
  let tmpDir: string;

  beforeEach(() => {
    tmpDir = mkdtempSync(join(tmpdir(), 'snap-write-'));
  });

  afterEach(() => {
    rmSync(tmpDir, { recursive: true, force: true });
  });

  it('only captures changed files in memory.before and memory.after', async () => {
    const before: FileSnapshot = {
      'pref-terse.md': 'frontmatter: {}\n---\noriginal content',
    };
    const after: FileSnapshot = {
      'pref-terse.md': 'frontmatter: {}\n---\nupdated content',
    };

    const result = await writeSnapshot(makeDeps(tmpDir), before, after, makeMeta(), undefined, undefined, FIXED_DATE);

    const sDir = result.dir;
    // Before file is present with original content
    const beforeFile = join(sDir, 'memory.before', 'pref-terse.md');
    expect(existsSync(beforeFile)).toBe(true);
    expect(readFileSync(beforeFile, 'utf-8')).toBe('frontmatter: {}\n---\noriginal content');

    // After file is present with updated content
    const afterFile = join(sDir, 'memory.after', 'pref-terse.md');
    expect(existsSync(afterFile)).toBe(true);
    expect(readFileSync(afterFile, 'utf-8')).toBe('frontmatter: {}\n---\nupdated content');
  });

  it('does not create memory.before or memory.after dirs when no files changed', async () => {
    const result = await writeSnapshot(
      makeDeps(tmpDir), {}, {}, makeMeta({ memoryChanges: 0 }), undefined, undefined, FIXED_DATE,
    );
    expect(existsSync(join(result.dir, 'memory.before'))).toBe(false);
    expect(existsSync(join(result.dir, 'memory.after'))).toBe(false);
  });

  it('only writes files that changed (not the full memory dir)', async () => {
    // Simulate 5 memory files, but only 1 changed
    const before: FileSnapshot = { 'changed.md': 'old body' };
    const after: FileSnapshot = { 'changed.md': 'new body' };

    const result = await writeSnapshot(makeDeps(tmpDir), before, after, makeMeta(), undefined, undefined, FIXED_DATE);

    const sDir = result.dir;
    const beforeFiles = existsSync(join(sDir, 'memory.before'))
      ? require('fs').readdirSync(join(sDir, 'memory.before'))
      : [];
    const afterFiles = existsSync(join(sDir, 'memory.after'))
      ? require('fs').readdirSync(join(sDir, 'memory.after'))
      : [];

    // Should only contain 'changed.md', not any other hypothetical files
    expect(beforeFiles).toEqual(['changed.md']);
    expect(afterFiles).toEqual(['changed.md']);
  });

  it('appends exactly one row to index.jsonl in the expected shape', async () => {
    const before: FileSnapshot = { 'a.md': 'a content' };
    const after: FileSnapshot = { 'a.md': 'a updated', 'b.md': 'new b' };

    await writeSnapshot(
      makeDeps(tmpDir),
      before,
      after,
      makeMeta({ memoryChanges: 2, pieceName: 'research' }),
      undefined,
      undefined,
      FIXED_DATE,
    );

    const indexPath = join(tmpDir, USER_ID, '.reflection-history', 'index.jsonl');
    expect(existsSync(indexPath)).toBe(true);
    const lines = readFileSync(indexPath, 'utf-8').trim().split('\n');
    expect(lines).toHaveLength(1);

    const row = JSON.parse(lines[0]!);
    expect(row.ts).toBe(EXPECTED_ISO_TS);
    expect(row.snapshotId).toBe(EXPECTED_TS_DIR);
    expect(row.jobId).toBe('j-001');
    expect(row.pieceName).toBe('research');
    expect(row.memoryChanges).toBe(2);
    expect(row.pieceEdited).toBe(false);
    expect(row.reverted).toBe(false);
  });

  it('writes multiple rows when called multiple times', async () => {
    const date1 = new Date('2026-05-11T10:00:00Z');
    const date2 = new Date('2026-05-11T11:00:00Z');
    const meta1 = makeMeta({ originalJobId: 'j-001' });
    const meta2 = makeMeta({ originalJobId: 'j-002' });

    await writeSnapshot(makeDeps(tmpDir), {}, {}, meta1, undefined, undefined, date1);
    await writeSnapshot(makeDeps(tmpDir), {}, {}, meta2, undefined, undefined, date2);

    const indexPath = join(tmpDir, USER_ID, '.reflection-history', 'index.jsonl');
    const lines = readFileSync(indexPath, 'utf-8').trim().split('\n');
    expect(lines).toHaveLength(2);
    expect(JSON.parse(lines[0]!).jobId).toBe('j-001');
    expect(JSON.parse(lines[1]!).jobId).toBe('j-002');
  });

  it('meta.json contains outcome, reasoning, model, tokens', async () => {
    const meta = makeMeta({
      outcome: 'partial',
      reasoning: 'Learned one thing but rejected another.',
      modelUsed: 'claude-haiku',
      tokensIn: 999,
      tokensOut: 42,
      ratingAtTime: 'good',
    });

    const result = await writeSnapshot(makeDeps(tmpDir), {}, {}, meta, undefined, undefined, FIXED_DATE);

    const metaJson = JSON.parse(readFileSync(join(result.dir, 'meta.json'), 'utf-8'));
    expect(metaJson.outcome).toBe('partial');
    expect(metaJson.reasoning).toBe('Learned one thing but rejected another.');
    expect(metaJson.modelUsed).toBe('claude-haiku');
    expect(metaJson.tokensIn).toBe(999);
    expect(metaJson.tokensOut).toBe(42);
    expect(metaJson.ratingAtTime).toBe('good');
    expect(metaJson.snapshotId).toBe(EXPECTED_TS_DIR);
    expect(metaJson.ts).toBe(EXPECTED_ISO_TS);
  });

  it('does NOT write llm-raw.json when storeLlmRaw=false', async () => {
    const meta = makeMeta({ llmRaw: { memory_changes: [], piece_changes: {} } });
    const result = await writeSnapshot(makeDeps(tmpDir, false), {}, {}, meta, undefined, undefined, FIXED_DATE);
    expect(existsSync(join(result.dir, 'llm-raw.json'))).toBe(false);
  });

  it('writes llm-raw.json when storeLlmRaw=true', async () => {
    const rawPayload = { memory_changes: [{ op: 'add', name: 'x' }], piece_changes: {} };
    const meta = makeMeta({ llmRaw: rawPayload });
    const result = await writeSnapshot(makeDeps(tmpDir, true), {}, {}, meta, undefined, undefined, FIXED_DATE);

    const llmRawPath = join(result.dir, 'llm-raw.json');
    expect(existsSync(llmRawPath)).toBe(true);
    const parsed = JSON.parse(readFileSync(llmRawPath, 'utf-8'));
    expect(parsed).toEqual(rawPayload);
  });

  it('writes piece.before.yaml and piece.after.yaml when pieceEdited=true', async () => {
    const meta = makeMeta({ pieceEdited: true, pieceName: 'research' });
    const result = await writeSnapshot(
      makeDeps(tmpDir),
      {},
      {},
      meta,
      'before yaml content',
      'after yaml content',
      FIXED_DATE,
    );

    expect(readFileSync(join(result.dir, 'piece.before.yaml'), 'utf-8')).toBe('before yaml content');
    expect(readFileSync(join(result.dir, 'piece.after.yaml'), 'utf-8')).toBe('after yaml content');
  });

  it('does NOT write piece files when pieceEdited=false', async () => {
    const meta = makeMeta({ pieceEdited: false });
    const result = await writeSnapshot(
      makeDeps(tmpDir),
      {},
      {},
      meta,
      'should not be written',
      'should not be written',
      FIXED_DATE,
    );

    expect(existsSync(join(result.dir, 'piece.before.yaml'))).toBe(false);
    expect(existsSync(join(result.dir, 'piece.after.yaml'))).toBe(false);
  });

  it('writes diff.txt with a human-readable summary', async () => {
    const before: FileSnapshot = { 'existing.md': 'old' };
    const after: FileSnapshot = { 'existing.md': 'new', 'added.md': 'brand new' };

    const result = await writeSnapshot(makeDeps(tmpDir), before, after, makeMeta(), undefined, undefined, FIXED_DATE);

    const diff = readFileSync(join(result.dir, 'diff.txt'), 'utf-8');
    expect(diff).toContain('added.md');
    expect(diff).toContain('existing.md');
  });

  it('returns snapshotId and dir with correct naming', async () => {
    const result = await writeSnapshot(makeDeps(tmpDir), {}, {}, makeMeta(), undefined, undefined, FIXED_DATE);
    expect(result.snapshotId).toBe(EXPECTED_TS_DIR);
    expect(result.dir).toContain(EXPECTED_TS_DIR);
    expect(existsSync(result.dir)).toBe(true);
  });
});

// ── revertSnapshot tests ───────────────────────────────────────────────────────

describe('revertSnapshot', () => {
  let tmpDir: string;

  beforeEach(() => {
    tmpDir = mkdtempSync(join(tmpdir(), 'snap-revert-'));
  });

  afterEach(() => {
    rmSync(tmpDir, { recursive: true, force: true });
  });

  async function writeAndGetSnapshot(opts: {
    before?: FileSnapshot;
    after?: FileSnapshot;
    pieceEdited?: boolean;
    pieceBeforeYaml?: string;
    pieceAfterYaml?: string;
    jobId?: string;
  } = {}): Promise<{ snapshotId: string; liveMemDir: string }> {
    const before = opts.before ?? { 'changed.md': 'original content' };
    const after = opts.after ?? { 'changed.md': 'updated content' };
    const meta = makeMeta({
      originalJobId: opts.jobId ?? 'j-001',
      pieceEdited: opts.pieceEdited ?? false,
    });

    const result = await writeSnapshot(
      makeDeps(tmpDir),
      before,
      after,
      meta,
      opts.pieceBeforeYaml,
      opts.pieceAfterYaml,
      FIXED_DATE,
    );

    // Set up the live memory directory as if the applier already wrote the after state
    const liveMemDir = join(tmpDir, USER_ID, 'memory');
    mkdirSync(liveMemDir, { recursive: true });
    for (const [file, content] of Object.entries(after)) {
      writeFileSync(join(liveMemDir, file), content, 'utf-8');
    }

    return { snapshotId: result.snapshotId, liveMemDir };
  }

  it('idempotent: second revert returns reverted=false and does not re-apply', async () => {
    const { snapshotId } = await writeAndGetSnapshot();

    // First revert
    const r1 = await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId);
    expect(r1.reverted).toBe(true);

    // Second revert — should be no-op
    const r2 = await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId);
    expect(r2.reverted).toBe(false);

    // The index.jsonl should have exactly one reverted:true row (not two)
    const indexPath = join(tmpDir, USER_ID, '.reflection-history', 'index.jsonl');
    const lines = readFileSync(indexPath, 'utf-8').trim().split('\n').filter(Boolean);
    const revertRows = lines
      .map((l) => JSON.parse(l))
      .filter((r: Record<string, unknown>) => r['reverted'] === true);
    expect(revertRows).toHaveLength(1);
  });

  it('restores memory.before/* over the live memory directory', async () => {
    const { snapshotId, liveMemDir } = await writeAndGetSnapshot({
      before: { 'changed.md': 'original content' },
      after: { 'changed.md': 'updated content' },
    });

    // Live memory dir currently has the "after" version
    expect(readFileSync(join(liveMemDir, 'changed.md'), 'utf-8')).toBe('updated content');

    await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId);

    // After revert, should be restored to "original content"
    expect(readFileSync(join(liveMemDir, 'changed.md'), 'utf-8')).toBe('original content');
  });

  it('deletes files that were ADDED by the reflection (only in after, not in before)', async () => {
    // "added.md" appears only in after, not before → should be deleted on revert
    const { snapshotId, liveMemDir } = await writeAndGetSnapshot({
      before: { 'existing.md': 'existing content' },
      after: { 'existing.md': 'updated existing', 'added.md': 'brand new file' },
    });

    // Confirm the added file exists in live memory
    expect(existsSync(join(liveMemDir, 'added.md'))).toBe(true);

    await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId);

    // added.md should be gone
    expect(existsSync(join(liveMemDir, 'added.md'))).toBe(false);
    // existing.md should be restored to its original content
    expect(readFileSync(join(liveMemDir, 'existing.md'), 'utf-8')).toBe('existing content');
  });

  it('restores piece.before.yaml when it exists', async () => {
    const { snapshotId } = await writeAndGetSnapshot({
      before: {},
      after: {},
      pieceEdited: true,
      pieceBeforeYaml: 'original piece yaml',
      pieceAfterYaml: 'modified piece yaml',
    });

    // Set up the live pieces dir with the "after" version
    const livePiecesDir = join(tmpDir, USER_ID, 'pieces');
    mkdirSync(livePiecesDir, { recursive: true });
    writeFileSync(join(livePiecesDir, 'chat.yaml'), 'modified piece yaml', 'utf-8');

    await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId);

    // Should be restored to the before version
    expect(readFileSync(join(livePiecesDir, 'chat.yaml'), 'utf-8')).toBe('original piece yaml');
  });

  it('appends {reverted:true, refersTo:snapshotId} row to index.jsonl', async () => {
    const { snapshotId } = await writeAndGetSnapshot();

    await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId);

    const indexPath = join(tmpDir, USER_ID, '.reflection-history', 'index.jsonl');
    const lines = readFileSync(indexPath, 'utf-8').trim().split('\n').filter(Boolean);
    // Line 0: the original snapshot row
    // Line 1: the revert row
    expect(lines).toHaveLength(2);

    const revertRow = JSON.parse(lines[1]!);
    expect(revertRow.reverted).toBe(true);
    expect(revertRow.refersTo).toBe(snapshotId);
    expect(typeof revertRow.ts).toBe('string');
  });

  it('handles a snapshot with no memory changes gracefully (no memory dir)', async () => {
    const result = await writeSnapshot(
      makeDeps(tmpDir),
      {},
      {},
      makeMeta({ memoryChanges: 0 }),
      undefined,
      undefined,
      FIXED_DATE,
    );

    // Should not throw even if memory.before doesn't exist
    const r = await revertSnapshotForUser(makeDeps(tmpDir), USER_ID, result.snapshotId);
    expect(r.reverted).toBe(true);
  });

  it('concurrent revert calls: only one wins, second sees the index entry and aborts', async () => {
    const { snapshotId } = await writeAndGetSnapshot();

    // Fire both reverts concurrently — they must serialize via the lock
    const [r1, r2] = await Promise.all([
      revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId),
      revertSnapshotForUser(makeDeps(tmpDir), USER_ID, snapshotId),
    ]);

    // Exactly one should succeed, the other sees the index entry and aborts
    const successCount = [r1, r2].filter((r) => r.reverted).length;
    const abortCount = [r1, r2].filter((r) => !r.reverted).length;
    expect(successCount).toBe(1);
    expect(abortCount).toBe(1);

    // Only one reverted:true row in index.jsonl
    const indexPath = join(tmpDir, USER_ID, '.reflection-history', 'index.jsonl');
    const lines = readFileSync(indexPath, 'utf-8').trim().split('\n').filter(Boolean);
    const revertRows = lines
      .map((l) => JSON.parse(l))
      .filter((r: Record<string, unknown>) => r['reverted'] === true);
    expect(revertRows).toHaveLength(1);
  });
});

// ── listSnapshots and readSnapshot ────────────────────────────────────────────

describe('listSnapshots', () => {
  let tmpDir: string;

  beforeEach(() => {
    tmpDir = mkdtempSync(join(tmpdir(), 'snap-list-'));
  });

  afterEach(() => {
    rmSync(tmpDir, { recursive: true, force: true });
  });

  it('returns entries sorted most-recent first', async () => {
    const date1 = new Date('2026-05-10T09:00:00Z');
    const date2 = new Date('2026-05-11T10:00:00Z');

    await writeSnapshot(makeDeps(tmpDir), {}, {}, makeMeta({ originalJobId: 'j-001' }), undefined, undefined, date1);
    await writeSnapshot(makeDeps(tmpDir), {}, {}, makeMeta({ originalJobId: 'j-002' }), undefined, undefined, date2);

    const entries = listSnapshots(makeDeps(tmpDir), USER_ID);
    expect(entries[0]!.jobId).toBe('j-002');
    expect(entries[1]!.jobId).toBe('j-001');
  });

  it('respects the limit option', async () => {
    for (let i = 1; i <= 5; i++) {
      const d = new Date(`2026-05-${String(i).padStart(2, '0')}T10:00:00Z`);
      await writeSnapshot(makeDeps(tmpDir), {}, {}, makeMeta({ originalJobId: `j-00${i}` }), undefined, undefined, d);
    }

    const limited = listSnapshots(makeDeps(tmpDir), USER_ID, { limit: 3 });
    expect(limited).toHaveLength(3);
  });

  it('returns empty array when no history exists', () => {
    const entries = listSnapshots(makeDeps(tmpDir), USER_ID);
    expect(entries).toEqual([]);
  });
});

describe('readSnapshot', () => {
  let tmpDir: string;

  beforeEach(() => {
    tmpDir = mkdtempSync(join(tmpdir(), 'snap-read-'));
  });

  afterEach(() => {
    rmSync(tmpDir, { recursive: true, force: true });
  });

  it('returns null for a non-existent snapshotId', () => {
    const detail = readSnapshot(makeDeps(tmpDir), USER_ID, 'nonexistent-id');
    expect(detail).toBeNull();
  });

  it('returns full detail including beforeFiles and afterFiles', async () => {
    const before: FileSnapshot = { 'entry.md': 'original' };
    const after: FileSnapshot = { 'entry.md': 'updated' };

    const { snapshotId } = await writeSnapshot(
      makeDeps(tmpDir),
      before,
      after,
      makeMeta({ reasoning: 'test reasoning' }),
      undefined,
      undefined,
      FIXED_DATE,
    );

    const detail = readSnapshot(makeDeps(tmpDir), USER_ID, snapshotId);
    expect(detail).not.toBeNull();
    expect(detail!.reasoning).toBe('test reasoning');
    expect(detail!.beforeFiles['entry.md']).toBe('original');
    expect(detail!.afterFiles['entry.md']).toBe('updated');
    expect(detail!.snapshotId).toBe(snapshotId);
  });

  it('includes pieceBeforeYaml and pieceAfterYaml when pieceEdited=true', async () => {
    const { snapshotId } = await writeSnapshot(
      makeDeps(tmpDir),
      {},
      {},
      makeMeta({ pieceEdited: true }),
      'before yaml',
      'after yaml',
      FIXED_DATE,
    );

    const detail = readSnapshot(makeDeps(tmpDir), USER_ID, snapshotId);
    expect(detail!.pieceBeforeYaml).toBe('before yaml');
    expect(detail!.pieceAfterYaml).toBe('after yaml');
  });
});
