/**
 * Maps NDJSON events from the Rust agent binary to AG-UI protocol events.
 */

import type { AguiEvent } from "../types/agui-types";
import type { NdjsonEvent } from "./rust-agent-adapter";

export class RustEventMapper {
  hasEmittedText = false;
  private messageCounter = 0;
  private streamingMessageId: string | null = null;

  private nextMessageId(): string {
    this.messageCounter++;
    return `msg_rust_${Date.now()}_${this.messageCounter}`;
  }

  *mapEvent(event: NdjsonEvent): Generator<AguiEvent> {
    switch (event.type) {
      case "started":
        console.log(`[rust-agent] started with tools: [${event.tools.join(",")}]`);
        break;

      case "thinking":
        yield {
          type: "CUSTOM",
          name: "thinking",
          value: { content: event.content, iteration: event.iteration },
        };
        break;

      case "tool_call": {
        // Close any open streaming text message before tool calls
        yield* this.closeStreamingMessage();

        yield {
          type: "TOOL_CALL_START",
          toolCallId: event.tool_call_id,
          toolCallName: event.tool_name,
          parentMessageId: "",
        };
        if (event.arguments) {
          yield {
            type: "TOOL_CALL_ARGS",
            toolCallId: event.tool_call_id,
            delta: event.arguments,
          };
        }
        yield {
          type: "TOOL_CALL_END",
          toolCallId: event.tool_call_id,
        };
        break;
      }

      case "tool_result": {
        const resultMsgId = this.nextMessageId();
        yield {
          type: "TOOL_CALL_RESULT",
          toolCallId: event.tool_call_id,
          messageId: resultMsgId,
          content: event.content,
          role: "tool",
        };
        break;
      }

      case "text_delta": {
        if (event.content) {
          // Open a streaming message if not already open
          if (!this.streamingMessageId) {
            this.streamingMessageId = this.nextMessageId();
            yield { type: "TEXT_MESSAGE_START", messageId: this.streamingMessageId, role: "assistant" };
          }
          yield { type: "TEXT_MESSAGE_CONTENT", messageId: this.streamingMessageId, delta: event.content };
          this.hasEmittedText = true;
        }
        break;
      }

      case "completed": {
        // Close any open streaming message
        yield* this.closeStreamingMessage();

        // Use final_content as fallback if no text_delta was emitted
        if (!this.hasEmittedText && event.final_content && event.final_content.trim()) {
          const msgId = this.nextMessageId();
          yield { type: "TEXT_MESSAGE_START", messageId: msgId, role: "assistant" };
          yield { type: "TEXT_MESSAGE_CONTENT", messageId: msgId, delta: event.final_content };
          yield { type: "TEXT_MESSAGE_END", messageId: msgId };
          this.hasEmittedText = true;
        }
        console.log(
          `[rust-agent] completed: tokens=${event.tokens_used} iterations=${event.iterations} tool_calls=${event.tool_calls}`,
        );
        break;
      }

      case "failed": {
        yield* this.closeStreamingMessage();

        const msgId = this.nextMessageId();
        yield { type: "TEXT_MESSAGE_START", messageId: msgId, role: "assistant" };
        yield { type: "TEXT_MESSAGE_CONTENT", messageId: msgId, delta: `Agent error: ${event.error}` };
        yield { type: "TEXT_MESSAGE_END", messageId: msgId };
        this.hasEmittedText = true;
        yield { type: "RUN_ERROR", message: event.error };
        break;
      }
    }
  }

  private *closeStreamingMessage(): Generator<AguiEvent> {
    if (this.streamingMessageId) {
      yield { type: "TEXT_MESSAGE_END", messageId: this.streamingMessageId };
      this.streamingMessageId = null;
    }
  }
}