/**
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import { afterEach, beforeEach, describe, it } from '@jest/globals';
import * as assert from 'assert';
import { deleteApp, initializeApp, type App } from 'firebase-admin/app';
import { getFirestore } from 'firebase-admin/firestore';
import { FirestoreStreamManager } from '../src/stream-manager/firestore';

describe('FirestoreStreamManager', () => {
  let app: App;
  let streamManager: FirestoreStreamManager;

  beforeEach(() => {
    process.env.FIRESTORE_EMULATOR_HOST = '127.0.0.1:8080';
    app = initializeApp({
      projectId: 'genkit-test',
    });
    streamManager = new FirestoreStreamManager({
      firebaseApp: app,
      collection: 'genkit-streams',
    });
  });

  afterEach(async () => {
    const db = getFirestore(app);
    const collections = await db.listCollections();
    for (const collection of collections) {
      await db.recursiveDelete(collection);
    }
    await deleteApp(app);
  });

  it('should open a stream and write chunks', async () => {
    const streamId = 'test-stream-1';
    const stream = await streamManager.open(streamId);

    await stream.write({ foo: 'bar' });
    await stream.write({ bar: 'baz' });

    const db = getFirestore(app);
    const snapshot = await db.collection('genkit-streams').doc(streamId).get();
    const data = snapshot.data();
    assert.ok(data);

    assert.strictEqual(data.stream.length, 2);
    assert.strictEqual(data.stream[0].type, 'chunk');
    assert.deepStrictEqual(data.stream[0].chunk, { foo: 'bar' });
    assert.ok(data.stream[0].uuid);
    assert.strictEqual(data.stream[1].type, 'chunk');
    assert.deepStrictEqual(data.stream[1].chunk, { bar: 'baz' });
    assert.ok(data.stream[1].uuid);
    assert.ok(data.createdAt);
    assert.ok(data.updatedAt);
  });

  it('should preserve duplicate chunks', async () => {
    const streamId = 'test-stream-dupes';
    const stream = await streamManager.open(streamId);

    await stream.write({ foo: 'bar' });
    await stream.write({ foo: 'bar' });

    const db = getFirestore(app);
    const snapshot = await db.collection('genkit-streams').doc(streamId).get();
    const data = snapshot.data();
    assert.ok(data);

    assert.strictEqual(data.stream.length, 2);
    assert.deepStrictEqual(data.stream[0].chunk, { foo: 'bar' });
    assert.deepStrictEqual(data.stream[1].chunk, { foo: 'bar' });
    assert.notStrictEqual(data.stream[0].uuid, data.stream[1].uuid);
  });

  it('should open a stream and write done', async () => {
    const streamId = 'test-stream-2';
    const stream = await streamManager.open(streamId);

    await stream.done({ result: 'success' });

    const db = getFirestore(app);
    const snapshot = await db.collection('genkit-streams').doc(streamId).get();
    const data = snapshot.data();
    assert.ok(data);

    assert.deepStrictEqual(data.stream, [
      { type: 'done', output: { result: 'success' } },
    ]);
    assert.ok(data.createdAt);
    assert.ok(data.updatedAt);
  });

  it('should open a stream and write error', async () => {
    const streamId = 'test-stream-3';
    const stream = await streamManager.open(streamId);

    await stream.error(new Error('test error'));

    const db = getFirestore(app);
    const snapshot = await db.collection('genkit-streams').doc(streamId).get();
    const data = snapshot.data();
    assert.ok(data);

    assert.strictEqual(data.stream.length, 1);
    assert.strictEqual(data.stream[0].type, 'error');
    assert.strictEqual(data.stream[0].err.message, 'test error');
    assert.ok(data.createdAt);
    assert.ok(data.updatedAt);
  });

  it('should subscribe to a stream', (done) => {
    const streamId = 'test-stream-4';
    const chunks: any[] = [];
    let finalOutput: any;

    const db = getFirestore(app);
    const streamDoc = db.collection('genkit-streams').doc(streamId);
    streamDoc
      .set({
        stream: [
          { type: 'chunk', chunk: { foo: 'bar' } },
          { type: 'chunk', chunk: { bar: 'baz' } },
          { type: 'done', output: { result: 'success' } },
        ],
      })
      .then(() => {
        streamManager.subscribe(streamId, {
          onChunk: (chunk) => {
            chunks.push(chunk);
          },
          onDone: (output) => {
            finalOutput = output;
            assert.deepStrictEqual(chunks, [{ foo: 'bar' }, { bar: 'baz' }]);
            assert.deepStrictEqual(finalOutput, { result: 'success' });
            done();
          },
          onError(error) {
            console.log(error);
          },
        });
      });
  });

  it('should throw an error when subscribing to a non-existent stream', async () => {
    await assert.rejects(
      streamManager.subscribe('non-existent-stream', {
        onChunk: () => {},
        onDone: () => {},
        onError: () => {},
      }),
      (err: any) => {
        assert.strictEqual(err.name, 'StreamNotFoundError');
        return true;
      }
    );
  });

  it('should time out when no chunks are received', (done) => {
    const streamId = 'test-stream-5';
    streamManager = new FirestoreStreamManager({
      firebaseApp: app,
      collection: 'genkit-streams',
      timeout: 100,
    });
    streamManager.open(streamId).then(() => {
      streamManager.subscribe(streamId, {
        onChunk: () => {
          assert.fail('should not have received a chunk');
        },
        onDone: () => {
          assert.fail('should not have received done');
        },
        onError: (err) => {
          assert.strictEqual(err.status, 'DEADLINE_EXCEEDED');
          done();
        },
      });
    });
  });
});