/**
 * In-memory chat transport for stories and tests. Replays scripted replies.
 */

import type {
  ChatMessage,
  ChatStreamEvent,
  ChatTransport,
  CreateSessionOptions,
  HistoryPage,
  SendOptions,
  SessionInfo,
  StreamOptions,
} from '../../types';
import { createId } from '../ids';

export interface MockTransportOptions {
  /** Each entry is the assistant's reply for one user turn. Strings are split
   *  into chunks; arrays are taken as the exact event sequence (after a
   *  prepended `message_start` and before a synthetic `message_end`). */
  replies?: Array<string | ChatStreamEvent[]>;
  latencyMs?: number;
  /** Initial history returned by `createSession`. */
  initialMessages?: ChatMessage[];
  shouldFail?: (attempt: number) => boolean;
}

const DEFAULT_REPLY = 'Hi there!';

function splitForStream(text: string, parts = 4): string[] {
  if (!text) return [];
  const chunkSize = Math.max(1, Math.ceil(text.length / parts));
  const out: string[] = [];
  for (let i = 0; i < text.length; i += chunkSize) {
    out.push(text.slice(i, i + chunkSize));
  }
  return out;
}

const sleep = (ms: number) => new Promise((r) => setTimeout(r, ms));

export function createMockTransport(opts: MockTransportOptions = {}): ChatTransport {
  const replies = opts.replies?.length ? opts.replies : [DEFAULT_REPLY];
  const latency = opts.latencyMs ?? 30;
  const history: ChatMessage[] = [...(opts.initialMessages ?? [])];
  let turn = 0;
  let attempt = 0;

  return {
    async createSession(_opts?: CreateSessionOptions): Promise<SessionInfo> {
      await sleep(latency);
      return {
        sessionId: createId('s'),
        messages: history.length ? [...history] : undefined,
        hasMore: false,
        cursor: null,
        resumed: history.length > 0,
      };
    },

    async loadHistory(_sid, _cursor, _limit): Promise<HistoryPage> {
      await sleep(latency);
      return { messages: [], hasMore: false, nextCursor: null };
    },

    async *stream(
      _sid: string,
      content: string,
      options: StreamOptions,
    ): AsyncGenerator<ChatStreamEvent, void, void> {
      attempt += 1;
      if (opts.shouldFail?.(attempt)) {
        throw new Error('mock transport scripted failure');
      }

      // Record the user message for any subsequent loadHistory.
      history.push({
        id: createId('u'),
        role: 'user',
        content,
        createdAt: Date.now(),
      });

      const messageId = createId('a');
      yield { type: 'message_start', messageId, sessionId: _sid };

      const reply = replies[turn % replies.length];
      turn += 1;

      if (typeof reply === 'string') {
        for (const piece of splitForStream(reply)) {
          if (options.signal.aborted) return;
          await sleep(latency);
          yield { type: 'chunk', delta: piece };
        }
        yield {
          type: 'message_end',
          patch: { tokensIn: content.length, tokensOut: reply.length },
        };
      } else {
        for (const ev of reply) {
          if (options.signal.aborted) return;
          await sleep(latency);
          yield ev;
        }
        // If the script didn't end the message itself, do it for them.
        const lastType = reply[reply.length - 1]?.type;
        if (lastType !== 'message_end' && lastType !== 'error') {
          yield { type: 'message_end' };
        }
      }
    },

    async send(_sid, content, _options?: SendOptions): Promise<ChatMessage> {
      await sleep(latency);
      const reply = replies[turn % replies.length];
      turn += 1;
      const text =
        typeof reply === 'string'
          ? reply
          : reply
              .filter((e) => e.type === 'chunk')
              .map((e) => (e as { delta: string }).delta)
              .join('');
      return {
        id: createId('a'),
        role: 'assistant',
        content: text || DEFAULT_REPLY,
        createdAt: Date.now(),
      };
    },

    async closeSession(_sid: string): Promise<void> {
      // no-op
    },
  };
}