import { type LanguageModel, Output, streamText } from "ai";
import { z } from "zod";
import { type AiMapStreamEvent, serializeAiMapEvent } from "./ndjson";

const mappingElement = z.object({
  id: z.string().min(1),
  group: z.string().min(1),
});

export interface AiMapStreamOptions {
  model: LanguageModel;
  prompt: string;
  // Optional system message (the mapping rules). Included in the `meta` event's
  // prompt so the client's "Data sent to AI" panel shows the full prompt.
  system?: string;
  // The root ids we sent — any element the model emits with an id outside this
  // set is dropped (never trust an id the model invented).
  validIds: readonly string[];
  abortSignal?: AbortSignal;
}

// Runs the model with streamText + Output.array and returns an NDJSON streaming
// Response (see ./ndjson for the wire format). Reasoning chunks are emitted live
// via onChunk as the model thinks; each validated array element is emitted as it
// completes via elementStream. The first line is a `meta` event carrying the
// exact prompt so the client can show what was sent. Shared by the local
// on-device route and the hub cloud route so both stream identically.
export function streamAiMapNdjson({
  model,
  prompt,
  system,
  validIds,
  abortSignal,
}: AiMapStreamOptions): Response {
  const validSet = new Set(validIds);
  const encoder = new TextEncoder();

  const stream = new ReadableStream<Uint8Array>({
    async start(controller) {
      const emit = (event: AiMapStreamEvent) => {
        controller.enqueue(encoder.encode(serializeAiMapEvent(event)));
      };
      emit({
        type: "meta",
        prompt: system ? `${system}\n\n${prompt}` : prompt,
      });
      try {
        const result = streamText({
          model,
          system,
          output: Output.array({ element: mappingElement }),
          prompt,
          temperature: 0,
          maxRetries: 1,
          abortSignal,
          onChunk({ chunk }) {
            if (chunk.type === "reasoning-delta" && chunk.text) {
              emit({ type: "reasoning", text: chunk.text });
            }
          },
        });
        for await (const element of result.elementStream) {
          if (element && validSet.has(element.id) && element.group.trim()) {
            emit({
              type: "mapping",
              id: element.id,
              group: element.group.trim(),
            });
          }
        }
        emit({ type: "done" });
      } catch (error) {
        emit({
          type: "error",
          message: error instanceof Error ? error.message : String(error),
        });
      } finally {
        controller.close();
      }
    },
  });

  return new Response(stream, {
    headers: {
      "Content-Type": "application/x-ndjson; charset=utf-8",
      "Cache-Control": "no-store",
    },
  });
}