import { complete, stream } from "@earendil-works/pi-ai/compat";
import type { Api, AssistantMessage, AssistantMessageEvent, Context, Model, ProviderStreamOptions } from "@earendil-works/pi-ai/compat";
import type { ToolDefinition } from "@earendil-works/pi-coding-agent";

/** A complete-only injected client remains supported; production uses compat streaming. */
export type SideModelClient = {
  complete?: (model: Model<Api>, context: Context, options: ProviderStreamOptions) => Promise<AssistantMessage>;
  stream?: (model: Model<Api>, context: Context, options: ProviderStreamOptions) => AsyncIterable<AssistantMessageEvent>;
};

export type SideToolFactory = (cwd: string) => ToolDefinition<any, any>[];
export type SideProgressEvent =
  | { type: "model-request"; request: number }
  | { type: "text"; request: number; text: string }
  | { type: "tool-start"; name: string; id: string }
  | { type: "tool-end"; name: string; id: string; isError: boolean }
  | { type: "model-response"; request: number; stopReason: string };
export type SideRuntime = { modelClient: SideModelClient; toolFactory: SideToolFactory; now: () => number; onProgress?: (event: SideProgressEvent) => void };
export const defaultSideModelClient: SideModelClient = { stream, complete };
export type SideStreamClient = Required<Pick<SideModelClient, "stream">>;
export const compatSideStreamClient: SideStreamClient = { stream };

export type StreamTerminalEvent = "done" | "error" | "aborted";
/** `terminalEvent` is present only when the provider emitted a terminal event. */
export type StreamOutcome =
  | { kind: "done"; message: AssistantMessage; text: string; terminalEvent: "done" }
  | { kind: "aborted"; message?: AssistantMessage; text: string; terminalEvent?: "aborted" }
  | { kind: "error"; message?: AssistantMessage; text: string; error: string; terminalEvent?: "error" };

function textOf(message: AssistantMessage | undefined): string {
  return message?.content.filter((block): block is { type: "text"; text: string } => block.type === "text").map((block) => block.text).join("\n") ?? "";
}
function closeIterator(iterator: AsyncIterator<AssistantMessageEvent>): void {
  try {
    const closing = iterator.return?.();
    if (closing) void Promise.resolve(closing).catch(() => {});
  } catch {
    // Iterator cleanup is best effort; its rejection must not become unhandled.
  }
}
async function nextOrAbort(iterator: AsyncIterator<AssistantMessageEvent>, signal: AbortSignal): Promise<{ kind: "next"; value: IteratorResult<AssistantMessageEvent> } | { kind: "aborted" } | { kind: "rejected"; error: unknown }> {
  if (signal.aborted) return { kind: "aborted" };
  let listener: (() => void) | undefined;
  const aborted = new Promise<{ kind: "aborted" }>((resolve) => { listener = () => resolve({ kind: "aborted" }); signal.addEventListener("abort", listener, { once: true }); });
  try {
    const next = Promise.resolve(iterator.next()).then((value) => ({ kind: "next" as const, value }), (error) => ({ kind: "rejected" as const, error }));
    return await Promise.race([next, aborted]);
  } catch (error) {
    return { kind: "rejected", error };
  } finally {
    if (listener) signal.removeEventListener("abort", listener);
  }
}

/** Consume the compat event protocol, always closing its iterator without delaying cancellation. */
export async function consumeSideStream(events: AsyncIterable<AssistantMessageEvent>, signal: AbortSignal, onText?: (text: string) => void, onEvent?: (event: AssistantMessageEvent) => void): Promise<StreamOutcome> {
  const iterator = events[Symbol.asyncIterator]();
  const textBlocks = new Map<number, string>();
  let last: AssistantMessage | undefined;
  let callbackFailed = false;
  let callbackFailure: unknown;
  const visibleText = () => [...textBlocks.entries()]
    .sort(([left], [right]) => left - right)
    .map(([, text]) => text)
    .join("\n");
  const accumulatedText = () => visibleText() || textOf(last);
  const terminalText = (message: AssistantMessage) => textOf(message) || visibleText();

  try {
    while (true) {
      const next = await nextOrAbort(iterator, signal);
      if (next.kind === "aborted") return { kind: "aborted", message: last, text: accumulatedText() };
      if (next.kind === "rejected") {
        if (signal.aborted) return { kind: "aborted", message: last, text: accumulatedText() };
        return { kind: "error", message: last, text: accumulatedText(), error: next.error instanceof Error ? next.error.message : String(next.error) };
      }
      if (next.value.done) {
        if (signal.aborted) return { kind: "aborted", message: last, text: accumulatedText() };
        return { kind: "error", message: last, text: accumulatedText(), error: "stream ended without done or error" };
      }
      if (signal.aborted) return { kind: "aborted", message: last, text: accumulatedText() };

      const event = next.value.value;
      last = "partial" in event ? event.partial : event.type === "done" ? event.message : event.error;
      try {
        onEvent?.(event);
      } catch (error) {
        callbackFailed = true;
        callbackFailure = error;
        throw error;
      }

      if (event.type === "text_delta") {
        textBlocks.set(event.contentIndex, (textBlocks.get(event.contentIndex) ?? "") + event.delta);
        try {
          onText?.(visibleText());
        } catch (error) {
          callbackFailed = true;
          callbackFailure = error;
          throw error;
        }
        continue;
      }
      if (event.type === "text_end") {
        // Providers may omit deltas or correct them at end; the final block wins.
        textBlocks.set(event.contentIndex, event.content);
        try {
          onText?.(visibleText());
        } catch (error) {
          callbackFailed = true;
          callbackFailure = error;
          throw error;
        }
        continue;
      }
      if (event.type === "done") return { kind: "done", message: event.message, text: terminalText(event.message), terminalEvent: "done" };
      if (event.type === "error") {
        const text = terminalText(event.error);
        if (event.reason === "aborted") return { kind: "aborted", message: event.error, text, terminalEvent: "aborted" };
        return { kind: "error", message: event.error, text, error: event.error.errorMessage || "provider error", terminalEvent: "error" };
      }
    }
  } catch (error) {
    // Consumer callback failures are programming/UI failures, not provider errors.
    // Rethrowing lets callers settle promptly even when the provider's next() stalls.
    if (callbackFailed) throw callbackFailure;
    if (signal.aborted) return { kind: "aborted", message: last, text: accumulatedText() };
    return { kind: "error", message: last, text: accumulatedText(), error: error instanceof Error ? error.message : String(error) };
  } finally {
    closeIterator(iterator);
  }
}

export function probeSideStream(client: SideStreamClient, model: Model<Api>, context: Context, options: ProviderStreamOptions, signal: AbortSignal, onEvent: (event: AssistantMessageEvent) => void): Promise<"done" | "error" | "aborted"> {
  return consumeSideStream(client.stream(model, context, { ...options, signal }), signal, undefined, onEvent).then((outcome) => outcome.kind);
}
export async function observeSideStream(events: AsyncIterable<AssistantMessageEvent>, signal: AbortSignal, onEvent: (event: AssistantMessageEvent) => void): Promise<"done" | "error" | "aborted"> {
  const outcome = await consumeSideStream(events, signal, undefined, onEvent);
  return outcome.kind;
}