/**
 * 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 {
  Action,
  AsyncTaskQueue,
  StreamNotFoundError,
  type ActionContext,
  type ActionStreamInput,
  type StreamManager,
  type z,
} from 'genkit/beta';
import {
  getCallableJSON,
  getHttpStatus,
  type ContextProvider,
  type RequestData,
} from 'genkit/context';
import { NextRequest, NextResponse } from 'next/server.js';
export { NextRequest, NextResponse, z, type Action, type ActionContext };

const delimiter = '\n\n';

async function subscribeToStream<S, O>(
  streamManager: StreamManager,
  streamId: string
): Promise<NextResponse | null> {
  try {
    const encoder = new TextEncoder();
    const { readable, writable } = new TransformStream();
    const writer = writable.getWriter();
    await streamManager.subscribe(streamId, {
      onChunk: (chunk) => {
        writer.write(
          encoder.encode(
            'data: ' + JSON.stringify({ message: chunk }) + delimiter
          )
        );
      },
      onDone: (output) => {
        writer.write(
          encoder.encode(
            'data: ' + JSON.stringify({ result: output }) + delimiter
          )
        );
        writer.write(encoder.encode('END'));
        writer.close();
      },
      onError: (err) => {
        console.error(
          `Streaming request failed with error: ${(err as Error).message}\n${
            (err as Error).stack
          }`
        );
        writer.write(
          encoder.encode(
            `error: ${JSON.stringify({
              error: getCallableJSON(err),
            })}${delimiter}`
          )
        );
        writer.write(encoder.encode('END'));
        writer.close();
      },
    });
    return new NextResponse(readable, {
      status: 200,
      headers: {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        Connection: 'keep-alive',
        'Transfer-Encoding': 'chunked',
        'x-genkit-stream-id': streamId,
      },
    });
  } catch (e: any) {
    if (e instanceof StreamNotFoundError) {
      return new NextResponse(null, { status: 204 });
    }
    if (e.status === 'DEADLINE_EXCEEDED') {
      const encoder = new TextEncoder();
      const { readable, writable } = new TransformStream();
      const writer = writable.getWriter();
      writer.write(
        encoder.encode(
          `error: ${JSON.stringify({
            error: getCallableJSON(e),
          })}${delimiter}`
        )
      );
      writer.write(encoder.encode('END'));
      writer.close();
      return new NextResponse(readable, {
        status: 200,
        headers: {
          'Content-Type': 'text/event-stream',
          'Cache-Control': 'no-cache',
          Connection: 'keep-alive',
          'Transfer-Encoding': 'chunked',
        },
      });
    }
    throw e;
  }
}

async function getContext<C extends ActionContext, T>(
  request: NextRequest,
  input: T,
  provider: ContextProvider<C, T> | undefined
): Promise<C> {
  // Type cast is necessary because there is no runtime way to generate a context if C is provided to appRoute
  // but contextProvider is missing. When I'm less sleepy/busy I'll see if I can make this a type error.
  const context = {} as C;
  if (!provider) {
    return context;
  }

  const r: RequestData = {
    method: request.method as RequestData['method'],
    headers: {},
    input,
  };
  request.headers.forEach((val, key) => {
    r.headers[key.toLowerCase()] = val;
  });
  return await provider(r);
}

function appRoute<
  C extends ActionContext = ActionContext,
  I extends z.ZodTypeAny = z.ZodTypeAny,
  O extends z.ZodTypeAny = z.ZodTypeAny,
  S extends z.ZodTypeAny = z.ZodTypeAny,
  Init extends z.ZodTypeAny = z.ZodTypeAny,
>(
  action: Action<I, O, S, any, Init>,
  opts?: {
    contextProvider?: ContextProvider<C, I>;
    streamManager?: StreamManager;
  }
) {
  return async (req: NextRequest): Promise<NextResponse> => {
    let context: C = {} as C;
    const { data: input, init } = (await req.json()) as {
      data: z.infer<I>;
      init?: z.infer<Init>;
    };
    const streamId = req.headers.get('x-genkit-stream-id');
    if (req.headers.get('accept') !== 'text/event-stream') {
      try {
        context = await getContext(req, input, opts?.contextProvider);
      } catch (e) {
        console.error('Error gathering context for running action:', e);
        return NextResponse.json(
          { error: getCallableJSON(e) },
          { status: getHttpStatus(e) }
        );
      }
      try {
        const resp = await action.run(input, {
          context,
          abortSignal: req.signal,
          init,
        });
        const response = NextResponse.json({ result: resp.result });
        if (opts?.streamManager && streamId) {
          response.headers.set('x-genkit-stream-id', streamId);
        }
        return response;
      } catch (e) {
        // For security reasons, log the error rather than responding with it.
        console.error('Error calling action:', e);
        return NextResponse.json(
          { error: getCallableJSON(e) },
          { status: getHttpStatus(e) }
        );
      }
    }

    try {
      context = await getContext(req, input, opts?.contextProvider);
    } catch (e) {
      console.error('Error gathering context for streaming action:', e);
      return new NextResponse(
        `error: ${JSON.stringify(getCallableJSON(e))}${delimiter}END`,
        { status: getHttpStatus(e) }
      );
    }
    const streamManager = opts?.streamManager;
    if (streamManager && streamId) {
      const response = await subscribeToStream(streamManager, streamId);
      if (response) {
        return response;
      }
    }

    const streamIdToUse = randomUUID();
    const encoder = new TextEncoder();
    const { readable, writable } = new TransformStream();

    // Not using a dangling promise causes this closure to block on the stream being drained,
    // which doesn't happen until the NextResponse is consumed later in the cosure.
    // TODO: Add ping comments at regular intervals between streaming responses to mitigate
    // timeouts.
    (async (): Promise<void> => {
      const writer = writable.getWriter();
      const taskQueue = new AsyncTaskQueue();
      let durableStream: ActionStreamInput<S, O> | undefined = undefined;
      if (streamManager) {
        durableStream = await streamManager.open(streamIdToUse);
      }
      try {
        const output = action.run(input, {
          context,
          abortSignal: req.signal,
          init,
          onChunk: (chunk) => {
            if (durableStream) {
              taskQueue.enqueue(() => durableStream!.write(chunk));
            }
            taskQueue.enqueue(() =>
              writer.write(
                encoder.encode(
                  `data: ${JSON.stringify({ message: chunk })}${delimiter}`
                )
              )
            );
          },
        });
        const finalOutput = await output;
        if (durableStream) {
          taskQueue.enqueue(() => durableStream!.done(finalOutput.result));
        }
        taskQueue.enqueue(() =>
          writer.write(
            encoder.encode(
              `data: ${JSON.stringify({ result: finalOutput.result })}${delimiter}`
            )
          )
        );
        taskQueue.enqueue(() => writer.write(encoder.encode('END')));
      } catch (err) {
        if (durableStream) {
          taskQueue.enqueue(() => durableStream!.error(err));
        }
        console.error('Error streaming action:', err);
        taskQueue.enqueue(() =>
          writer.write(
            encoder.encode(
              `error: ${JSON.stringify(getCallableJSON(err))}` + '\n\n'
            )
          )
        );
        taskQueue.enqueue(() => writer.write(encoder.encode('END')));
      } finally {
        await taskQueue.merge();
        await writer.close();
      }
    })();

    const headers = {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache',
      Connection: 'keep-alive',
      'Transfer-Encoding': 'chunked',
    };
    if (streamManager) {
      headers['x-genkit-stream-id'] = streamIdToUse;
    }

    return new NextResponse(readable, {
      status: 200,
      headers,
    });
  };
}

export default appRoute;
export { appRoute };