import type { Chat, GeneratorController } from "@lmstudio/sdk";
import {
  type ChatCompletionMessageToolCall,
  type ChatCompletionTool,
  type ChatCompletionToolMessageParam,
} from "openai/resources/index.js";
import type { OpenAICompatAssistantMessage, OpenAICompatMessageParam } from "../types";
import {
  REASONING_CONTENT_KEY,
  REASONING_END_TAG,
  REASONING_START_TAG,
  TURN_SUMMARY_MARKER,
} from "../constants";

/* -------------------------------------------------------------------------- */
/*                             Small builders                                 */
/* -------------------------------------------------------------------------- */

/** Parse a streamed tool-call argument string into an object (or null if malformed). */
export function parseToolCallArguments(rawArguments: string): Record<string, unknown> | null {
  if (rawArguments.trim().length === 0) return {};
  try {
    const parsed = JSON.parse(rawArguments);
    return typeof parsed === "object" && parsed !== null ? (parsed as Record<string, unknown>) : null;
  } catch {
    return null;
  }
}

/** Strip the in-chat turn-summary tail from assistant text before resending to the model. */
export function stripTurnSummary(text: string): string {
  const idx = text.lastIndexOf(TURN_SUMMARY_MARKER);
  return idx === -1 ? text : text.slice(0, idx);
}

/* -------------------------------------------------------------------------- */
/*                          Reasoning block handling                          */
/* -------------------------------------------------------------------------- */

/** Split text into visible content and concatenated reasoning segments delimited by tags. */
export function extractReasoningBlocks(text: string): {
  content: string;
  reasoningContent: string | undefined;
} {
  const reasoningParts: string[] = [];
  let visibleContent = "";
  let cursor = 0;

  while (cursor < text.length) {
    const startIndex = text.indexOf(REASONING_START_TAG, cursor);
    if (startIndex === -1) {
      visibleContent += text.slice(cursor);
      break;
    }

    visibleContent += text.slice(cursor, startIndex);
    const reasoningStartIndex = startIndex + REASONING_START_TAG.length;
    const endIndex = text.indexOf(REASONING_END_TAG, reasoningStartIndex);
    if (endIndex === -1) {
      reasoningParts.push(text.slice(reasoningStartIndex));
      cursor = text.length;
      break;
    }

    reasoningParts.push(text.slice(reasoningStartIndex, endIndex));
    cursor = endIndex + REASONING_END_TAG.length;
  }

  const reasoningContent = reasoningParts.join("");
  return {
    content: visibleContent,
    reasoningContent: reasoningContent.length > 0 ? reasoningContent : undefined,
  };
}

/**
 * Normalize heterogeneous reasoning payloads (string | array | object) into a plain string.
 * Recurses through wrapper shapes until a usable text fragment is found.
 */
export function normalizeReasoningDelta(value: unknown, preferredKey?: string | null): string | null {
  if (typeof value === "string") return value;

  if (Array.isArray(value)) {
    const combined = value
      .map((item) => normalizeReasoningDelta(item, preferredKey))
      .filter((item): item is string => item !== null && item.length > 0)
      .join("");
    return combined.length > 0 ? combined : null;
  }

  if (value === null || typeof value !== "object") return null;

  const record = value as Record<string, unknown>;
  const candidateKeys = [preferredKey, "text", "content", "reasoning_content", "delta", "parts"].filter(
    (key): key is string => typeof key === "string" && key.length > 0,
  );
  for (const key of candidateKeys) {
    const normalized = normalizeReasoningDelta(record[key], preferredKey);
    if (normalized !== null && normalized.length > 0) return normalized;
  }

  return null;
}

/* -------------------------------------------------------------------------- */
/*                             History conversion                             */
/* -------------------------------------------------------------------------- */

