import { QueryClient } from '@tanstack/react-query';
import { beforeEach, describe, expect, it, vi } from 'vitest';

import type { BatchIdWithType, HistoryEntity } from './generationHistoryTypes';
import {
  generationQueryKeys,
  removeCompletedGenerations,
  updateGenerationHistoryCache,
} from './queryKeys';

// Mock the logWebUserEvent module
vi.mock('@/logging/logWebUserEvent', () => ({
  default: vi.fn(),
}));

describe('updateGenerationHistoryCache', () => {
  let queryClient: QueryClient;
  const clipId = 'test-clip-id';

  beforeEach(() => {
    queryClient = new QueryClient({
      defaultOptions: {
        queries: { retry: false },
      },
    });
  });

  it('should do nothing if completedBatches is empty', () => {
    const oldHistory: HistoryEntity[] = [
      {
        batchId: 'batch-1',
        status: 'processing',
        type: 'image',
        prompt: 'Test prompt',
        createdAt: '2024-01-01',
        items: [],
      },
    ];

    queryClient.setQueryData(generationQueryKeys.history(clipId), oldHistory);

    updateGenerationHistoryCache(queryClient, clipId, {});

    const result = queryClient.getQueryData(
      generationQueryKeys.history(clipId)
    );
    expect(result).toEqual(oldHistory);
  });

  it('should update items for matching batch IDs', () => {
    const oldHistory: HistoryEntity[] = [
      {
        batchId: 'batch-1',
        status: 'processing',
        type: 'image',
        prompt: 'Test prompt',
        createdAt: '2024-01-01',
        items: [
          {
            id: 'item-1',
            clipId,
            type: 'image',
            status: 'processing',
            isLiked: false,
          },
        ],
      },
    ];

    queryClient.setQueryData(generationQueryKeys.history(clipId), oldHistory);

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId,
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/image.jpg',
        },
      ],
    };

    updateGenerationHistoryCache(queryClient, clipId, completedBatches);

    const result = queryClient.getQueryData<HistoryEntity[]>(
      generationQueryKeys.history(clipId)
    );

    expect(result).toBeDefined();
    expect(result![0].items).toEqual(completedBatches['batch-1']);
  });

  it('should update batch status to complete when all items are complete', () => {
    const oldHistory: HistoryEntity[] = [
      {
        batchId: 'batch-1',
        status: 'processing',
        type: 'image',
        prompt: 'Test prompt',
        createdAt: '2024-01-01',
        items: [
          {
            id: 'item-1',
            clipId,
            type: 'image',
            status: 'processing',
            isLiked: false,
          },
          {
            id: 'item-2',
            clipId,
            type: 'image',
            status: 'processing',
            isLiked: false,
          },
        ],
      },
    ];

    queryClient.setQueryData(generationQueryKeys.history(clipId), oldHistory);

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId,
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/image1.jpg',
        },
        {
          id: 'item-2',
          clipId,
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/image2.jpg',
        },
      ],
    };

    updateGenerationHistoryCache(queryClient, clipId, completedBatches);

    const result = queryClient.getQueryData<HistoryEntity[]>(
      generationQueryKeys.history(clipId)
    );

    expect(result).toBeDefined();
    expect(result![0].status).toBe('complete');
    expect(result![0].items).toEqual(completedBatches['batch-1']);
  });

  it('should keep batch status as processing when some items are still processing', () => {
    const oldHistory: HistoryEntity[] = [
      {
        batchId: 'batch-1',
        status: 'processing',
        type: 'image',
        prompt: 'Test prompt',
        createdAt: '2024-01-01',
        items: [
          {
            id: 'item-1',
            clipId,
            type: 'image',
            status: 'processing',
            isLiked: false,
          },
          {
            id: 'item-2',
            clipId,
            type: 'image',
            status: 'processing',
            isLiked: false,
          },
        ],
      },
    ];

    queryClient.setQueryData(generationQueryKeys.history(clipId), oldHistory);

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId,
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/image1.jpg',
        },
        {
          id: 'item-2',
          clipId,
          type: 'image' as const,
          status: 'processing' as const,
          isLiked: false,
        },
      ],
    };

    updateGenerationHistoryCache(queryClient, clipId, completedBatches);

    const result = queryClient.getQueryData<HistoryEntity[]>(
      generationQueryKeys.history(clipId)
    );

    expect(result).toBeDefined();
    expect(result![0].status).toBe('processing');
  });

  it('should not modify batches that are not in completedBatches', () => {
    const oldHistory: HistoryEntity[] = [
      {
        batchId: 'batch-1',
        status: 'processing',
        type: 'image',
        prompt: 'Test prompt 1',
        createdAt: '2024-01-01',
        items: [
          {
            id: 'item-1',
            clipId,
            type: 'image',
            status: 'processing',
            isLiked: false,
          },
        ],
      },
      {
        batchId: 'batch-2',
        status: 'processing',
        type: 'video',
        prompt: 'Test prompt 2',
        createdAt: '2024-01-02',
        items: [
          {
            id: 'item-2',
            clipId,
            type: 'video',
            status: 'processing',
            isLiked: false,
          },
        ],
      },
    ];

    queryClient.setQueryData(generationQueryKeys.history(clipId), oldHistory);

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId,
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/image.jpg',
        },
      ],
    };

    updateGenerationHistoryCache(queryClient, clipId, completedBatches);

    const result = queryClient.getQueryData<HistoryEntity[]>(
      generationQueryKeys.history(clipId)
    );

    expect(result).toBeDefined();
    expect(result![0].items).toEqual(completedBatches['batch-1']);
    expect(result![1].items).toEqual(oldHistory[1].items); // Unchanged
  });

  it('should return early if oldHistory is undefined', () => {
    updateGenerationHistoryCache(queryClient, clipId, {
      'batch-1': [
        {
          id: 'item-1',
          clipId,
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/image.jpg',
        },
      ],
    });

    const result = queryClient.getQueryData(
      generationQueryKeys.history(clipId)
    );
    expect(result).toBeUndefined();
  });

  it('should handle mixed complete and error statuses', () => {
    const oldHistory: HistoryEntity[] = [
      {
        batchId: 'batch-1',
        status: 'processing',
        type: 'video',
        prompt: 'Test prompt',
        createdAt: '2024-01-01',
        items: [
          {
            id: 'item-1',
            clipId,
            type: 'video',
            status: 'processing',
            isLiked: false,
          },
          {
            id: 'item-2',
            clipId,
            type: 'video',
            status: 'processing',
            isLiked: false,
          },
        ],
      },
    ];

    queryClient.setQueryData(generationQueryKeys.history(clipId), oldHistory);

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId,
          type: 'video' as const,
          status: 'complete' as const,
          isLiked: false,
          url: 'https://example.com/video1.mp4',
          videoUploadId: 'upload-1',
        },
        {
          id: 'item-2',
          clipId,
          type: 'video' as const,
          status: 'error' as const,
          isLiked: false,
        },
      ],
    };

    updateGenerationHistoryCache(queryClient, clipId, completedBatches);

    const result = queryClient.getQueryData<HistoryEntity[]>(
      generationQueryKeys.history(clipId)
    );

    expect(result).toBeDefined();
    expect(result![0].status).toBe('complete'); // All items have terminal status
    expect(result![0].items).toEqual(completedBatches['batch-1']);
  });
});

