import { renderToStaticMarkup } from 'react-dom/server';
import { describe, expect, it } from 'vitest';
import type { ModelLoadProgress, ModelShardProgress } from '../../lib/contracts';
import { ComposerActivity, describeLoadProgress, overallLoadProgress } from './ComposerActivity';

function shard(index: number, state: ModelShardProgress['state']): ModelShardProgress {
  return {
    index,
    path: `shard-${index + 1}.gguf`,
    loadedBytes: state === 'queued' ? 0 : 100,
    verifiedBytes: state === 'complete' || state === 'cached' ? 100 : 0,
    totalBytes: 100,
    state,
  };
}

function progress(overrides: Partial<ModelLoadProgress> = {}): ModelLoadProgress {
  return {
    modelId: 'bonsai-27b',
    modelLabel: 'Bonsai 27B',
    stage: 'download',
    stageProgress: 0.5,
    downloadedBytes: 400,
    totalBytes: 800,
    residentBytes: 0,
    nativeStage: null,
    shards: [
      shard(0, 'complete'),
      shard(1, 'downloading'),
      ...Array.from({ length: 6 }, (_, index) => shard(index + 2, 'queued')),
    ],
    ...overrides,
  };
}

describe('ComposerActivity', () => {
  it('combines the load pipeline into one weighted overall progress value', () => {
    expect(overallLoadProgress(progress())).toBeCloseTo(0.39625);
    expect(overallLoadProgress(progress({ stage: 'cache', stageProgress: 0.2 }))).toBeCloseTo(0.39625);
    expect(overallLoadProgress(progress({ stage: 'load', stageProgress: 0.5 }))).toBeCloseTo(0.95);
    expect(overallLoadProgress(progress({ stage: 'ready' }))).toBe(1);
  });

  it('annotates the active shard and total ready shard count', () => {
    expect(describeLoadProgress(progress())).toEqual({
      label: 'Downloading shard 2 of 8',
      detail: '1 of 8 shards ready · 400 B / 800 B',
    });
  });

  it('renders exactly one progressbar during load and disappears afterwards', () => {
    const loading = renderToStaticMarkup(
      <ComposerActivity modelId="bonsai-27b" streamState="loading" progress={progress()} />,
    );
    expect(loading.match(/role="progressbar"/g)).toHaveLength(1);
    expect(loading).not.toContain('Model transfer');
    expect(loading).not.toContain('Load pipeline');

    const ready = renderToStaticMarkup(
      <ComposerActivity modelId="bonsai-27b" streamState="complete" progress={progress({ stage: 'ready' })} />,
    );
    expect(ready).toBe('');
  });
});