import type { ChatRequest, ChatResponse, ChatStreamEvent, ModelClient, ModelToolCall, Message } from '../types.js';
import { normalizeMessagesForLocal } from '../media.js';

export interface LlamaCppClientConfig {
  endpoint: string;
  model: string;
  authToken?: string;
  /** Per-request timeout in ms. Default: 180000 (3 minutes). */
  requestTimeout?: number;
}

interface LlamaCppToolCall {
  id: string;
  type: 'function';
  function: {
    name: string;
    arguments: string; // JSON string
  };
}

interface LlamaCppResponse {
  choices: Array<{
    message: {
      content: string | null;
      tool_calls?: LlamaCppToolCall[];
    };
    finish_reason?: string;
  }>;
  usage: { prompt_tokens: number; completion_tokens: number };
}

interface LlamaCppStreamChunk {
  choices: Array<{
    delta?: {
      content?: string;
      tool_calls?: Array<{
        index: number;
        id?: string;
        type?: string;
        function?: {
          name?: string;
          arguments?: string;
        };
      }>;
    };
    finish_reason?: string | null;
  }>;
  usage?: { prompt_tokens: number; completion_tokens: number };
}

function sanitizeToolParametersForLlamaCpp(toolName: string, parameters: unknown): unknown {
  // llama.cpp is stricter than most tool-call APIs about JSON schema.
  // In particular, some builds reject extra optional properties for common tools.
  // Keep the full schema for most tools, but reduce known-problematic ones.
  if (toolName !== 'web_search') {
    return parameters;
  }

  if (!parameters || typeof parameters !== 'object') {
    return parameters;
  }

  const schema = parameters as {
    type?: unknown;
    properties?: Record<string, unknown>;
    required?: unknown;
  };

  const querySchema = schema.properties?.query;
  if (!querySchema) {
    return parameters;
  }

  return {
    type: 'object',
    properties: { query: querySchema },
    required: ['query'],
  };
}

/** Message format for OpenAI-compatible chat completions API. */
interface LlamaCppChatMessage {
  role: 'system' | 'user' | 'assistant' | 'tool';
  content: string;
  tool_calls?: Array<{
    id: string;
    type: 'function';
    function: { name: string; arguments: string };
  }>;
  tool_call_id?: string;
}

/**
 * Normalize messages for llama.cpp's OpenAI-compatible API, converting
 * Anthropic-style tool_use/tool_result blocks for local model consumption.
 *
 * Uses a hybrid approach for maximum template compatibility:
 * - Assistant messages with tool_use blocks → assistant with tool_calls array
 *   (native format, so the model knows what it called)
 * - User messages with tool_result blocks → single user message with structured
 *   text (NOT role: 'tool', because many GGUF chat templates silently drop
 *   unknown roles, making tool results invisible to the model)
 * - All other messages → plain text via normalizeMessagesForLocal
 */
