import { Ollama, type Tool, type Message as OllamaMessage } from 'ollama';
import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient, ToolDefinition, ModelToolCall, Message } from '../types.js';
import { normalizeMessagesForLocal } from '../media.js';

/**
 * Normalize messages for Ollama, converting Anthropic-style tool_use/tool_result
 * blocks to Ollama's native tool calling format.
 *
 * - Assistant messages with tool_use blocks → assistant with tool_calls array
 * - User messages with tool_result blocks → one { role: 'tool' } message per result
 * - All other messages → plain text via normalizeMessagesForLocal
 */
export function normalizeMessagesForOllama(
  system: string | undefined,
  messages: Message[],
): OllamaMessage[] {
  // Check if any messages contain structured tool blocks
  const hasToolBlocks = messages.some(
    m => Array.isArray(m.content) && (m.content as Record<string, unknown>[]).some(
      b => b.type === 'tool_use' || b.type === 'tool_result',
    ),
  );

  // Fast path: no tool blocks, use the simple normalizer
  if (!hasToolBlocks) {
    return normalizeMessagesForLocal(system, messages) as OllamaMessage[];
  }

  const result: OllamaMessage[] = [];

  if (system) {
    result.push({ role: 'system', content: system });
  }

  // Track tool_use_id → tool_name so tool_result messages can set tool_name
  const toolNameMap = new Map<string, string>();

  for (const msg of messages) {
    if (typeof msg.content === 'string') {
      // Plain text message
      result.push({ role: msg.role, content: msg.content });
      continue;
    }

    const blocks = msg.content as Record<string, unknown>[];

    if (msg.role === 'assistant') {
      // Extract text and tool_use blocks
      const textParts: string[] = [];
      const toolCalls: Array<{ function: { name: string; arguments: Record<string, unknown> } }> = [];

      for (const block of blocks) {
        if (block.type === 'text' && typeof block.text === 'string') {
          textParts.push(block.text);
        } else if (block.type === 'tool_use') {
          const name = block.name as string;
          const id = block.id as string;
          if (id) {toolNameMap.set(id, name);}
          toolCalls.push({
            function: {
              name,
              arguments: (block.input ?? {}) as Record<string, unknown>,
            },
          });
        }
      }

      const ollamaMsg: OllamaMessage = {
        role: 'assistant',
        content: textParts.join('\n'),
      };
      if (toolCalls.length > 0) {
        ollamaMsg.tool_calls = toolCalls;
      }
      result.push(ollamaMsg);
    } else if (msg.role === 'user') {
      // Check if this is a tool_result message
      const toolResults = blocks.filter(b => b.type === 'tool_result');

      if (toolResults.length > 0) {
        // Convert each tool_result to a separate 'tool' role message
        for (const tr of toolResults) {
          const toolUseId = tr.tool_use_id as string;
          const ollamaMsg: OllamaMessage = {
            role: 'tool',
            content: (tr.content as string) ?? '',
          };
          // Set tool_name so Ollama associates this result with the correct tool call
          const toolName = toolUseId ? toolNameMap.get(toolUseId) : undefined;
          if (toolName) {
            ollamaMsg.tool_name = toolName;
          }
          result.push(ollamaMsg);
        }
      } else {
        // Regular user message with content parts (e.g. images)
        const textParts = blocks
          .filter(b => b.type === 'text' && typeof b.text === 'string')
          .map(b => b.text as string);
        if (textParts.length > 0) {
          result.push({ role: 'user', content: textParts.join('\n') });
        }
      }
    }
  }

  return result;
}

export interface OllamaClientConfig {
  host?: string;
  model: string;
  numGpu?: number;
}

export class OllamaClient implements ModelClient {
  private client: Ollama;
  private model: string;
  private numGpu: number;
  /** Cached result of model tool-support check. null = not yet checked. */
  private _supportsTools: boolean | null = null;

  constructor(config: OllamaClientConfig) {
    this.client = new Ollama({
      host: config.host ?? 'http://localhost:11434',
    });
    this.model = config.model;
    this.numGpu = config.numGpu ?? -1;
  }

