// src/ollama.ts
// Copyright (C) 2026 Rob Colbert <rob.colbert@openplatform.us>
// Licensed under the Apache License, Version 2.0

import assert from "node:assert";

import { Ollama } from "ollama";
import numeral from "numeral";
import {
  AiApi,
  IAiChatOptions,
  IAiChatResponse,
  IToolCall,
  IToolCallResult,
  IAiGenerateOptions,
  IAiGenerateResponse,
  IAiLogger,
  IAiModelConfig,
  IAiModelListResult,
  IAiModelProbeResult,
  IAiProvider,
  IAiResponseStreamFn,
} from "./api.js";
import { IAiEnvironment } from "./config/env.ts";
import type { Message as OllamaMessage } from "ollama";

export class OllamaAiApi extends AiApi {
  protected client: Ollama;

  constructor(env: IAiEnvironment, provider: IAiProvider, logger?: IAiLogger) {
    super(env, provider, logger);
    this.client = new Ollama({
      host: this.provider.baseUrl,
      headers: { Authorization: `Bearer ${this.provider.apiKey}` },
    });
  }

  async listModels(): Promise<IAiModelListResult> {
    await this.log.debug("OllamaAiApi.listModels called");

    const response = await this.client.list();
    await this.log.debug("Ollama list response", { models: response.models });

    const models = response.models.map((model) => {
      const parameterCount = this.parseParameterCount(
        model.details.parameter_size,
      );
      return {
        id: model.name,
        name: model.name,
        parameterLabel: model.details.parameter_size,
        parameterCount,
        contextWindow: undefined,
      };
    });

    return { models };
  }

  async probeModel(modelId: string): Promise<IAiModelProbeResult> {
    await this.log.debug("OllamaAiApi.probeModel called", { modelId });

    const response = await this.client.show({ model: modelId });
    await this.log.debug("Ollama show response", {
      modelId,
      capabilities: response.capabilities,
      details: response.details,
      modelInfo: response.model_info,
      template: response.template,
      system: response.system,
    });

    const capabilities = this.analyzeCapabilities(response, modelId);
    const settings = this.extractSettings(response);

    return {
      capabilities,
      settings,
    };
  }

  private parseParameterCount(parameterSize?: string): number | undefined {
    if (!parameterSize) return undefined;
    const match = parameterSize.match(/^([\d.]+)[BbMm]?$/);
    if (!match) return undefined;
    const value = parseFloat(match[1]);
    if (parameterSize.toLowerCase().includes("m")) {
      return value / 1000;
    }
    return value;
  }

  private analyzeCapabilities(
    response: Awaited<ReturnType<typeof this.client.show>>,
    modelId: string,
  ): IAiModelProbeResult["capabilities"] {
    const capabilities = response.capabilities || [];
    const modelInfo = response.model_info as unknown as
      | Record<string, unknown>
      | undefined;

    return {
      canCallTools:
        capabilities.includes("tools") ||
        capabilities.includes("function_calling"),
      hasVision:
        capabilities.includes("vision") ||
        !!modelInfo?.["vision_model"] ||
        !!modelInfo?.["clip"],
      hasEmbedding: capabilities.includes("embeddings"),
      hasThinking: capabilities.includes("reasoning"),
      isInstructTuned:
        modelId.toLowerCase().includes("instruct") ||
        modelId.toLowerCase().includes("chat") ||
        modelId.toLowerCase().includes("-it"),
    };
  }

  private extractSettings(
    response: Awaited<ReturnType<typeof this.client.show>>,
  ): IAiModelProbeResult["settings"] {
    const parameters = response.parameters || "";
    const settings: IAiModelProbeResult["settings"] = {};

    const temperatureMatch = parameters.match(/temperature\s+(\d+\.?\d*)/i);
    if (temperatureMatch) {
      settings.temperature = parseFloat(temperatureMatch[1]);
    }

    const topPMatch = parameters.match(/top_p\s+(\d+\.?\d*)/i);
    if (topPMatch) {
      settings.topP = parseFloat(topPMatch[1]);
    }

    const topKMatch = parameters.match(/top_k\s+(\d+)/i);
    if (topKMatch) {
      settings.topK = parseInt(topKMatch[1], 10);
    }

    const numCtxMatch = parameters.match(/num_ctx\s+(\d+)/i);
    if (numCtxMatch) {
      settings.numCtx = parseInt(numCtxMatch[1], 10);
    }

    return Object.keys(settings).length > 0 ? settings : undefined;
  }

