import type { ChatMessage, Chunk, StepId } from "../contracts/conversation.js";
import type { Logger, Span } from "../contracts/logging.js";
import type {
  ProviderContract,
  ProviderEvent,
  ProviderStreamOptions,
  Usage,
} from "../contracts/provider.js";
import type {
  EventEmitter,
  RetryStrategy,
  RunTurnInput,
  RunTurnResult,
} from "../contracts/runtime.js";
import type { ToolCall, ToolContract } from "../contracts/tool.js";
import { createStepDispatcher, type StepDispatcher } from "./dispatch.js";
import {
  doneEvent,
  errorEvent,
  providerRetryEvent,
  reasoningDeltaEvent,
  stepCompleteEvent,
  textDeltaEvent,
  toolCallEvent,
  toolResultEvent,
  turnStartEvent,
  usageEvent,
} from "./events.js";

/** Max steps per turn. 0 = unlimited (the loop runs until the model stops
 *  making tool calls or the abort signal fires). */
export const MAX_STEPS = 0;

function zeroUsage(): Usage {
  return { inputTokens: 0, outputTokens: 0 };
}

function addUsage(a: Usage, b: Usage): Usage {
  const inputTokens = a.inputTokens + b.inputTokens;
  const outputTokens = a.outputTokens + b.outputTokens;

  if (a.cacheReadTokens !== undefined || b.cacheReadTokens !== undefined) {
    const cacheReadTokens = (a.cacheReadTokens ?? 0) + (b.cacheReadTokens ?? 0);
    if (a.cacheWriteTokens !== undefined || b.cacheWriteTokens !== undefined) {
      return {
        inputTokens,
        outputTokens,
        cacheReadTokens,
        cacheWriteTokens: (a.cacheWriteTokens ?? 0) + (b.cacheWriteTokens ?? 0),
      };
    }
    return { inputTokens, outputTokens, cacheReadTokens };
  }

  if (a.cacheWriteTokens !== undefined || b.cacheWriteTokens !== undefined) {
    return {
      inputTokens,
      outputTokens,
      cacheWriteTokens: (a.cacheWriteTokens ?? 0) + (b.cacheWriteTokens ?? 0),
    };
  }

  return { inputTokens, outputTokens };
}

function usageAttrs(usage: Usage): Record<string, string | number | boolean | null> {
  const attrs: Record<string, string | number | boolean | null> = {
    "usage.inputTokens": usage.inputTokens,
    "usage.outputTokens": usage.outputTokens,
  };
  if (usage.cacheReadTokens !== undefined) {
    attrs["usage.cacheReadTokens"] = usage.cacheReadTokens;
  }
  if (usage.cacheWriteTokens !== undefined) {
    attrs["usage.cacheWriteTokens"] = usage.cacheWriteTokens;
  }
  return attrs;
}

function appendTextDelta(chunks: Chunk[], delta: string): void {
  const lastIdx = chunks.length - 1;
  const last = chunks[lastIdx];
  if (last !== undefined && last.type === "text") {
    chunks[lastIdx] = { type: "text", text: last.text + delta };
  } else {
    chunks.push({ type: "text", text: delta });
  }
}

function appendThinkingDelta(chunks: Chunk[], delta: string): void {
  const lastIdx = chunks.length - 1;
  const last = chunks[lastIdx];
  if (last !== undefined && last.type === "thinking") {
    chunks[lastIdx] = { type: "thinking", text: last.text + delta };
  } else {
    chunks.push({ type: "thinking", text: delta });
  }
}

/**
 * Remove tool-call chunks from an assistant message, returning a new message
 * with only the non-tool-call chunks (text, thinking, error). Returns
 * `undefined` when all chunks were tool-calls (so the caller can omit the
 * message entirely). Used when a step is aborted to avoid persisting
 * incomplete tool calls whose placeholder "Aborted" results would create
 * orphaned `tool` messages in the next turn's history.
 */
function stripToolCallChunks(msg: ChatMessage): ChatMessage | undefined {
  const stripped = msg.chunks.filter((c) => c.type !== "tool-call");
  return stripped.length > 0 ? { role: msg.role, chunks: stripped } : undefined;
}