describe('removeCompletedGenerations', () => {
  let queryClient: QueryClient;

  beforeEach(() => {
    queryClient = new QueryClient({
      defaultOptions: {
        queries: { retry: false },
      },
    });
  });

  it('should do nothing if completedBatches is empty', () => {
    const oldPendingBatchIds: BatchIdWithType[] = [
      { id: 'batch-1', type: 'image' },
      { id: 'batch-2', type: 'video' },
    ];

    queryClient.setQueryData(
      generationQueryKeys.pendingBatchIds(),
      oldPendingBatchIds
    );

    removeCompletedGenerations(queryClient, {});

    const result = queryClient.getQueryData(
      generationQueryKeys.pendingBatchIds()
    );
    expect(result).toEqual(oldPendingBatchIds);
  });

  it('should remove completed batch IDs from pending list', () => {
    const oldPendingBatchIds: BatchIdWithType[] = [
      { id: 'batch-1', type: 'image' },
      { id: 'batch-2', type: 'video' },
      { id: 'batch-3', type: 'image' },
    ];

    queryClient.setQueryData(
      generationQueryKeys.pendingBatchIds(),
      oldPendingBatchIds
    );

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId: 'clip-1',
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
        },
      ],
      'batch-3': [
        {
          id: 'item-3',
          clipId: 'clip-1',
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
        },
      ],
    };

    removeCompletedGenerations(queryClient, completedBatches);

    const result = queryClient.getQueryData<BatchIdWithType[]>(
      generationQueryKeys.pendingBatchIds()
    );

    expect(result).toBeDefined();
    expect(result).toHaveLength(1);
    expect(result![0]).toEqual({ id: 'batch-2', type: 'video' });
  });

  it('should handle all batches being completed', () => {
    const oldPendingBatchIds: BatchIdWithType[] = [
      { id: 'batch-1', type: 'image' },
      { id: 'batch-2', type: 'video' },
    ];

    queryClient.setQueryData(
      generationQueryKeys.pendingBatchIds(),
      oldPendingBatchIds
    );

    const completedBatches = {
      'batch-1': [
        {
          id: 'item-1',
          clipId: 'clip-1',
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
        },
      ],
      'batch-2': [
        {
          id: 'item-2',
          clipId: 'clip-1',
          type: 'video' as const,
          status: 'complete' as const,
          isLiked: false,
        },
      ],
    };

    removeCompletedGenerations(queryClient, completedBatches);

    const result = queryClient.getQueryData<BatchIdWithType[]>(
      generationQueryKeys.pendingBatchIds()
    );

    expect(result).toBeDefined();
    expect(result).toHaveLength(0);
  });

  it('should return early if oldPendingBatchIds is undefined', () => {
    removeCompletedGenerations(queryClient, {
      'batch-1': [
        {
          id: 'item-1',
          clipId: 'clip-1',
          type: 'image' as const,
          status: 'complete' as const,
          isLiked: false,
        },
      ],
    });

    const result = queryClient.getQueryData(
      generationQueryKeys.pendingBatchIds()
    );
    expect(result).toBeUndefined();
  });

  it('should only remove batches that are in completedBatches', () => {
    const oldPendingBatchIds: BatchIdWithType[] = [
      { id: 'batch-1', type: 'image' },
      { id: 'batch-2', type: 'video' },
      { id: 'batch-3', type: 'image' },
      { id: 'batch-4', type: 'video' },
    ];

    queryClient.setQueryData(
      generationQueryKeys.pendingBatchIds(),
      oldPendingBatchIds
    );

    const completedBatches = {
      'batch-2': [
        {
          id: 'item-2',
          clipId: 'clip-1',
          type: 'video' as const,
          status: 'complete' as const,
          isLiked: false,
        },
      ],
    };

    removeCompletedGenerations(queryClient, completedBatches);

    const result = queryClient.getQueryData<BatchIdWithType[]>(
      generationQueryKeys.pendingBatchIds()
    );

    expect(result).toBeDefined();
    expect(result).toHaveLength(3);
    expect(result).toEqual([
      { id: 'batch-1', type: 'image' },
      { id: 'batch-3', type: 'image' },
      { id: 'batch-4', type: 'video' },
    ]);
  });
});