  async generate(
    model: IAiModelConfig,
    options: IAiGenerateOptions,
    streamCallback?: IAiResponseStreamFn,
  ): Promise<IAiGenerateResponse> {
    await this.log.debug("OllamaAiApi.generate called", {
      provider: model.provider.name,
      modelId: model.modelId,
    });

    const response = await this.client.generate({
      model: model.modelId,
      prompt: options.prompt,
      system: options.systemPrompt,
      stream: true,
    });

    let lastChunk;
    for await (const chunk of response) {
      await this.log.debug("stream chunk received", { chunk });
      lastChunk = chunk;

      if (streamCallback) {
        if (chunk.thinking) {
          await streamCallback({
            type: 'thinking',
            data: chunk.thinking,
          });
        }
        if (chunk.response) {
          await streamCallback({
            type: 'response',
            data: chunk.response,
          });
        }
      }
    }
    assert(lastChunk, "no stream response chunks received");

    return {
      done: lastChunk.done,
      doneReason: lastChunk.done_reason,
      response: lastChunk.response,
      thinking: lastChunk.thinking,
      stats: {
        duration: {
          seconds: lastChunk.total_duration,
          text: numeral(lastChunk.total_duration).format("hh:mm:ss"),
        },
        tokenCounts: {
          input: lastChunk.prompt_eval_count,
          response: lastChunk.eval_count,
          thinking: 0,
        },
      },
    };
  }