interface StepContext {
  readonly provider: ProviderContract;
  readonly messages: ChatMessage[];
  readonly tools: readonly ToolContract[];
  readonly toolMap: Map<string, ToolContract>;
  readonly dispatch: RunTurnInput["dispatch"];
  readonly emit: EventEmitter;
  readonly signal: AbortSignal;
  readonly conversationId: string;
  readonly turnId: string;
  readonly stepId: StepId;
  readonly logger: Logger;
  readonly turnSpan: Span | undefined;
  readonly toolSpans: Map<string, Span>;
  readonly cwd: string | undefined;
  readonly computerId: string | undefined;
  readonly now: (() => number) | undefined;
  /** Per-turn provider options (model, systemPrompt, …) threaded to stream(). */
  readonly providerOpts: ProviderStreamOptions | undefined;
  /** Optional injected retry strategy (omit = no retry, backward-compatible). */
  readonly retry: RetryStrategy | undefined;
}

interface TimingState {
  ttftSpan: Span | undefined;
  decodeSpan: Span | undefined;
  firstTokenSeen: boolean;
  streamStartMs: number | undefined;
  firstTokenMs: number | undefined;
}

interface StepResult {
  readonly assistantMessage: ChatMessage | undefined;
  readonly toolCalls: ToolCall[];
  readonly toolMessages: ChatMessage[];
  readonly usage: Usage;
  readonly finishReason: string;
}

function processEvent(
  event: ProviderEvent,
  chunks: Chunk[],
  toolCalls: ToolCall[],
  dispatcher: StepDispatcher,
  ctx: StepContext,
  stepSpan: Span | undefined,
  timing: TimingState,
  toolDispatchTimes: Map<string, number>,
): void {
  switch (event.type) {
    case "text-delta":
      if (!timing.firstTokenSeen) {
        timing.firstTokenSeen = true;
        if (ctx.now !== undefined) {
          timing.firstTokenMs = ctx.now();
        }
        try {
          timing.ttftSpan?.end({ attrs: { firstToken: true } });
        } catch {
          // Swallow — D7.
        }
        timing.ttftSpan = undefined;
        try {
          timing.decodeSpan = stepSpan?.child("decode");
        } catch {
          // Swallow — D7.
        }
      }
      appendTextDelta(chunks, event.delta);
      ctx.emit(textDeltaEvent(ctx.conversationId, ctx.turnId, event.delta));
      break;
    case "reasoning-delta":
      if (!timing.firstTokenSeen) {
        timing.firstTokenSeen = true;
        if (ctx.now !== undefined) {
          timing.firstTokenMs = ctx.now();
        }
        try {
          timing.ttftSpan?.end({ attrs: { firstToken: true } });
        } catch {
          // Swallow — D7.
        }
        timing.ttftSpan = undefined;
        try {
          timing.decodeSpan = stepSpan?.child("decode");
        } catch {
          // Swallow — D7.
        }
      }
      appendThinkingDelta(chunks, event.delta);
      ctx.emit(reasoningDeltaEvent(ctx.conversationId, ctx.turnId, event.delta));
      break;
    case "tool-call": {
      const call: ToolCall = {
        id: event.toolCallId,
        name: event.toolName,
        input: event.input,
      };
      toolCalls.push(call);
      chunks.push({
        type: "tool-call",
        toolCallId: event.toolCallId,
        toolName: event.toolName,
        input: event.input,
        stepId: ctx.stepId,
      });
      ctx.emit(
        toolCallEvent(
          ctx.conversationId,
          ctx.turnId,
          ctx.stepId,
          event.toolCallId,
          event.toolName,
          event.input,
        ),
      );

      // Capture dispatch time for tool-call durationMs
      if (ctx.now !== undefined) {
        toolDispatchTimes.set(event.toolCallId, ctx.now());
      }

      // Open a tool-call span as a child of the step span (attrs: name, toolCallId)
      try {
        const tcSpan =
          stepSpan !== undefined
            ? stepSpan.child("tool-call", {
                name: event.toolName,
                toolCallId: event.toolCallId,
              })
            : ctx.logger.span("tool-call", {
                name: event.toolName,
                toolCallId: event.toolCallId,
              });
        ctx.toolSpans.set(event.toolCallId, tcSpan);
      } catch {
        // Swallow — D7: logging never breaks the turn.
      }

      if (ctx.dispatch.eager) {
        dispatcher.submit(call);
      }
      break;
    }
    case "usage":
      ctx.emit(usageEvent(ctx.conversationId, ctx.turnId, event.usage, ctx.stepId));
      break;
    case "finish":
      break;
    case "error":
      // Handled by the retry loop in executeStep (not here): an error event
      // is intercepted before processEvent so the step can decide whether to
      // retry (suppressing the error) or give up (emit it). processEvent
      // never receives an "error" event.
      break;
  }
}

