import { readStream } from "@/app/utilities/readStream";

export async function stream(model: string, prompt: string, system: string): Promise<ReadableStream<any>> {
  return await buildStream(model, prompt, system);
}

function parseSequentialJSON(jsonString: string): { response: string }[] {
  // Add a comma between the }{ to transform it into a valid JSON array
  const validJsonString = `[${jsonString.replace(/}\s*{/g, '},{')}]`;

  try {
    // Parse the valid JSON string
    const jsonArray = JSON.parse(validJsonString);
    return jsonArray;
  } catch (error) {
    throw new Error("Failed to parse.")
  }
}

async function buildStream(model: string, prompt: string, system: string): Promise<ReadableStream<any>> {
  const encoder = new TextEncoder();
  return new ReadableStream({
    async start(controller) {
      try {
        const { OLLAMA_HOST } = process.env;
        console.log(OLLAMA_HOST);
        const endpoint = '/api/generate'; // API endpoint

        const url = new URL(endpoint, OLLAMA_HOST).toString();

        console.log(url);

        const requestBody = {
          prompt,
          system,
          model,
          stream: true,
        };

        const response = await fetch(url, {
          method: 'POST',
          body: JSON.stringify(requestBody),
          headers: {
            'Content-Type': 'application/json'
          }
        });

        console.log(`Response ok: ${response.ok}`);

        if (!response.ok) {
          throw new Error(`HTTP error! Status: ${response.status}`);
        }

        console.log("Response is OK");

        // Handle the stream
        const reader = response.body;
        if (reader == null) {
          throw new Error("Reader is null.");
        }
        let buffer = '';
        await readStream(
          reader,
          value => {
            try {
              buffer += value;
              const next = parseSequentialJSON(buffer).map(element => element.response).join("");
              controller.enqueue(encoder.encode(next));
              buffer = '';
            } catch { }
          }
        );

        controller.close();
      } catch (e) {
        console.log('An error occurred while communicating with ollama.');
        throw e;
      }
    },
  });
}