  async chat(
    model: IAiModelConfig,
    options: IAiChatOptions,
    streamCallback?: IAiResponseStreamFn,
  ): Promise<IAiChatResponse> {
    await this.log.debug("OllamaAiApi.chat called", {
      provider: model.provider.name,
      modelId: model.modelId,
    });

    // VALIDATE: Ensure we have at least one message with content
    if (!options.userPrompt || !options.userPrompt.trim()) {
      throw new Error("userPrompt is required and cannot be empty");
    }

    const maxIterations = options.maxToolIterations ?? 5;
    let iteration = 0;

    // Build messages array like OpenAI does
    const messages: OllamaMessage[] = [];
    
    // Add system prompt if present
    if (options.systemPrompt) {
      messages.push({
        role: 'system',
        content: options.systemPrompt,
      });
    }
    
    // Add context messages
    if (options.context) {
      for (const msg of options.context) {
        if (msg.content && msg.content.trim()) {
          if (msg.role === 'tool') {
            messages.push({
              role: 'tool',
              content: msg.content,
              tool_name: msg.toolName,
            });
          } else {
            messages.push({
              role: msg.role as 'user' | 'assistant' | 'system',
              content: msg.content,
            });
          }
        }
      }
    }
    
    // Add user prompt (required)
    messages.push({
      role: 'user',
      content: options.userPrompt,
    });
    
    // VALIDATE: Ensure messages array is not empty before calling API
    if (messages.length === 0) {
      throw new Error("Messages array is empty - cannot call Ollama API with no messages");
    }
    
    // DEBUG: Log what we're sending to Ollama
    await this.log.debug("Ollama chat request", {
      messagesCount: messages.length,
      messages: messages.map(m => ({ role: m.role, contentLength: m.content?.length || 0 })),
      userPrompt: options.userPrompt?.slice(0, 100),
      contextCount: options.context?.length || 0,
    });
    
    const allToolCallResults: IToolCallResult[] = [];
    const allToolCalls: IToolCall[] = [];
    let totalAccumulatedResponse = "";
    let totalAccumulatedThinking = "";

    while (iteration < maxIterations) {
      iteration++;

      const ollamaTools = options.tools
        ? options.tools.map((tool) => ({
            type: tool.definition.type,
            function: {
              name: tool.definition.function.name,
              description: tool.definition.function.description,
              parameters: tool.definition.function.parameters,
            },
          }))
        : undefined;

      const response = await this.client.chat({
        model: model.modelId,
        messages,
        stream: true,
        think: model.params.reasoning,
        tools: ollamaTools,
      });

      let lastChunk;
      let accumulatedThinking = "";
      let accumulatedResponse = "";
      const streamedToolCalls: Array<{
        callId: string;
        function: { name: string; arguments: any };
      }> = [];

      for await (const chunk of response) {
        await this.log.debug("stream chunk received", { chunk });
        lastChunk = chunk;

        if (streamCallback) {
          if (chunk.message.thinking) {
            accumulatedThinking += chunk.message.thinking;
            await streamCallback({
              type: 'thinking',
              data: chunk.message.thinking,
            });
          }
          if (chunk.message.content) {
            accumulatedResponse += chunk.message.content;
            await streamCallback({
              type: 'response',
              data: chunk.message.content,
            });
          }
          if (chunk.message.tool_calls) {
            for (const tc of chunk.message.tool_calls) {
              const params = JSON.stringify(tc.function.arguments);
              const callId = `tool_${tc.function.name}_${Date.now()}`;
              
              const toolCall: IToolCall = {
                callId,
                function: {
                  name: tc.function.name,
                  arguments: JSON.stringify(tc.function.arguments),
                },
              };
              streamedToolCalls.push(toolCall);
              allToolCalls.push(toolCall);
              
              await streamCallback({
                type: 'toolCall',
                data: params,
                toolCallId: callId,
                toolName: tc.function.name,
                params,
              });
            }
          }
        }
      }
      assert(lastChunk, "no response chunks received");

      // Use accumulated thinking/response for final response
      const finalThinking = accumulatedThinking || lastChunk.message.thinking;
      const finalResponse = accumulatedResponse || lastChunk.message.content;
      
      // Accumulate across iterations
      totalAccumulatedResponse += finalResponse || "";
      totalAccumulatedThinking += finalThinking || "";

      // Use accumulated tool calls from stream
      const toolCalls = streamedToolCalls;

      if (!toolCalls || toolCalls.length === 0) {
        return {
          done: lastChunk.done,
          doneReason: lastChunk.done_reason,
          response: totalAccumulatedResponse,
          thinking: totalAccumulatedThinking,
          toolCalls: allToolCalls.length > 0 ? allToolCalls : undefined,
          toolCallResults: allToolCallResults.length > 0 ? allToolCallResults : undefined,
          stats: {
            duration: {
              seconds: lastChunk.total_duration,
              text: numeral(lastChunk.total_duration).format("hh:mm:ss"),
            },
            tokenCounts: {
              input: lastChunk.prompt_eval_count,
              response: lastChunk.eval_count,
              thinking: 0,
            },
          },
        };
      }

      const toolCallResults = await this.executeToolCalls(
        toolCalls,
        options.tools || [],
      );
      allToolCallResults.push(...toolCallResults);

      // DEBUG: Log tool results being added to context
      await this.log.debug("tool results ready for context", {
        toolCallResults: toolCallResults.map(r => ({
          callId: r.callId,
          functionName: r.functionName,
          resultLength: r.result?.length || 0,
          hasError: !!r.error,
        })),
      });

      const assistantMsg: OllamaMessage = {
        role: "assistant",
        content: accumulatedResponse || lastChunk.message.content,
      };
      if (lastChunk.message.thinking) {
        assistantMsg.thinking = lastChunk.message.thinking;
      }
      if (lastChunk.message.tool_calls) {
        assistantMsg.tool_calls = lastChunk.message.tool_calls;
      }
      messages.push(assistantMsg);

      for (const result of toolCallResults) {
        const toolContent = result.error 
          ? `Error executing ${result.functionName}: ${result.error}`
          : result.result;
        
        const toolMsg = {
          role: "tool" as const,
          content: toolContent,
        };
        await this.log.debug("adding tool result to messages", {
          contentLength: toolMsg.content?.length || 0,
          hasContent: !!(toolMsg.content && toolMsg.content.length),
        });
        messages.push(toolMsg);
      }

      // DEBUG: Log full messages array before next iteration
      await this.log.debug("messages array for next Ollama API call", {
        messageCount: messages.length,
        messages: messages.map(m => ({
          role: m.role,
          contentLength: m.content?.length || 0,
          tool_name: (m as any).tool_name,
          contentPreview: m.content?.slice(0, 200),
        })),
      });

      // VALIDATE: Ensure tool results are in messages
      const toolMessages = messages.filter(m => m.role === 'tool');
      if (toolMessages.length === 0 && toolCallResults.length > 0) {
        await this.log.error("CRITICAL: tool results NOT in messages array", {
          toolCallResultsCount: toolCallResults.length,
          messagesCount: messages.length,
        });
      }
    }

    return {
      done: false,
      doneReason: "max_tool_iterations_reached",
      response: "",
      thinking: undefined,
      toolCalls: allToolCalls.length > 0 ? allToolCalls : undefined,
      toolCallResults: allToolCallResults,
      stats: {
        duration: {
          seconds: 0,
          text: "00:00:00",
        },
        tokenCounts: {
          input: 0,
          response: 0,
          thinking: 0,
        },
      },
    };
  }
}
