import { useRef } from 'react';
import type {
  ModelId,
  ModelLoadProgress,
  StreamState,
} from '../../lib/contracts';
import { formatBytes, formatPercent } from '../../lib/format';

interface ComposerActivityProps {
  modelId: ModelId;
  streamState: StreamState;
  progress: ModelLoadProgress;
}

function ratio(value: number, total: number): number {
  return total > 0 ? Math.min(1, Math.max(0, value / total)) : 0;
}

function humanizeStage(value: string | null): string {
  if (!value) return '';
  return value.replaceAll('_', ' ').replace(/\b\w/g, (letter) => letter.toUpperCase());
}

export function overallLoadProgress(progress: ModelLoadProgress): number {
  const stageProgress = Math.min(1, Math.max(0, progress.stageProgress));
  const loadedShardBytes = progress.shards.reduce((total, shard) => (
    total + (shard.state === 'cached' || shard.state === 'complete'
      ? shard.totalBytes
      : Math.min(shard.loadedBytes, shard.totalBytes))
  ), 0);
  const verifiedBytes = progress.shards.reduce((total, shard) => (
    total + (shard.state === 'cached' || shard.state === 'complete'
      ? shard.totalBytes
      : Math.min(shard.verifiedBytes, shard.totalBytes))
  ), 0);
  const downloadRatio = ratio(Math.max(progress.downloadedBytes, loadedShardBytes), progress.totalBytes);
  const verifiedRatio = ratio(verifiedBytes, progress.totalBytes);
  const transferProgress = Math.min(0.9, 0.08 + downloadRatio * 0.57 + verifiedRatio * 0.25);

  switch (progress.stage) {
    case 'manifest':
      return stageProgress * 0.03;
    case 'cache':
      return Math.max(0.03 + stageProgress * 0.05, transferProgress);
    case 'download':
    case 'verify':
      return transferProgress;
    case 'load':
      return 0.9 + stageProgress * 0.1;
    case 'ready':
      return 1;
    default:
      return 0;
  }
}

export function describeLoadProgress(progress: ModelLoadProgress): { label: string; detail: string } {
  const shardCount = progress.shards.length;
  const readyShards = progress.shards.filter((shard) => shard.state === 'cached' || shard.state === 'complete').length;
  const activeShard = progress.shards.find((shard) => shard.state === 'downloading' || shard.state === 'verifying');
  const shardDetail = shardCount > 0 ? `${readyShards} of ${shardCount} shards ready` : 'Preparing shard map';
  const byteDetail = progress.totalBytes > 0
    ? `${formatBytes(progress.downloadedBytes)} / ${formatBytes(progress.totalBytes)}`
    : '';

  switch (progress.stage) {
    case 'manifest':
      return { label: 'Reading model manifest', detail: progress.modelLabel || 'Pinned release metadata' };
    case 'cache':
      return { label: 'Checking local cache', detail: shardCount > 0 ? shardDetail : 'Looking for existing shards' };
    case 'download':
      return {
        label: activeShard ? `Downloading shard ${activeShard.index + 1} of ${shardCount}` : 'Downloading model shards',
        detail: [shardDetail, byteDetail].filter(Boolean).join(' · '),
      };
    case 'verify':
      return {
        label: activeShard ? `Verifying shard ${activeShard.index + 1} of ${shardCount}` : 'Verifying model shards',
        detail: [shardDetail, byteDetail].filter(Boolean).join(' · '),
      };
    case 'load':
      return {
        label: progress.nativeStage ? humanizeStage(progress.nativeStage) : 'Allocating model',
        detail: progress.residentBytes > 0 ? `${formatBytes(progress.residentBytes)} resident` : 'Building the local graph',
      };
    default:
      return { label: 'Preparing model', detail: progress.modelLabel };
  }
}

export function ComposerActivity({ modelId, streamState, progress }: ComposerActivityProps) {
  const displayedProgress = useRef<{ modelId: ModelId | null; value: number }>({ modelId: null, value: 0 });
  const isLoading = streamState === 'loading' && progress.modelId === modelId;
  if (!isLoading) {
    displayedProgress.current = { modelId: null, value: 0 };
    return null;
  }

  const measuredProgress = overallLoadProgress(progress);
  const overallProgress = displayedProgress.current.modelId === modelId
    ? Math.max(displayedProgress.current.value, measuredProgress)
    : measuredProgress;
  displayedProgress.current = { modelId, value: overallProgress };
  const status = describeLoadProgress(progress);
  const percentage = Math.round(overallProgress * 100);

  return (
    <section className="composer-activity" aria-label="Model load progress">
      <div className="activity-status" aria-live="polite" aria-atomic="true">
        <div>
          <span className="activity-pulse" aria-hidden="true" />
          <span>
            <strong>{status.label}</strong>
            <small>{status.detail}</small>
          </span>
        </div>
        <b>{formatPercent(overallProgress)}</b>
      </div>
      <div
        className="activity-track"
        role="progressbar"
        aria-label={`${status.label}: ${status.detail}`}
        aria-valuemin={0}
        aria-valuemax={100}
        aria-valuenow={percentage}
      >
        <span style={{ width: `${percentage}%` }} />
      </div>
    </section>
  );
}