/**
 * 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 { getDatabase } from 'firebase-admin/database';
import { RtdbStreamManager } from '../src/stream-manager/rtdb';

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

  beforeEach(() => {
    app = initializeApp({
      databaseURL: 'http://127.0.0.1:9000?ns=genkit-test',
    });
    streamManager = new RtdbStreamManager({
      firebaseApp: app,
    });
  });

  afterEach(async () => {
    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 = getDatabase(app);
    const snapshot = await db.ref(`genkit-streams/${streamId}`).get();
    const data = snapshot.val();
    const values = Object.values(data.stream);

    assert.deepStrictEqual(values, [
      { type: 'chunk', chunk: { foo: 'bar' } },
      { type: 'chunk', chunk: { bar: 'baz' } },
    ]);
    assert.ok(data.metadata.createdAt);
    assert.ok(data.metadata.updatedAt);
  });

  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 = getDatabase(app);
    const snapshot = await db.ref(`genkit-streams/${streamId}`).get();
    const data = snapshot.val();
    const values = Object.values(data.stream);

    assert.deepStrictEqual(values, [
      { type: 'done', output: { result: 'success' } },
    ]);
    assert.ok(data.metadata.createdAt);
    assert.ok(data.metadata.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 = getDatabase(app);
    const snapshot = await db.ref(`genkit-streams/${streamId}`).get();
    const data = snapshot.val();
    const values = Object.values(data.stream) as any[];

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

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

    const db = getDatabase(app);
    const streamRef = db.ref(`genkit-streams/${streamId}`);
    streamRef.set({ metadata: { createdAt: Date.now() } }).then(async () => {
      await 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);
        },
      });

      const streamDataRef = streamRef.child('stream');
      streamDataRef.push({ type: 'chunk', chunk: { foo: 'bar' } });
      streamDataRef.push({ type: 'chunk', chunk: { bar: 'baz' } });
      streamDataRef.push({ type: 'done', output: { result: 'success' } });
    });
  });

  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 RtdbStreamManager({
      firebaseApp: app,
      refPrefix: '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();
        },
      });
    });
  });
});