import { createAssistantMessageEventStream, streamSimpleOpenAICompletions } from "@earendil-works/pi-ai";

function convertThinkingBlocks(message) {
  return {
    ...message,
    content: (message.content ?? []).map((block) => {
      if (block?.type !== "thinking") return block;
      return { type: "text", text: block.thinking ?? "" };
    }),
  };
}

function convertEvent(event) {
  if (event.type === "thinking_start") {
    return { type: "text_start", contentIndex: event.contentIndex, partial: convertThinkingBlocks(event.partial) };
  }
  if (event.type === "thinking_delta") {
    return { type: "text_delta", contentIndex: event.contentIndex, delta: event.delta, partial: convertThinkingBlocks(event.partial) };
  }
  if (event.type === "thinking_end") {
    return { type: "text_end", contentIndex: event.contentIndex, content: event.content, partial: convertThinkingBlocks(event.partial) };
  }
  if (event.type === "done") {
    return { ...event, message: convertThinkingBlocks(event.message) };
  }
  if (event.type === "error") {
    return { ...event, error: convertThinkingBlocks(event.error) };
  }
  if (event.partial) {
    return { ...event, partial: convertThinkingBlocks(event.partial) };
  }
  return event;
}

function streamAritmolab(model, context, options) {
  const source = streamSimpleOpenAICompletions(model, context, options);
  const stream = createAssistantMessageEventStream();

  (async () => {
    try {
      for await (const event of source) {
        stream.push(convertEvent(event));
      }
    } catch (error) {
      const message = {
        role: "assistant",
        content: [],
        api: model.api,
        provider: model.provider,
        model: model.id,
        usage: {
          input: 0,
          output: 0,
          cacheRead: 0,
          cacheWrite: 0,
          totalTokens: 0,
          cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
        },
        stopReason: options?.signal?.aborted ? "aborted" : "error",
        errorMessage: error instanceof Error ? error.message : String(error),
        timestamp: Date.now(),
      };
      stream.push({ type: "error", reason: message.stopReason, error: message });
      stream.end();
    }
  })();

  return stream;
}

export default function (pi) {
  pi.registerProvider("aritmolab", {
    name: "AritmoLab",
    baseUrl: "https://ml-aritmolab.policlinicosandonato.it/v1",
    apiKey: "aritmolab",
    api: "openai-completions",
    streamSimple: streamAritmolab,
    models: [
      {
        id: "qwen3.6-35b-a3b",
        name: "AritmoLab Qwen3.6 35B A3B",
        reasoning: false,
        input: ["text"],
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
        contextWindow: 131072,
        maxTokens: 16384,
        compat: {
          supportsDeveloperRole: false,
          supportsReasoningEffort: false,
          supportsStore: false,
          maxTokensField: "max_tokens",
        },
      },
    ],
  });
}