export function normalizeMessagesForLlamaCpp(
  system: string | undefined,
  messages: Message[],
): LlamaCppChatMessage[] {
  // 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 LlamaCppChatMessage[];
  }

  const result: LlamaCppChatMessage[] = [];

  // Track tool_use_id → tool_name for labeling results
  const toolNameMap = new Map<string, string>();

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

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

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

    if (msg.role === 'assistant') {
      const textParts: string[] = [];
      const toolCalls: LlamaCppChatMessage['tool_calls'] = [];

      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);}
          let argsStr: string;
          try {
            argsStr = JSON.stringify(block.input);
          } catch {
            argsStr = '{}';
          }
          toolCalls!.push({
            id,
            type: 'function',
            function: {
              name,
              arguments: argsStr,
            },
          });
        }
      }

      const chatMsg: LlamaCppChatMessage = {
        role: 'assistant',
        content: textParts.join('\n'),
      };
      if (toolCalls!.length > 0) {
        chatMsg.tool_calls = toolCalls;
      }
      result.push(chatMsg);
    } else if (msg.role === 'user') {
      const toolResults = blocks.filter(b => b.type === 'tool_result');

      if (toolResults.length > 0) {
        // Send tool results as a user message with clear text formatting.
        // Many GGUF chat templates don't handle role: 'tool', silently
        // dropping those messages. A user message always works.
        const parts: string[] = [];
        for (const tr of toolResults) {
          const toolUseId = tr.tool_use_id as string;
          const toolName = toolUseId ? toolNameMap.get(toolUseId) : undefined;
          const content = (tr.content as string) ?? '';
          const isError = tr.is_error ? ' (error)' : '';
          const label = toolName ?? 'unknown';
          parts.push(`[Tool "${label}" result${isError}]\n${content}`);
        }
        result.push({ role: 'user', content: parts.join('\n\n') });
      } else {
        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 class LlamaCppClient implements ModelClient {
  private endpoint: string;
  private model: string;
  private authToken?: string;
  private requestTimeout: number;

  constructor(config: LlamaCppClientConfig) {
    this.endpoint = config.endpoint.replace(/\/$/, '');
    this.model = config.model;
    this.authToken = config.authToken;
    this.requestTimeout = config.requestTimeout ?? 180_000; // 3 minutes default
  }

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

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
    };

    if (this.authToken) {
      headers['Authorization'] = `Bearer ${this.authToken}`;
    }

    let response: Response;
    try {
      const body: Record<string, unknown> = {
        model: this.model,
        messages,
        max_tokens: request.maxTokens ?? 2048,
      };

      // Pass tool definitions to the API if provided
      if (request.tools && request.tools.length > 0) {
        body.tools = request.tools.map(t => ({
          type: 'function' as const,
          function: {
            name: t.name,
            description: t.description,
            parameters: sanitizeToolParametersForLlamaCpp(t.name, t.input_schema),
          },
        }));
      }

      const controller = new AbortController();
      const timer = setTimeout(() => controller.abort(), this.requestTimeout);
      try {
        response = await fetch(`${this.endpoint}/v1/chat/completions`, {
          method: 'POST',
          headers,
          body: JSON.stringify(body),
          signal: controller.signal,
        });
      } finally {
        clearTimeout(timer);
      }
    } catch (error) {
      if (error instanceof DOMException && error.name === 'AbortError') {
        throw new Error(`llama-server request timed out after ${Math.round(this.requestTimeout / 1000)}s`);
      }
      if (error instanceof TypeError && error.message.includes('fetch failed')) {
        throw new Error(`llama-server not running at ${this.endpoint}`);
      }
      throw error;
    }

    if (!response.ok) {
      const text = await response.text();
      throw new Error(`llama-server error (${response.status}): ${text}`);
    }

    const data = (await response.json()) as LlamaCppResponse;

    // Parse tool calls from the response, if present
    const toolCalls: ModelToolCall[] = data.choices[0]?.message?.tool_calls?.map((tc) => ({
      id: tc.id ?? `llamacpp_tc_${Math.random().toString(36).slice(2, 8)}`,
      name: tc.function.name,
      args: JSON.parse(tc.function.arguments),
    })) ?? [];

    // Set stopReason to 'tool_use' when tool_calls are present
    const stopReason = toolCalls.length > 0 ? 'tool_use' : (data.choices[0]?.finish_reason ?? 'stop');

    return {
      content: data.choices[0]?.message?.content ?? '',
      stopReason,
      usage: {
        inputTokens: data.usage?.prompt_tokens ?? 0,
        outputTokens: data.usage?.completion_tokens ?? 0,
      },
      ...(toolCalls.length > 0 ? { toolCalls } : {}),
    };
  }

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

    const headers: Record<string, string> = {
      'Content-Type': 'application/json',
    };

    if (this.authToken) {
      headers['Authorization'] = `Bearer ${this.authToken}`;
    }

    try {
      const body: Record<string, unknown> = {
        model: this.model,
        messages,
        max_tokens: request.maxTokens ?? 2048,
        stream: true,
      };

      // Pass tool definitions to the API if provided
      if (request.tools && request.tools.length > 0) {
        body.tools = request.tools.map(t => ({
          type: 'function' as const,
          function: {
            name: t.name,
            description: t.description,
            parameters: sanitizeToolParametersForLlamaCpp(t.name, t.input_schema),
          },
        }));
      }

      const response = await fetch(`${this.endpoint}/v1/chat/completions`, {
        method: 'POST',
        headers,
        body: JSON.stringify(body),
      });

      if (!response.ok) {
        const text = await response.text();
        throw new Error(`llama-server error (${response.status}): ${text}`);
      }

      if (!response.body) {
        throw new Error('No response body for streaming');
      }

      const reader = response.body.getReader();
      const decoder = new TextDecoder();
      let buffer = '';
      let usage = { inputTokens: 0, outputTokens: 0 };
      // Accumulate tool call deltas across streamed chunks
      const toolCallAccumulators: Map<number, { id: string; name: string; arguments: string }> = new Map();

      while (true) {
        const { done, value } = await reader.read();
        if (done) {break;}

        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split('\n');
        buffer = lines.pop() ?? '';

        for (const line of lines) {
          const trimmed = line.trim();
          if (!trimmed || !trimmed.startsWith('data: ')) {continue;}

          const data = trimmed.slice(6);
          if (data === '[DONE]') {continue;}

          try {
            const chunk = JSON.parse(data) as LlamaCppStreamChunk;

            if (chunk.choices[0]?.delta?.content) {
              yield { type: 'content', content: chunk.choices[0].delta.content };
            }

            // Accumulate tool call deltas from the stream
            if (chunk.choices[0]?.delta?.tool_calls) {
              for (const tc of chunk.choices[0].delta.tool_calls) {
                if (!toolCallAccumulators.has(tc.index)) {
                  toolCallAccumulators.set(tc.index, {
                    id: tc.id ?? `llamacpp_tc_${tc.index}`,
                    name: tc.function?.name ?? '',
                    arguments: '',
                  });
                }
                const acc = toolCallAccumulators.get(tc.index)!;
                if (tc.function?.name) {acc.name = tc.function.name;}
                if (tc.function?.arguments) {acc.arguments += tc.function.arguments;}
              }
            }

            if (chunk.usage) {
              usage = {
                inputTokens: chunk.usage.prompt_tokens,
                outputTokens: chunk.usage.completion_tokens,
              };
            }
          } catch {
            // Skip malformed JSON
          }
        }
      }

      // Yield completed tool calls before the done event
      for (const [, acc] of toolCallAccumulators) {
        yield {
          type: 'tool_use',
          toolCall: {
            id: acc.id,
            name: acc.name,
            args: JSON.parse(acc.arguments),
          },
        };
      }

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