import { describe, it, expect, vi } from 'vitest';
import type { StageContext, LLMProvider, CacheProvider, StageLogger } from '../src/proxymodel/types.js';
import passthrough from '../src/proxymodel/stages/passthrough.js';
import paginate from '../src/proxymodel/stages/paginate.js';
import sectionSplit from '../src/proxymodel/stages/section-split.js';
import summarizeTree from '../src/proxymodel/stages/summarize-tree.js';
import { BUILT_IN_STAGES } from '../src/proxymodel/stages/index.js';

function mockCtx(original: string, config: Record<string, unknown> = {}, llmAvailable = false): StageContext {
  const llmResponses: string[] = [];

  const mockLlm: LLMProvider = {
    async complete(prompt) {
      const response = `Summary of: ${prompt.slice(0, 40)}...`;
      llmResponses.push(response);
      return response;
    },
    available: () => llmAvailable,
  };

  const cache = new Map<string, string>();
  const mockCache: CacheProvider = {
    async getOrCompute(key, compute) {
      if (cache.has(key)) return cache.get(key)!;
      const val = await compute();
      cache.set(key, val);
      return val;
    },
    hash(content) { return content.slice(0, 8); },
    async get(key) { return cache.get(key) ?? null; },
    async set(key, value) { cache.set(key, value); },
  };

  const mockLog: StageLogger = {
    debug: vi.fn(),
    info: vi.fn(),
    warn: vi.fn(),
    error: vi.fn(),
  };

  return {
    contentType: 'toolResult',
    sourceName: 'test/tool',
    projectName: 'test',
    sessionId: 'sess-1',
    originalContent: original,
    llm: mockLlm,
    cache: mockCache,
    log: mockLog,
    getSystemPrompt: async (_name: string, fallback: string) => fallback,
    config,
  };
}

describe('passthrough stage', () => {
  it('returns content unchanged', async () => {
    const result = await passthrough('hello world', mockCtx('hello world'));
    expect(result.content).toBe('hello world');
    expect(result.sections).toBeUndefined();
  });
});

describe('paginate stage', () => {
  it('returns small content unchanged', async () => {
    const result = await paginate('small', mockCtx('small', { pageSize: 100 }));
    expect(result.content).toBe('small');
    expect(result.sections).toBeUndefined();
  });

  it('splits large content into pages', async () => {
    const content = 'line\n'.repeat(500); // ~2500 chars
    const result = await paginate(content, mockCtx(content, { pageSize: 500 }));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBeGreaterThan(1);
    expect(result.content).toContain('pages');
    // Each section should have an id like page-1, page-2
    expect(result.sections![0].id).toBe('page-1');
  });

  it('page sections contain the actual content', async () => {
    const content = 'A'.repeat(1000) + '\n' + 'B'.repeat(1000);
    const result = await paginate(content, mockCtx(content, { pageSize: 500 }));
    expect(result.sections).toBeDefined();
    // Joining all section content should approximate the original
    const reassembled = result.sections!.map((s) => s.content).join('');
    expect(reassembled.length).toBe(content.length);
  });
});

