/**
 * Pydantic-AI SSE event mapper.
 *
 * Translates the event shape emitted by pydantic-AI–style Django backends
 * (text_delta / tool_call / tool_result / done / error / approval_required)
 * into the canonical `ChatStreamEvent` stream consumed by the Chat reducer.
 *
 * Backends that don't expose a stable `tool_call_id` are supported via a
 * per-stream FIFO queue keyed by tool name. The earlier "Map<name, toolId>"
 * approach lost the first toolId when a tool was invoked twice in one turn
 * (e.g. `list_tasks` for search and again for confirmation) — using a queue
 * keeps each call/result pair correctly matched.
 */

import type { ChatStreamEvent } from '../../../types';
import type { ParseSSEOptions } from '../sse';

export interface PydanticAIEvent {
  type:
    | 'text_delta'
    | 'tool_call'
    | 'tool_result'
    | 'done'
    | 'error'
    | 'approval_required';
  delta?: string;
  tool?: string;
  args?: unknown;
  result?: unknown;
  /**
   * Structured frontend payload — present on `tool_result` when the tool has
   * a `result_schema`. The LLM sees only `result` (compact text); the
   * frontend uses `data` to render rich UI (e.g. vehicle cards).
   */
  data?: unknown;
  total_tokens?: number;
  error?: string;
  tool_call_id?: string;
  session_id?: string;
}

/** Per-stream FIFO queue keyed by tool name. Created via `createToolIdQueue`. */
export interface ToolIdQueue {
  /** Allocate a new toolId for a `tool_call` event and enqueue it under `name`. */
  push(name: string): string;
  /** Pop the oldest toolId for `name` (or return an orphan marker if none). */
  shift(name: string): string;
  /** Reset all queues (e.g. on stream close). */
  clear(): void;
}

export function createToolIdQueue(): ToolIdQueue {
  const queues = new Map<string, string[]>();
  let counter = 0;
  return {
    push(name) {
      const id = `${name}-${counter++}-${Date.now()}`;
      const q = queues.get(name) ?? [];
      q.push(id);
      queues.set(name, q);
      return id;
    },
    shift(name) {
      return queues.get(name)?.shift() ?? `${name}-orphan-${counter++}`;
    },
    clear() {
      queues.clear();
    },
  };
}

/**
 * Translate a single pydantic-AI event into zero or more `ChatStreamEvent`s.
 * Pass a `ToolIdQueue` shared across the lifetime of one stream so that
 * `tool_call` / `tool_result` pairs match correctly.
 */
export function* mapPydanticAIEvent(
  ev: PydanticAIEvent,
  toolIds: ToolIdQueue,
): Generator<ChatStreamEvent> {
  switch (ev.type) {
    case 'text_delta':
      if (ev.delta) yield { type: 'chunk', delta: ev.delta };
      return;

    case 'tool_call': {
      const name = ev.tool ?? 'tool';
      const toolId = toolIds.push(name);
      yield { type: 'tool_call_start', toolId, name, input: ev.args };
      return;
    }

    case 'tool_result': {
      const name = ev.tool ?? 'tool';
      const toolId = toolIds.shift(name);
      // `data` is the structured JSON for the frontend; `result` is the LLM-facing text.
      const output: unknown = ev.data !== undefined ? ev.data : ev.result;
      yield { type: 'tool_call_end', toolId, output, status: 'success' };
      return;
    }

    case 'done':
      yield { type: 'message_end', patch: { tokensOut: ev.total_tokens } };
      return;

    case 'error':
      yield { type: 'error', code: 'backend_error', message: ev.error ?? 'Unknown error' };
      return;

    case 'approval_required':
      // Surfaced via a separate side-channel (see ResumableTransport in
      // host packages). Not translated to a ChatStreamEvent here.
      return;
  }
}

/**
 * Convenience factory: returns a `ParseSSEOptions['map']` callback that
 * decodes raw SSE frames as `PydanticAIEvent` JSON and yields zero or
 * more `ChatStreamEvent`s through `mapPydanticAIEvent`.
 *
 * Allocates an internal `ToolIdQueue` — call this once per stream.
 */
export function createPydanticAISSEMap(): NonNullable<ParseSSEOptions['map']> {
  const toolIds = createToolIdQueue();
  return (raw) => {
    if (!raw.data) return null;
    let parsed: PydanticAIEvent;
    try {
      parsed = JSON.parse(raw.data) as PydanticAIEvent;
    } catch {
      return null;
    }
    const out: ChatStreamEvent[] = [];
    for (const evt of mapPydanticAIEvent(parsed, toolIds)) {
      out.push(evt);
    }
    if (out.length === 0) return null;
    if (out.length === 1) return out[0]!;
    return out;
  };
}