import { OllamaChatResponse, ToolCall } from '../types/ollama';
import { parseXmlToolCalls } from '../parsers';
import { logger } from '../utils/logger';
import { sanitizeContent } from '../utils/content-sanitizer';

function buildStandardToolCalls(parsedCalls: ReturnType<typeof parseXmlToolCalls>): ToolCall[] {
  return parsedCalls.map((call, index) => ({
    id: `call_${Date.now()}_${index}`,
    type: 'function',
    function: {
      name: call.name,
      arguments: call.args || {},
    }
  }));
}

/**
 * Rewrites the Ollama response to include structured tool calls if missing
 * but present in XML tags within the content.
 */
export function rewriteResponse(response: OllamaChatResponse): OllamaChatResponse {
  // If the response isn't properly formed or has no message, return as is
  if (!response || !response.message) {
    return response;
  }

  // If the response already has tool_calls, normalize them (add missing fields)
  if (response.message.tool_calls && response.message.tool_calls.length > 0) {
    response.message.tool_calls = response.message.tool_calls.map((call: any) => {
      const normalized: any = { ...call };
      // Ensure 'type' field exists
      if (!normalized.type) {
        normalized.type = 'function';
      }
      // Remove Ollama-specific 'function.index' if present (not part of OpenAI spec)
      if (normalized.function && normalized.function.index !== undefined) {
        delete normalized.function.index;
      }
      return normalized;
    });
    return response;
  }

  const content = response.message.content || '';
  const thinking = response.message.thinking || '';
  const parsedCalls = [
    ...parseXmlToolCalls(content),
    ...parseXmlToolCalls(thinking),
  ];

  if (parsedCalls.length > 0) {
    logger.info(`Rewriting response: found ${parsedCalls.length} tool calls in content/thinking`);

    response.message.tool_calls = buildStandardToolCalls(parsedCalls);
    response.message.content = '';

    if (response.message.thinking) {
      response.message.thinking = sanitizeContent(
        response.message.thinking
          .replace(/<function=([^>]+)>([\s\S]*?)<\/function>/g, '')
          .replace(/<tool_call>([\s\S]*?)<\/tool_call>/g, '')
      );
    }
    return response;
  }
  
  // Sanitize plain text responses too
  if (response.message.content) {
    response.message.content = sanitizeContent(response.message.content);
  }
  if (response.message.thinking) {
    response.message.thinking = sanitizeContent(response.message.thinking);
  }

  return response;
}