import type { Message, ToolCallEvent } from "@/lib/api";

export const STREAMING_ASSISTANT_METADATA_KIND = "streaming_assistant";

export type StreamingAssistantMetadata = {
  kind: typeof STREAMING_ASSISTANT_METADATA_KIND;
  ephemeralTrace: string | null;
  traceRevision: number;
};

export type StreamingAttemptState = {
  segmentText: string;
  initiatedToolCallIds: string[];
};

export type StreamingToolCallTransition = {
  attempt: StreamingAttemptState;
  startedNewToolCall: boolean;
  traceText: string | null;
};

export function createStreamingAssistantMetadata(): StreamingAssistantMetadata {
  return {
    kind: STREAMING_ASSISTANT_METADATA_KIND,
    ephemeralTrace: null,
    traceRevision: 0,
  };
}

export function asStreamingAssistantMetadata(value: unknown): StreamingAssistantMetadata | null {
  if (!value || typeof value !== "object" || Array.isArray(value)) return null;
  const record = value as Record<string, unknown>;
  if (record.kind !== STREAMING_ASSISTANT_METADATA_KIND) return null;

  return {
    kind: STREAMING_ASSISTANT_METADATA_KIND,
    ephemeralTrace: typeof record.ephemeralTrace === "string" ? record.ephemeralTrace : null,
    traceRevision:
      typeof record.traceRevision === "number" && Number.isFinite(record.traceRevision)
        ? Math.max(0, Math.floor(record.traceRevision))
        : 0,
  };
}

export function createStreamingAttemptState(): StreamingAttemptState {
  return {
    segmentText: "",
    initiatedToolCallIds: [],
  };
}

export function appendStreamingDelta(attempt: StreamingAttemptState, delta: string): StreamingAttemptState {
  if (!delta) return attempt;
  return {
    ...attempt,
    segmentText: attempt.segmentText + delta,
  };
}

export function transitionStreamingToolCall(
  attempt: StreamingAttemptState,
  event: Pick<ToolCallEvent, "toolCallId" | "status">
): StreamingToolCallTransition {
  if (event.status !== "initiated" || attempt.initiatedToolCallIds.includes(event.toolCallId)) {
    return { attempt, startedNewToolCall: false, traceText: null };
  }

  const traceText = attempt.segmentText.trim();
  return {
    attempt: {
      segmentText: "",
      initiatedToolCallIds: attempt.initiatedToolCallIds.concat(event.toolCallId),
    },
    startedNewToolCall: true,
    traceText: traceText || null,
  };
}

function updateStreamingAssistant(
  messages: Message[],
  assistantMessagePrefix: string,
  update: (message: Message, metadata: StreamingAssistantMetadata) => Message
) {
  let didUpdate = false;
  const nextMessages = messages.map((message, index, all) => {
    const isTarget = index === all.length - 1 && message.id.startsWith(assistantMessagePrefix);
    if (!isTarget) return message;
    didUpdate = true;
    return update(message, asStreamingAssistantMetadata(message.metadata) ?? createStreamingAssistantMetadata());
  });
  return didUpdate ? nextMessages : messages;
}

export function setStreamingAssistantSegment(messages: Message[], assistantMessagePrefix: string, segmentText: string) {
  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
    ...message,
    content: segmentText,
    metadata: {
      ...metadata,
      ephemeralTrace: null,
    } satisfies StreamingAssistantMetadata,
  }));
}

export function promoteStreamingAssistantTrace(messages: Message[], assistantMessagePrefix: string, traceText: string) {
  const normalizedTrace = traceText.trim();
  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
    ...message,
    content: "",
    metadata: normalizedTrace
      ? ({
          ...metadata,
          ephemeralTrace: normalizedTrace,
          traceRevision: metadata.traceRevision + 1,
        } satisfies StreamingAssistantMetadata)
      : metadata,
  }));
}

export function finalizeStreamingAssistant(messages: Message[], assistantMessagePrefix: string, finalText: string) {
  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
    ...message,
    content: finalText,
    metadata: {
      ...metadata,
      ephemeralTrace: null,
    } satisfies StreamingAssistantMetadata,
  }));
}

export function clearStreamingAssistantTrace(messages: Message[], assistantMessagePrefix: string) {
  return updateStreamingAssistant(messages, assistantMessagePrefix, (message, metadata) => ({
    ...message,
    metadata: {
      ...metadata,
      ephemeralTrace: null,
    } satisfies StreamingAssistantMetadata,
  }));
}