  /**
   * Check whether the current model supports tool calling via Ollama's
   * /api/show endpoint. Result is cached for the lifetime of this client.
   * On error (e.g. old Ollama version without capabilities), defaults to
   * true (optimistic — let the server decide).
   */
  private async checkToolSupport(): Promise<boolean> {
    if (this._supportsTools !== null) {return this._supportsTools;}
    try {
      const info = await this.client.show({ model: this.model });
      const caps: string[] = (info as any).capabilities ?? [];
      this._supportsTools = caps.includes('tools');
    } catch {
      // Old Ollama or network issue — assume tools are supported
      this._supportsTools = true;
    }
    return this._supportsTools;
  }

  /**
   * Convert Flynn ToolDefinition[] to Ollama Tool[] format.
   */
  private convertTools(tools: ToolDefinition[]): Tool[] {
    return tools.map(t => ({
      type: 'function',
      function: {
        name: t.name,
        description: t.description,
        parameters: {
          type: t.input_schema.type,
          required: t.input_schema.required,
          properties: t.input_schema.properties as Record<string, any>,
        },
      },
    }));
  }

  async chat(request: ChatRequest): Promise<ChatResponse> {
    const messages = normalizeMessagesForOllama(request.system, request.messages);

    // Build the chat params, optionally including tools
    const chatParams: Parameters<typeof this.client.chat>[0] = {
      model: this.model,
      messages,
      options: {
        num_gpu: this.numGpu,
      },
    };

    if (request.tools && request.tools.length > 0 && await this.checkToolSupport()) {
      chatParams.tools = this.convertTools(request.tools);
    }

    const response = await this.client.chat(chatParams);

    // Extract content, checking for thinking field from reasoning models
    let content = response.message.content;
    let thinkingContent: string | undefined;
    const thinking = (response.message as any).thinking;
    if (thinking && typeof thinking === 'string') {
      if (!content) {
        // If no regular content, use thinking as content
        content = thinking;
      }
      thinkingContent = thinking;
    }

    // Parse tool_calls from the response
    const toolCalls: ModelToolCall[] = response.message.tool_calls?.map((tc, i) => ({
      id: `ollama_tc_${i}`,
      name: tc.function.name,
      args: tc.function.arguments,
    })) ?? [];

    const hasToolCalls = toolCalls.length > 0;

    return {
      content,
      stopReason: hasToolCalls ? 'tool_use' : (response.done_reason ?? 'stop'),
      usage: {
        inputTokens: response.prompt_eval_count ?? 0,
        outputTokens: response.eval_count ?? 0,
      },
      ...(hasToolCalls ? { toolCalls } : {}),
      ...(thinkingContent ? { thinkingContent } : {}),
    };
  }

  async *chatStream(request: ChatRequest): AsyncIterable<ChatStreamEvent> {
    const messages = normalizeMessagesForOllama(request.system, request.messages);

    try {
      // Build tools array if provided and model supports them
      const supportsTools = await this.checkToolSupport();
      const tools = request.tools && request.tools.length > 0 && supportsTools
        ? this.convertTools(request.tools)
        : undefined;

      const stream = await this.client.chat({
        model: this.model,
        messages,
        stream: true,
        options: {
          num_gpu: this.numGpu,
        },
        ...(tools ? { tools } : {}),
      });

      let inputTokens = 0;
      let outputTokens = 0;

      for await (const chunk of stream) {
        if (chunk.message?.content) {
          yield { type: 'content', content: chunk.message.content };
        }

        // Handle thinking field from reasoning models (e.g., deepseek-r1)
        const thinking = (chunk.message as any)?.thinking;
        if (thinking && typeof thinking === 'string') {
          yield { type: 'content', content: thinking };
        }

        if (chunk.prompt_eval_count) {
          inputTokens = chunk.prompt_eval_count;
        }
        if (chunk.eval_count) {
          outputTokens = chunk.eval_count;
        }

        if (chunk.done) {
          // Handle tool_calls in the final chunk
          const toolCalls = (chunk.message as any)?.tool_calls;
          if (toolCalls && Array.isArray(toolCalls)) {
            for (let i = 0; i < toolCalls.length; i++) {
              const tc = toolCalls[i];
              yield {
                type: 'tool_use',
                toolCall: {
                  id: `ollama_tc_${i}`,
                  name: tc.function.name,
                  args: tc.function.arguments,
                },
              };
            }
          }

          yield {
            type: 'done',
            usage: {
              inputTokens,
              outputTokens,
            },
          };
        }
      }
    } catch (error) {
      yield {
        type: 'error',
        error: error instanceof Error ? error : new Error(String(error)),
      };
    }
  }
}
