/**
 * 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 { randomUUID } from 'crypto';
import { App } from 'firebase-admin/app';
import {
  DocumentReference,
  FieldValue,
  getFirestore,
  type Firestore,
} from 'firebase-admin/firestore';
import {
  GenkitError,
  StreamNotFoundError,
  type ActionStreamInput,
  type ActionStreamSubscriber,
  type StreamManager,
} from 'genkit/beta';

export interface FirestoreStreamManagerOptions {
  firebaseApp?: App;
  db?: Firestore;
  collection: string;
  timeout?: number;
}

class FirestoreActionStream<S, O> implements ActionStreamInput<S, O> {
  private streamDoc;

  constructor(streamDoc: DocumentReference) {
    this.streamDoc = streamDoc;
  }

  private async update(data: any): Promise<void> {
    await this.streamDoc.update({
      ...data,
      updatedAt: FieldValue.serverTimestamp(),
    });
  }

  async write(chunk: S): Promise<void> {
    await this.update({
      // We add a random ID to the chunk to prevent Firestore from deduplicating chunks
      // that have the same content.
      stream: FieldValue.arrayUnion({
        type: 'chunk',
        chunk,
        uuid: randomUUID(),
      }),
    });
  }

  async done(output: O): Promise<void> {
    await this.update({
      stream: FieldValue.arrayUnion({ type: 'done', output }),
    });
  }

  async error(err: any): Promise<void> {
    const serializableError = {
      message: err.message,
      stack: err.stack,
      ...err,
    };
    await this.update({
      stream: FieldValue.arrayUnion({ type: 'error', err: serializableError }),
    });
  }
}

export class FirestoreStreamManager implements StreamManager {
  private db: Firestore;
  private collection: string;
  private timeout: number;

  constructor(opts: FirestoreStreamManagerOptions) {
    this.collection = opts.collection;
    this.db =
      opts.db ??
      (opts.firebaseApp ? getFirestore(opts.firebaseApp) : getFirestore());
    this.timeout = opts.timeout ?? 60000;
  }

  async open<S, O>(streamId: string): Promise<ActionStreamInput<S, O>> {
    const streamDoc = this.db.collection(this.collection).doc(streamId);
    await streamDoc.set({
      createdAt: FieldValue.serverTimestamp(),
      stream: [],
    });
    return new FirestoreActionStream(streamDoc);
  }

  async subscribe<S, O>(
    streamId: string,
    callbacks: ActionStreamSubscriber<S, O>
  ): Promise<{
    unsubscribe: () => void;
  }> {
    const streamDoc = this.db.collection(this.collection).doc(streamId);
    const snapshot = await streamDoc.get();
    if (!snapshot.exists) {
      throw new StreamNotFoundError(`Stream ${streamId} not found.`);
    }
    let lastIndex = -1;
    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 unsubscribe = streamDoc.onSnapshot((snapshot) => {
      resetTimeout();
      const data = snapshot.data();
      if (!data) {
        return;
      }
      const stream = data.stream || [];
      for (let i = lastIndex + 1; i < stream.length; i++) {
        const event = stream[i];
        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();
        }
      }
      lastIndex = stream.length - 1;
    });
    resetTimeout();
    return { unsubscribe };
  }
}