/**
 * 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 { App } from 'firebase-admin/app';
import { Database, getDatabase } from 'firebase-admin/database';
import {
  GenkitError,
  StreamNotFoundError,
  type ActionStreamInput,
  type ActionStreamSubscriber,
  type StreamManager,
} from 'genkit/beta';

export interface RtdbStreamManagerOptions {
  firebaseApp?: App;
  db?: Database;
  refPrefix?: string;
  timeout?: number;
}

class RtdbActionStream<S, O> implements ActionStreamInput<S, O> {
  private db: Database;
  private streamRef: string;
  private metadataRef: string;

  constructor(db: Database, streamRef: string) {
    this.db = db;
    this.streamRef = `${streamRef}/stream`;
    this.metadataRef = `${streamRef}/metadata`;
  }

  private async update(): Promise<void> {
    await this.db.ref(this.metadataRef).update({ updatedAt: Date.now() });
  }

  async write(chunk: S): Promise<void> {
    await this.db.ref(this.streamRef).push({ type: 'chunk', chunk });
    await this.update();
  }

  async done(output: O): Promise<void> {
    await this.db.ref(this.streamRef).push({ type: 'done', output });
    await this.update();
  }

  async error(err: any): Promise<void> {
    const serializableError = {
      message: err.message,
      stack: err.stack,
      ...err,
    };
    await this.db
      .ref(this.streamRef)
      .push({ type: 'error', err: serializableError });
    await this.update();
  }
}

export class RtdbStreamManager implements StreamManager {
  private db: Database;
  private refRoot: string;
  private timeout: number;

  constructor(opts: RtdbStreamManagerOptions) {
    this.refRoot = opts.refPrefix ?? 'genkit-streams';
    // refRoot should have a trailing slash for simplicity.
    if (this.refRoot && !this.refRoot.endsWith('/')) {
      this.refRoot += '/';
    }
    this.db =
      opts.db ??
      (opts.firebaseApp ? getDatabase(opts.firebaseApp) : getDatabase());
    this.timeout = opts.timeout ?? 60000;
  }

  async open<S, O>(streamId: string): Promise<ActionStreamInput<S, O>> {
    const streamRef = this.refRoot + streamId;
    await this.db.ref(streamRef).remove();
    await this.db.ref(`${streamRef}/metadata`).set({ createdAt: Date.now() });
    return new RtdbActionStream(this.db, streamRef);
  }

  async subscribe<S, O>(
    streamId: string,
    callbacks: ActionStreamSubscriber<S, O>
  ): Promise<{
    unsubscribe: () => void;
  }> {
    const streamRef = this.db.ref(`${this.refRoot}${streamId}`);
    const snapshot = await streamRef.get();
    if (!snapshot.exists()) {
      throw new StreamNotFoundError(`Stream ${streamId} not found.`);
    }
    const streamDataRef = streamRef.child('stream');
    let timeoutId: NodeJS.Timeout;
    const resetTimeout = () => {
      clearTimeout(timeoutId);
      timeoutId = setTimeout(() => {
        callbacks.onError?.(
          new GenkitError({
            status: 'DEADLINE_EXCEEDED',
            message: 'Stream timed out.',
          })
        );
        unsubscribe();
      }, this.timeout);
    };
    const handleEvent = (snapshot) => {
      resetTimeout();
      if (!snapshot.exists()) {
        return;
      }
      const event = snapshot.val();
      if (event.type === 'chunk') {
        callbacks.onChunk?.(event.chunk);
      } else if (event.type === 'done') {
        clearTimeout(timeoutId);
        callbacks.onDone?.(event.output);
        unsubscribe();
      } else if (event.type === 'error') {
        clearTimeout(timeoutId);
        callbacks.onError?.(event.err);
        unsubscribe();
      }
    };
    const unsubscribe = () => {
      streamDataRef.off('child_added', handleEvent);
    };
    streamDataRef.on('child_added', handleEvent);
    resetTimeout();
    return { unsubscribe };
  }
}