/** Convert an LM Studio chat history into OpenAI-compatible wire messages. */
export async function toOpenAIMessages(history: Chat): Promise<OpenAICompatMessageParam[]> {
  const messages: OpenAICompatMessageParam[] = [];

  for (const message of history) {
    switch (message.getRole()) {
      case "system":
        messages.push({ role: "system", content: message.getText() });
        break;

      case "user":
        messages.push({ role: "user", content: message.getText() });
        break;

      case "assistant": {
        const assistantText = stripTurnSummary(message.getText());
        const { content, reasoningContent } = extractReasoningBlocks(assistantText);

        const toolCalls: ChatCompletionMessageToolCall[] = message
          .getToolCallRequests()
          .map((toolCall) => ({
            id: toolCall.id ?? "",
            type: "function" as const,
            function: {
              name: toolCall.name,
              arguments: JSON.stringify(toolCall.arguments ?? {}),
            },
          }));

        const assistantMessage: OpenAICompatAssistantMessage = {
          role: "assistant",
          content,
          ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
        };

        if (reasoningContent) {
          assistantMessage[REASONING_CONTENT_KEY] = reasoningContent;
        }

        messages.push(assistantMessage);
        break;
      }

      case "tool": {
        message.getToolCallResults().forEach((toolCallResult) => {
          messages.push({
            role: "tool",
            tool_call_id: toolCallResult.toolCallId ?? "",
            content: toolCallResult.content,
          } as ChatCompletionToolMessageParam);
        });
        break;
      }
    }
  }

  return messages;
}

/** Project available controller tools into the OpenAI tool schema (undefined when none). */
export function toOpenAITools(ctl: GeneratorController): ChatCompletionTool[] | undefined {
  const tools = ctl.getToolDefinitions().map<ChatCompletionTool>((t) => ({
    type: "function",
    function: {
      name: t.function.name,
      description: t.function.description,
      parameters: t.function.parameters ?? {},
    },
  }));
  return tools.length ? tools : undefined;
}

/* -------------------------------------------------------------------------- */
/*                          Ordering / coherence pass                         */
/* -------------------------------------------------------------------------- */

/** Loose structural view used only inside the sanitizing pass. */
type SanitizeMsg = {
  role: string;
  content?: string;
  tool_calls?: Array<{ id?: string }>;
  tool_call_id?: string;
  [key: string]: unknown;
};

/**
 * Defensive coherence pass over the assembled wire messages.
 *
 * Ensures every assistant message advertising `tool_calls` is immediately followed by the
 * matching `tool` result messages, and that no `tool` message is stranded ahead of its caller.
 *
 * Intentional edge behaviour: if an assistant message advertises `tool_calls` whose ids have NO
 * corresponding `tool` results (e.g. results were pruned), the dangling `tool_calls` are stripped.
 * If non-blank textual content remains it is kept as a plain assistant message; otherwise the
 * message is dropped entirely. This stops providers from rejecting malformed tool sequences.
 */
export function sanitizeToolMessageOrdering(messages: OpenAICompatMessageParam[]): OpenAICompatMessageParam[] {
  const src = messages as SanitizeMsg[];
  const result: SanitizeMsg[] = [];

  for (let i = 0; i < src.length; i++) {
    const message = src[i];

    if (message.role === "tool") continue;

    const toolCalls = Array.isArray(message.tool_calls) ? message.tool_calls : [];

    if (message.role !== "assistant" || toolCalls.length === 0) {
      result.push(message);
      continue;
    }

    const expectedIds = new Set(
      toolCalls.map((tc) => tc?.id).filter((id): id is string => typeof id === "string" && id.length > 0),
    );

    const toolResults: SanitizeMsg[] = [];
    let j = i + 1;
    while (j < src.length && src[j].role === "tool") {
      const toolMessage = src[j];
      if (toolMessage.tool_call_id != null && expectedIds.has(toolMessage.tool_call_id)) {
        toolResults.push(toolMessage);
        expectedIds.delete(toolMessage.tool_call_id);
      }
      j++;
    }

    if (expectedIds.size === 0 && toolResults.length > 0) {
      result.push(message);
      result.push(...toolResults);
      i = j - 1;
      continue;
    }

    const copy: SanitizeMsg = { ...message };
    delete copy.tool_calls;
    if (typeof copy.content === "string" && copy.content.trim().length > 0) {
      result.push(copy);
    }
    i = j - 1;
  }

  return result as OpenAICompatMessageParam[];
}