async function executeStep(ctx: StepContext): Promise<StepResult> {
  const chunks: Chunk[] = [];
  const toolCalls: ToolCall[] = [];
  const toolDispatchTimes = new Map<string, number>();
  let stepUsage = zeroUsage();
  let finishReason = "stop";

  // Open a step span as a child of the turn span; capture the verbatim
  // pre-mutation prompt via a "prompt" child span whose body holds the
  // serialized messages+tools.
  let stepSpan: Span | undefined;
  try {
    stepSpan = ctx.turnSpan !== undefined ? ctx.turnSpan.child("step") : ctx.logger.span("step");
    const promptBody = JSON.stringify({ messages: ctx.messages, tools: ctx.tools });
    const promptSpan = stepSpan.child(
      "prompt",
      {
        messageCount: ctx.messages.length,
        toolCount: ctx.tools.length,
      },
      promptBody,
    );
    promptSpan.end();
  } catch {
    // Swallow — D7.
  }

  const dispatcher = createStepDispatcher(
    ctx.toolMap,
    ctx.dispatch,
    ctx.signal,
    ctx.emit,
    ctx.conversationId,
    ctx.turnId,
    ctx.toolSpans,
    ctx.cwd,
    ctx.computerId,
  );

  const timing: TimingState = {
    ttftSpan: undefined,
    decodeSpan: undefined,
    firstTokenSeen: false,
    streamStartMs: ctx.now !== undefined ? ctx.now() : undefined,
    firstTokenMs: undefined,
  };

  // Open TTFT span when spans are enabled
  try {
    if (stepSpan !== undefined) {
      timing.ttftSpan = stepSpan.child("ttft");
    }
  } catch {
    // Swallow — D7.
  }

  // Retry loop: wrap provider.stream() consumption. Retries are ONLY
  // attempted when no content was emitted yet this step (the safety
  // invariant — never duplicate partial output). On a retryable error —
  // either an EMITTED `error` ProviderEvent with `retryable === true`, OR a
  // THROWN error (retryable-by-default when pre-content) — with !hadContent:
  // ask retry.delayFor(attempt); if it returns a delay → emit a transient
  // provider-retry AgentEvent, sleep via the injected retry.sleep (abortable),
  // attempt++, re-call provider.stream(); if it returns undefined (budget
  // exhausted) → give up. Non-retryable emitted errors (retryable === false or
  // absent), errors after content, and the no-retry-configured case all fall
  // through to "give up" — identical to the pre-retry behavior.
  let hadContent = false;
  let attempt = 0;
  while (true) {
    let errored = false;
    let wasThrown = false;
    let errorMessage: string | undefined;
    let errorCode: string | undefined;
    let errorRetryable: boolean | undefined;
    let thrownErr: unknown;

    try {
      const opts: ProviderStreamOptions = {
        ...ctx.providerOpts,
        ...(ctx.turnSpan !== undefined && stepSpan !== undefined ? { logger: stepSpan.log } : {}),
      };
      const stream = ctx.provider.stream(ctx.messages, ctx.tools, opts);
      for await (const event of stream) {
        if (ctx.signal.aborted) break;
        if (event.type === "error") {
          // Intercept: hold for the retry decision — don't push a chunk
          // or emit yet (a successful retry would leave a stale error).
          errored = true;
          errorMessage = event.message;
          errorCode = event.code;
          errorRetryable = event.retryable;
          break;
        }
        if (
          event.type === "text-delta" ||
          event.type === "reasoning-delta" ||
          event.type === "tool-call" ||
          event.type === "usage"
        ) {
          hadContent = true;
        }
        processEvent(
          event,
          chunks,
          toolCalls,
          dispatcher,
          ctx,
          stepSpan,
          timing,
          toolDispatchTimes,
        );
        if (event.type === "usage") {
          stepUsage = addUsage(stepUsage, event.usage);
        }
        if (event.type === "finish") {
          finishReason = event.reason;
        }
      }
    } catch (err) {
      errored = true;
      wasThrown = true;
      errorMessage = err instanceof Error ? err.message : String(err);
      errorCode = undefined;
      errorRetryable = undefined;
      thrownErr = err;
    }

    // Abort (during stream) → stop; the runTurn loop seals aborted.
    if (ctx.signal.aborted) {
      break;
    }

    // No error → step succeeded.
    if (!errored) {
      break;
    }

    // Retryable? A thrown error is retryable-by-default when pre-content;
    // an emitted error is retryable ONLY when `retryable === true` (absent
    // or false → not retried, per the contract).
    const isRetryable = wasThrown ? true : errorRetryable === true;
    if (ctx.retry !== undefined && !hadContent && isRetryable) {
      const delay = ctx.retry.delayFor(attempt);
      if (delay !== undefined) {
        // Emit the transient provider-retry event BEFORE the sleep so the
        // UI shows "⚠ retrying in Ns…" immediately. Not persisted as a
        // chat message — it never pollutes the prompt.
        ctx.emit(
          providerRetryEvent(
            ctx.conversationId,
            ctx.turnId,
            attempt,
            delay,
            errorMessage ?? "",
            errorCode,
          ),
        );
        // Abortable sleep. If the signal fires during sleep, the shell's
        // sleep rejects — we catch it and break so the turn seals aborted.
        try {
          await ctx.retry.sleep(delay, ctx.signal);
        } catch {
          // Abort during sleep (or unexpected sleep failure).
        }
        if (ctx.signal.aborted) {
          break;
        }
        attempt++;
        continue;
      }
      // delayFor returned undefined → budget exhausted → give up.
    }

    // Give up: emit the suppressed error and end the step. This is the
    // single emission point for a terminal provider error (non-retryable,
    // post-content, budget-exhausted, or no-retry-configured).
    const message = errorMessage ?? "";
    if (errorCode !== undefined) {
      chunks.push({ type: "error", message, code: errorCode });
    } else {
      chunks.push({ type: "error", message });
    }
    ctx.emit(errorEvent(ctx.conversationId, ctx.turnId, message, errorCode));
    finishReason = "error";
    try {
      stepSpan?.end({ err: thrownErr ?? new Error(message) });
    } catch {
      // Swallow — D7.
    }
    stepSpan = undefined;
    break;
  }

  // Close timing spans: if no first token was seen, end ttft with firstToken: false
  // If decode span is open, close it
  try {
    if (timing.ttftSpan !== undefined) {
      timing.ttftSpan.end({ attrs: { firstToken: false } });
      timing.ttftSpan = undefined;
    }
    if (timing.decodeSpan !== undefined) {
      timing.decodeSpan.end();
      timing.decodeSpan = undefined;
    }
  } catch {
    // Swallow — D7.
  }

  // Emit step-complete event with timing
  const streamEndMs = ctx.now !== undefined ? ctx.now() : undefined;
  if (timing.streamStartMs !== undefined && streamEndMs !== undefined) {
    const genTotalMs = streamEndMs - timing.streamStartMs;
    const stepTiming: { ttftMs?: number; decodeMs?: number; genTotalMs?: number } = {
      genTotalMs,
    };
    if (timing.firstTokenMs !== undefined) {
      stepTiming.ttftMs = timing.firstTokenMs - timing.streamStartMs;
      stepTiming.decodeMs = streamEndMs - timing.firstTokenMs;
    }
    ctx.emit(stepCompleteEvent(ctx.conversationId, ctx.turnId, ctx.stepId, stepTiming));
  } else {
    ctx.emit(stepCompleteEvent(ctx.conversationId, ctx.turnId, ctx.stepId));
  }

  if (!ctx.dispatch.eager) {
    for (const call of toolCalls) {
      dispatcher.submit(call);
    }
  }

  const results = await dispatcher.drain();

  // Close remaining tool-call spans
  for (const call of toolCalls) {
    const tcSpan = ctx.toolSpans.get(call.id);
    if (tcSpan !== undefined) {
      const result = results.get(call.id);
      try {
        tcSpan.end({
          attrs: {
            isError: result?.isError ?? false,
            contentLength: result?.content.length ?? 0,
          },
        });
      } catch {
        // Swallow — D7.
      }
      ctx.toolSpans.delete(call.id);
    }
  }

  const toolMessages: ChatMessage[] = [];
  for (const call of toolCalls) {
    const result = results.get(call.id);
    if (result !== undefined) {
      const isError = result.isError ?? false;
      const dispatchTime = toolDispatchTimes.get(call.id);
      const toolDurationMs =
        ctx.now !== undefined && dispatchTime !== undefined ? ctx.now() - dispatchTime : undefined;
      ctx.emit(
        toolResultEvent(
          ctx.conversationId,
          ctx.turnId,
          ctx.stepId,
          call.id,
          call.name,
          result.content,
          isError,
          toolDurationMs,
        ),
      );
      toolMessages.push({
        role: "tool",
        chunks: [
          {
            type: "tool-result",
            toolCallId: call.id,
            toolName: call.name,
            content: result.content,
            isError,
            stepId: ctx.stepId,
          },
        ],
      });
    }
  }

  // Close step span (if not already closed by error)
  if (stepSpan !== undefined) {
    try {
      stepSpan.end({
        attrs: {
          finishReason,
          ...usageAttrs(stepUsage),
        },
      });
    } catch {
      // Swallow — D7.
    }
  }

  const assistantMessage: ChatMessage | undefined =
    chunks.length > 0 ? { role: "assistant", chunks } : undefined;

  return { assistantMessage, toolCalls, toolMessages, usage: stepUsage, finishReason };
}

