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

/**
 * 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, do nothing
  if (response.message.tool_calls && response.message.tool_calls.length > 0) {
    return response;
  }

  const content = response.message.content;
  if (!content) {
    return response;
  }

  // Try to parse XML tool calls from content
  const parsedCalls = parseXmlToolCalls(content);

  if (parsedCalls.length > 0) {
    logger.info(`Rewriting response: found ${parsedCalls.length} tool calls in XML content`);
    
    // Construct standard tool_calls
    const standardToolCalls: ToolCall[] = parsedCalls.map((call, index) => {
      // Ensure arguments are correctly stringified as expected by standard OpenAI/Ollama APIs
      let argumentsString = '{}';
      try {
        argumentsString = JSON.stringify(call.args);
      } catch (e) {
        logger.error('Failed to stringify arguments for tool call', call.args);
      }

      return {
        id: `call_${Date.now()}_${index}`,
        type: 'function',
        function: {
          name: call.name,
          arguments: argumentsString,
        }
      };
    });

    // We can decide to either clear the content or keep it.
    // Usually, if we parsed tool calls, we clear the content to avoid confusion
    // But retaining it is also fine. Let's clear the XML parts or the whole content to be safe.
    response.message.tool_calls = standardToolCalls;
    response.message.content = ''; 
  }

  return response;
}