/* -------------------------------------------------------------------------- */
/*                       Conversation-thread keying                           */
/*                                                                            */
/* Derives a stable key for the conversation a generate() call belongs to.    */
/* Used to scope the persistent turn-stat accumulator across the multi-round  */
/* (tool handoff) regenerations of a single turn.                             */
/* -------------------------------------------------------------------------- */

/** Fast, dependency-free 53-bit string fingerprint (adequate for keying, not for security). */
function cyrb53(str: string, seed = 0): string {
  let h1 = 0xdeadbeef ^ seed;
  let h2 = 0x41c6ce57 ^ seed;
  for (let i = 0; i < str.length; i++) {
    const ch = str.charCodeAt(i);
    h1 = Math.imul(h1 ^ ch, 2654435761);
    h2 = Math.imul(h2 ^ ch, 1597334677);
  }
  h1 = Math.imul(h1 ^ (h1 >>> 16), 2246822507) ^ Math.imul(h2 ^ (h2 >>> 13), 3266489909);
  h2 = Math.imul(h2 ^ (h2 >>> 16), 2246822507) ^ Math.imul(h1 ^ (h1 >>> 13), 3266489909);
  return (4294967296 * (2097151 & h2) + (h1 >>> 0)).toString(36);
}

/**
 * Stable identifier for the conversation thread backing `history`.
 *
 * Built from the IMMUTABLE origin (system prompt + first user message) so it stays CONSTANT
 * while the transcript grows across a turn's rounds, yet differs between distinct
 * conversations. The growing tail is intentionally excluded so successive generate() calls of
 * the same turn land in the same bucket. Two threads opened with byte-identical origins are
 * indistinguishable from history alone and thus share a bucket (benign edge case).
 */
export function deriveConversationKey(history: Chat): string {
  const system = history.getSystemPrompt();
  let firstUser = "";
  for (const message of history) {
    if (message.getRole() === "user") {
      firstUser = message.getText();
      break;
    }
  }
  return cyrb53(`${system}\u0000${firstUser}`);
}

import type { Chat, GeneratorController } from "@lmstudio/sdk";
import {
  type ChatCompletionMessageToolCall,
  type ChatCompletionTool,
  type ChatCompletionToolMessageParam,
} from "openai/resources/index.js";
import type { OpenAICompatAssistantMessage, OpenAICompatMessageParam } from "../types";
import {
  REASONING_CONTENT_KEY,
  REASONING_END_TAG,
  REASONING_START_TAG,
  TURN_SUMMARY_MARKER,
} from "../constants";

/* -------------------------------------------------------------------------- */
/*                             Small builders                                 */
/* -------------------------------------------------------------------------- */

/** Parse a streamed tool-call argument string into an object (or null if malformed). */
export function parseToolCallArguments(rawArguments: string): Record<string, unknown> | null {
  if (rawArguments.trim().length === 0) return {};
  try {
    const parsed = JSON.parse(rawArguments);
    return typeof parsed === "object" && parsed !== null ? (parsed as Record<string, unknown>) : null;
  } catch {
    return null;
  }
}

/** Strip the in-chat turn-summary tail from assistant text before resending to the model. */
export function stripTurnSummary(text: string): string {
  const idx = text.lastIndexOf(TURN_SUMMARY_MARKER);
  return idx === -1 ? text : text.slice(0, idx);
}

/* -------------------------------------------------------------------------- */
/*                          Reasoning block handling                          */
/* -------------------------------------------------------------------------- */

/** Split text into visible content and concatenated reasoning segments delimited by tags. */
export function extractReasoningBlocks(text: string): {
  content: string;
  reasoningContent: string | undefined;
} {
  const reasoningParts: string[] = [];
  let visibleContent = "";
  let cursor = 0;

  while (cursor < text.length) {
    const startIndex = text.indexOf(REASONING_START_TAG, cursor);
    if (startIndex === -1) {
      visibleContent += text.slice(cursor);
      break;
    }

    visibleContent += text.slice(cursor, startIndex);
    const reasoningStartIndex = startIndex + REASONING_START_TAG.length;
    const endIndex = text.indexOf(REASONING_END_TAG, reasoningStartIndex);
    if (endIndex === -1) {
      reasoningParts.push(text.slice(reasoningStartIndex));
      cursor = text.length;
      break;
    }

    reasoningParts.push(text.slice(reasoningStartIndex, endIndex));
    cursor = endIndex + REASONING_END_TAG.length;
  }

  const reasoningContent = reasoningParts.join("");
  return {
    content: visibleContent,
    reasoningContent: reasoningContent.length > 0 ? reasoningContent : undefined,
  };
}

