export type AssistantMessageSections = {
  answer: string;
  thought: string | null;
  thoughtComplete: boolean;
};

export type ToolCall = {
  id: string;
  tool: string;
  params: Record<string, unknown>;
};

export type ContentSegment =
  | { type: "text"; content: string }
  | { type: "tool_call"; toolCall: ToolCall }
  | { type: "tool_call_pending" };

export type ParsedToolContent = {
  segments: ContentSegment[];
  toolCalls: ToolCall[];
};

const THINK_BLOCK_PATTERN = /<think>([\s\S]*?)<\/think>/gi;
const THINK_OPEN_TAG = "<think>";
const THINK_CLOSE_TAG = "</think>";

export function parseAssistantMessageSections(content: string): AssistantMessageSections {
  if (!content) {
    return { answer: "", thought: null, thoughtComplete: false };
  }

  const thoughtParts: string[] = [];
  let answer = content;

  answer = answer.replace(THINK_BLOCK_PATTERN, (_, thoughtContent: string) => {
    const trimmedThought = thoughtContent.trim();
    if (trimmedThought) {
      thoughtParts.push(trimmedThought);
    }

    return "\n";
  });

  const lastOpenIndex = answer.lastIndexOf(THINK_OPEN_TAG);
  const lastCloseIndex = answer.lastIndexOf(THINK_CLOSE_TAG);
  const hasUnclosedThought = lastOpenIndex !== -1 && lastOpenIndex > lastCloseIndex;

  if (hasUnclosedThought) {
    const streamingThought = answer.slice(lastOpenIndex + THINK_OPEN_TAG.length).trim();

    if (streamingThought) {
      thoughtParts.push(streamingThought);
    }

    answer = answer.slice(0, lastOpenIndex);
  }

  const normalizedAnswer = answer.replace(/\n{3,}/g, "\n\n").trim();
  const normalizedThought = thoughtParts.join("\n\n").trim();

  return {
    answer: normalizedAnswer,
    thought: normalizedThought || null,
    thoughtComplete: Boolean(normalizedThought) && !hasUnclosedThought
  };
}

const TOOL_CALL_OPEN_TAG = "<tool_call>";
const TOOL_CALL_PATTERN = /<tool_call>([\s\S]*?)<\/tool_call>/gi;
const PARTIAL_TOOL_TAG_TAIL = /<(?:t(?:o(?:o(?:l(?:_(?:c(?:a(?:l(?:l)?)?)?)?)?)?)?)?)?$/;

export function parseContentWithToolCalls(content: string): ParsedToolContent {
  if (!content) {
    return { segments: [], toolCalls: [] };
  }

  const segments: ContentSegment[] = [];
  const toolCalls: ToolCall[] = [];
  let lastIndex = 0;
  let toolCallIndex = 0;
  let match: RegExpExecArray | null;

  while ((match = TOOL_CALL_PATTERN.exec(content)) !== null) {
    const textBefore = content.slice(lastIndex, match.index);
    if (textBefore.trim()) {
      segments.push({ type: "text", content: textBefore.trim() });
    }

    try {
      const parsed = JSON.parse(match[1].trim()) as {
        tool?: string;
        params?: Record<string, unknown>;
      };
      const toolCall: ToolCall = {
        id: `tc-${toolCallIndex++}`,
        tool: parsed.tool ?? "unknown",
        params: parsed.params ?? {}
      };
      segments.push({ type: "tool_call", toolCall });
      toolCalls.push(toolCall);
    } catch {
      segments.push({ type: "text", content: match[0] });
    }

    lastIndex = match.index + match[0].length;
  }

  const remaining = content.slice(lastIndex);
  const unclosedIndex = remaining.lastIndexOf(TOOL_CALL_OPEN_TAG);

  if (unclosedIndex !== -1) {
    const textBefore = remaining.slice(0, unclosedIndex);
    if (textBefore.trim()) {
      segments.push({ type: "text", content: textBefore.trim() });
    }
    segments.push({ type: "tool_call_pending" });
  } else {
    const cleaned = remaining.replace(PARTIAL_TOOL_TAG_TAIL, "").trim();
    if (cleaned) {
      segments.push({ type: "text", content: cleaned });
    }
  }

  if (segments.length === 0) {
    segments.push({ type: "text", content });
  }

  return { segments, toolCalls };
}