describe('section-split stage', () => {
  it('splits markdown by headers', async () => {
    const md = '## Introduction\nSome intro text.\n## Methods\nMethod details.\n## Results\nResult data.';
    const result = await sectionSplit(md, mockCtx(md, { minSectionSize: 5 }));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBe(3);
    expect(result.sections![0].title).toBe('Introduction');
    expect(result.sections![1].title).toBe('Methods');
  });

  it('splits JSON array into elements', async () => {
    const arr = JSON.stringify([
      { id: 'flow1', label: 'Thermostat', nodes: [1, 2, 3] },
      { id: 'flow2', label: 'Lighting', nodes: [4, 5] },
      { id: 'flow3', label: 'Security', nodes: [6, 7, 8, 9] },
    ]);
    const result = await sectionSplit(arr, mockCtx(arr, { minSectionSize: 5 }));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBe(3);
    expect(result.sections![0].title).toBe('Thermostat');
    expect(result.sections![1].title).toBe('Lighting');
  });

  it('splits JSON object by keys', async () => {
    const obj = JSON.stringify({
      config: { port: 3000 },
      users: [{ name: 'alice' }, { name: 'bob' }],
      metadata: { version: '1.0' },
    });
    const result = await sectionSplit(obj, mockCtx(obj, { minSectionSize: 5 }));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBe(3);
    expect(result.sections!.map((s) => s.title)).toEqual(['config', 'users', 'metadata']);
  });

  it('returns small content unchanged', async () => {
    const result = await sectionSplit('tiny', mockCtx('tiny'));
    expect(result.content).toBe('tiny');
    expect(result.sections).toBeUndefined();
  });

  it('splits YAML by top-level keys', async () => {
    const yaml = 'name: test\n  sub: value\nversion: 1.0\n  build: 42\ndescription: hello world\n  more: stuff';
    const result = await sectionSplit(yaml, mockCtx(yaml, { minSectionSize: 5 }));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBe(3);
  });

  it('leaf content is exact original (not rewritten)', async () => {
    const arr = JSON.stringify([
      { id: 'flow1', label: 'Thermostat', data: { complex: true, nested: { deep: 'value' } } },
      { id: 'flow2', label: 'Lighting', data: { complex: false } },
    ]);
    const result = await sectionSplit(arr, mockCtx(arr, { minSectionSize: 5 }));
    // Each section content should be valid JSON matching the original item
    for (const section of result.sections!) {
      const parsed = JSON.parse(section.content);
      expect(parsed).toBeDefined();
      expect(parsed.id).toBeDefined();
    }
  });
});

describe('summarize-tree stage', () => {
  it('returns small content unchanged', async () => {
    const result = await summarizeTree('tiny content', mockCtx('tiny content'));
    expect(result.content).toBe('tiny content');
  });

  it('creates structural summary for JSON without LLM', async () => {
    const bigJson = JSON.stringify(Array.from({ length: 10 }, (_, i) => ({
      id: `item-${i}`,
      name: `Item ${i}`,
      data: 'x'.repeat(300),
    })));
    const result = await summarizeTree(bigJson, mockCtx(bigJson));
    expect(result.sections).toBeDefined();
    expect(result.sections!.length).toBeGreaterThan(0);
    expect(result.content).toContain('sections');
  });

  it('uses LLM for prose summaries when available', async () => {
    const prose = '## Security\n' + 'Important security details. '.repeat(200) +
      '\n## Performance\n' + 'Performance metrics and analysis. '.repeat(200);
    const result = await summarizeTree(prose, mockCtx(prose, {}, true));
    expect(result.sections).toBeDefined();
    expect(result.content).toContain('sections');
  });

  it('sections provide drill-down to full content', async () => {
    const items = Array.from({ length: 5 }, (_, i) => ({
      id: `flow-${i}`,
      label: `Flow ${i}`,
      config: { nodes: Array.from({ length: 20 }, (_, j) => ({ id: `node-${j}`, type: 'function' })) },
    }));
    const json = JSON.stringify(items);
    const result = await summarizeTree(json, mockCtx(json));

    // Drill-down sections should contain parseable JSON
    if (result.sections) {
      for (const section of result.sections) {
        // Content should be parseable (exact original JSON)
        try {
          JSON.parse(section.content);
        } catch {
          // Some sections may be ToC text, not raw JSON — that's OK
        }
      }
    }
  });
});

describe('BUILT_IN_STAGES registry', () => {
  it('contains all four built-in stages', () => {
    expect(BUILT_IN_STAGES.has('passthrough')).toBe(true);
    expect(BUILT_IN_STAGES.has('paginate')).toBe(true);
    expect(BUILT_IN_STAGES.has('section-split')).toBe(true);
    expect(BUILT_IN_STAGES.has('summarize-tree')).toBe(true);
    expect(BUILT_IN_STAGES.size).toBe(4);
  });

  it('all stages are callable functions', () => {
    for (const [name, handler] of BUILT_IN_STAGES) {
      expect(typeof handler).toBe('function');
    }
  });
});