/**
 * Normalize heterogeneous reasoning payloads (string | array | object) into a plain string.
 * Recurses through wrapper shapes until a usable text fragment is found.
 */
export function normalizeReasoningDelta(value: unknown, preferredKey?: string | null): string | null {
  if (typeof value === "string") return value;

  if (Array.isArray(value)) {
    const combined = value
      .map((item) => normalizeReasoningDelta(item, preferredKey))
      .filter((item): item is string => item !== null && item.length > 0)
      .join("");
    return combined.length > 0 ? combined : null;
  }

  if (value === null || typeof value !== "object") return null;

  const record = value as Record<string, unknown>;
  const candidateKeys = [preferredKey, "text", "content", "reasoning_content", "delta", "parts"].filter(
    (key): key is string => typeof key === "string" && key.length > 0,
  );
  for (const key of candidateKeys) {
    const normalized = normalizeReasoningDelta(record[key], preferredKey);
    if (normalized !== null && normalized.length > 0) return normalized;
  }

  return null;
}

/* -------------------------------------------------------------------------- */
/*                             History conversion                             */
/* -------------------------------------------------------------------------- */

/** Convert an LM Studio chat history into OpenAI-compatible wire messages. */
export async function toOpenAIMessages(history: Chat): Promise<OpenAICompatMessageParam[]> {
  const messages: OpenAICompatMessageParam[] = [];

  for (const message of history) {
    switch (message.getRole()) {
      case "system":
        messages.push({ role: "system", content: message.getText() });
        break;

      case "user":
        messages.push({ role: "user", content: message.getText() });
        break;

      case "assistant": {
        const assistantText = stripTurnSummary(message.getText());
        const { content, reasoningContent } = extractReasoningBlocks(assistantText);

        const toolCalls: ChatCompletionMessageToolCall[] = message
          .getToolCallRequests()
          .map((toolCall) => ({
            id: toolCall.id ?? "",
            type: "function" as const,
            function: {
              name: toolCall.name,
              arguments: JSON.stringify(toolCall.arguments ?? {}),
            },
          }));

        const assistantMessage: OpenAICompatAssistantMessage = {
          role: "assistant",
          content,
          ...(toolCalls.length ? { tool_calls: toolCalls } : {}),
        };

        if (reasoningContent) {
          assistantMessage[REASONING_CONTENT_KEY] = reasoningContent;
        }

        messages.push(assistantMessage);
        break;
      }

      case "tool": {
        message.getToolCallResults().forEach((toolCallResult) => {
          messages.push({
            role: "tool",
            tool_call_id: toolCallResult.toolCallId ?? "",
            content: toolCallResult.content,
          } as ChatCompletionToolMessageParam);
        });
        break;
      }
    }
  }

  return messages;
}

/** Project available controller tools into the OpenAI tool schema (undefined when none). */
export function toOpenAITools(ctl: GeneratorController): ChatCompletionTool[] | undefined {
  const tools = ctl.getToolDefinitions().map<ChatCompletionTool>((t) => ({
    type: "function",
    function: {
      name: t.function.name,
      description: t.function.description,
      parameters: t.function.parameters ?? {},
    },
  }));
  return tools.length ? tools : undefined;
}

/* -------------------------------------------------------------------------- */
/*                          Ordering / coherence pass                         */
/* -------------------------------------------------------------------------- */

/** Loose structural view used only inside the sanitizing pass. */
type SanitizeMsg = {
  role: string;
  content?: string;
  tool_calls?: Array<{ id?: string }>;
  tool_call_id?: string;
  [key: string]: unknown;
};