export async function runTurn(input: RunTurnInput): Promise<RunTurnResult> {
  const messages: ChatMessage[] = [...input.messages];
  const resultMessages: ChatMessage[] = [];
  let totalUsage = zeroUsage();
  let lastStepUsage: Usage | undefined;
  let finishReason = "stop";

  const toolMap = new Map<string, ToolContract>();
  for (const tool of input.tools) {
    toolMap.set(tool.name, tool);
  }

  const conversationId = input.conversationId;
  const turnId = input.turnId;
  const signal = input.signal ?? new AbortController().signal;
  const logger = input.logger;
  const now = input.now;

  // Record turn start time for durationMs on done
  const turnStartMs = now !== undefined ? now() : undefined;

  // Open a turn span (attrs: conversationId, turnId, model)
  let turnSpan: Span | undefined;
  if (logger !== undefined) {
    try {
      turnSpan = logger.span("turn", {
        conversationId,
        turnId,
        model: input.providerOpts?.model ?? input.provider.id,
      });
    } catch {
      // Swallow — D7.
    }
  }

  // Track open tool-call spans across steps so we can close them on abort
  const toolSpans = new Map<string, Span>();

  input.emit(turnStartEvent(conversationId, turnId));

  try {
    for (let step = 0; MAX_STEPS === 0 || step < MAX_STEPS; step++) {
      if (signal.aborted) {
        finishReason = "aborted";
        break;
      }

      const stepId = `${turnId}#${step}` as StepId;

      const stepResult = await executeStep({
        provider: input.provider,
        messages,
        tools: input.tools,
        toolMap,
        dispatch: input.dispatch,
        emit: input.emit,
        signal,
        conversationId,
        turnId,
        stepId,
        logger: turnSpan?.log ?? logger ?? createNoopLogger(),
        turnSpan,
        toolSpans,
        cwd: input.cwd,
        computerId: input.computerId,
        now,
        providerOpts: input.providerOpts,
        retry: input.retry,
      });

      totalUsage = addUsage(totalUsage, stepResult.usage);
      lastStepUsage = stepResult.usage;

      // When the signal is aborted mid-step, the tool results are
      // placeholders ({ content: "Aborted", isError: true }). If these
      // are persisted and included in the next turn's message history,
      // the provider sees a `tool` role message without a preceding
      // `assistant` message carrying `tool_calls` → 400 error.
      //
      // To prevent this, when the signal is aborted we:
      //  1. Strip tool-call chunks from the assistant message (keep
      //     text/thinking/error chunks so the partial response is
      //     preserved).
      //  2. Omit tool-result messages entirely (they are not persisted,
      //     not added to resultMessages, and not passed to onStepComplete).
      //
      // This keeps the conversation history clean: the assistant's
      // partial text is preserved, but no incomplete tool calls are
      // left dangling. The `done` event still carries
      // `reason: "aborted"`, so the turn seals cleanly.
      const stepAborted = signal.aborted;
      const assistantMessage =
        stepAborted && stepResult.assistantMessage !== undefined
          ? stripToolCallChunks(stepResult.assistantMessage)
          : stepResult.assistantMessage;
      const toolMessages = stepAborted ? [] : stepResult.toolMessages;

      if (assistantMessage !== undefined) {
        messages.push(assistantMessage);
        resultMessages.push(assistantMessage);
      }

      for (const msg of toolMessages) {
        messages.push(msg);
        resultMessages.push(msg);
      }

      // Incremental persistence: notify the caller that this step's
      // messages are finalized. The caller can persist them immediately
      // (assigning seq numbers during generation). The messages are the
      // SAME objects in resultMessages — the caller must NOT double-persist.
      if (input.onStepComplete !== undefined) {
        const stepMessages: ChatMessage[] = [];
        if (assistantMessage !== undefined) {
          stepMessages.push(assistantMessage);
        }
        for (const msg of toolMessages) {
          stepMessages.push(msg);
        }
        if (stepMessages.length > 0) {
          await input.onStepComplete(stepMessages);
        }
      }

      if (stepAborted) {
        finishReason = "aborted";
        break;
      }

      if (stepResult.toolCalls.length === 0) {
        finishReason = stepResult.finishReason;
        break;
      }

      if (MAX_STEPS > 0 && step === MAX_STEPS - 1) {
        finishReason = "max-steps";
        // No next step → no tool-result boundary. Leave any pending
        // steering messages for the caller (it owns the queue).
      } else {
        // Tool-result boundary: this step produced tool calls and we are
        // about to call provider.stream again. Drain steering messages
        // and append them after the tool results, before the next call.
        // The kernel owns no queue and names no feature — it just calls
        // the callback and appends. Emits nothing (caller emits the
        // `steering` AgentEvent in its own wrapper).
        const steering = input.drainSteering?.() ?? [];
        for (const msg of steering) {
          messages.push(msg);
        }
      }
    }
  } finally {
    // Close any orphaned tool-call spans (e.g. abort mid-tool)
    for (const [id, tcSpan] of toolSpans) {
      try {
        tcSpan.end({ attrs: { orphaned: true } });
      } catch {
        // Swallow — D7.
      }
      toolSpans.delete(id);
    }

    // Close the turn span
    if (turnSpan !== undefined) {
      try {
        turnSpan.end({
          attrs: {
            finishReason,
            ...usageAttrs(totalUsage),
          },
        });
      } catch {
        // Swallow — D7.
      }
    }
  }

  const turnDurationMs =
    turnStartMs !== undefined && now !== undefined ? now() - turnStartMs : undefined;
  const hasUsage =
    totalUsage.inputTokens > 0 ||
    totalUsage.outputTokens > 0 ||
    totalUsage.cacheReadTokens !== undefined ||
    totalUsage.cacheWriteTokens !== undefined;
  const contextSize =
    hasUsage && lastStepUsage !== undefined
      ? lastStepUsage.inputTokens + lastStepUsage.outputTokens
      : undefined;
  input.emit(
    doneEvent(
      conversationId,
      turnId,
      finishReason,
      turnDurationMs,
      hasUsage ? totalUsage : undefined,
      contextSize,
    ),
  );

  return { messages: resultMessages, usage: totalUsage, finishReason };
}

function createNoopLogger(): Logger {
  return {
    debug() {},
    info() {},
    warn() {},
    error() {},
    child() {
      return createNoopLogger();
    },
    span() {
      return {
        id: "noop",
        log: createNoopLogger(),
        setAttributes() {},
        addLink() {},
        child() {
          return this;
        },
        end() {},
      };
    },
  };
}