/**
 * Defensive coherence pass over the assembled wire messages.
 *
 * Ensures every assistant message advertising `tool_calls` is immediately followed by the
 * matching `tool` result messages, and that no `tool` message is stranded ahead of its caller.
 *
 * Intentional edge behaviour: if an assistant message advertises `tool_calls` whose ids have NO
 * corresponding `tool` results (e.g. results were pruned), the dangling `tool_calls` are stripped.
 * If non-blank textual content remains it is kept as a plain assistant message; otherwise the
 * message is dropped entirely. This stops providers from rejecting malformed tool sequences.
 */
export function sanitizeToolMessageOrdering(messages: OpenAICompatMessageParam[]): OpenAICompatMessageParam[] {
  const src = messages as SanitizeMsg[];
  const result: SanitizeMsg[] = [];

  for (let i = 0; i < src.length; i++) {
    const message = src[i];

    if (message.role === "tool") continue;

    const toolCalls = Array.isArray(message.tool_calls) ? message.tool_calls : [];

    if (message.role !== "assistant" || toolCalls.length === 0) {
      result.push(message);
      continue;
    }

    const expectedIds = new Set(
      toolCalls.map((tc) => tc?.id).filter((id): id is string => typeof id === "string" && id.length > 0),
    );

    const toolResults: SanitizeMsg[] = [];
    let j = i + 1;
    while (j < src.length && src[j].role === "tool") {
      const toolMessage = src[j];
      if (toolMessage.tool_call_id != null && expectedIds.has(toolMessage.tool_call_id)) {
        toolResults.push(toolMessage);
        expectedIds.delete(toolMessage.tool_call_id);
      }
      j++;
    }

    if (expectedIds.size === 0 && toolResults.length > 0) {
      result.push(message);
      result.push(...toolResults);
      i = j - 1;
      continue;
    }

    const copy: SanitizeMsg = { ...message };
    delete copy.tool_calls;
    if (typeof copy.content === "string" && copy.content.trim().length > 0) {
      result.push(copy);
    }
    i = j - 1;
  }

  return result as OpenAICompatMessageParam[];
}


/* -------------------------------------------------------------------------- */
/*                       Conversation-thread keying                           */
/*                                                                            */
/* Derives a stable key for the conversation a generate() call belongs to.    */
/* Used to scope the persistent turn-stat accumulator across the multi-round  */
/* (tool handoff) regenerations of a single turn.                             */
/* -------------------------------------------------------------------------- */

/** Fast, dependency-free 53-bit string fingerprint (adequate for keying, not for security). */
function cyrb53(str: string, seed = 0): string {
  let h1 = 0xdeadbeef ^ seed;
  let h2 = 0x41c6ce57 ^ seed;
  for (let i = 0; i < str.length; i++) {
    const ch = str.charCodeAt(i);
    h1 = Math.imul(h1 ^ ch, 2654435761);
    h2 = Math.imul(h2 ^ ch, 1597334677);
  }
  h1 = Math.imul(h1 ^ (h1 >>> 16), 2246822507) ^ Math.imul(h2 ^ (h2 >>> 13), 3266489909);
  h2 = Math.imul(h2 ^ (h2 >>> 16), 2246822507) ^ Math.imul(h1 ^ (h1 >>> 13), 3266489909);
  return (4294967296 * (2097151 & h2) + (h1 >>> 0)).toString(36);
}

/**
 * Stable identifier for the conversation thread backing `history`.
 *
 * Built from the IMMUTABLE origin (system prompt + first user message) so it stays CONSTANT
 * while the transcript grows across a turn's rounds, yet differs between distinct
 * conversations. The growing tail is intentionally excluded so successive generate() calls of
 * the same turn land in the same bucket. Two threads opened with byte-identical origins are
 * indistinguishable from history alone and thus share a bucket (benign edge case).
 */
export function deriveConversationKey(history: Chat): string {
  const system = history.getSystemPrompt();
  let firstUser = "";
  for (const message of history) {
    if (message.getRole() === "user") {
      firstUser = message.getText();
      break;
    }
  }
  return cyrb53(`${system}\u0000${firstUser}`);
}