'use strict';

const { execFile } = require('child_process');
const path = require('path');
const os = require('os');
const fs = require('fs');

const MLX_VENV = path.join(os.homedir(), '.mlx-env');
const MLX_PYTHON = path.join(MLX_VENV, 'bin', 'python3');

// Map of known model aliases to HuggingFace MLX repos
const MLX_MODEL_REGISTRY = {
  'llama3.1:8b':    'mlx-community/Meta-Llama-3.1-8B-Instruct-4bit',
  'llama3.1:8b-instruct-q4_K_M': 'mlx-community/Meta-Llama-3.1-8B-Instruct-4bit',
  'qwen2.5:7b':     'mlx-community/Qwen2.5-7B-Instruct-4bit',
  'qwen2.5:7b-instruct-q4_K_M': 'mlx-community/Qwen2.5-7B-Instruct-4bit',
  'phi4':            'mlx-community/phi-4-4bit',
  'phi4:latest':     'mlx-community/phi-4-4bit',
  'gemma3:4b':       'mlx-community/gemma-3-4b-it-4bit',
  'llama3.2:3b':     'mlx-community/Llama-3.2-3B-Instruct-4bit',
  'mistral:7b':      'mlx-community/Mistral-7B-Instruct-v0.3-4bit',
};

function execAsync(cmd, args, opts = {}) {
  return new Promise((resolve, reject) => {
    execFile(cmd, args, { timeout: 300000, maxBuffer: 10 * 1024 * 1024, ...opts }, (err, stdout, stderr) => {
      if (err) {
        err.stderr = stderr;
        return reject(err);
      }
      resolve({ stdout, stderr });
    });
  });
}

/**
 * Resolve a model name to an MLX HuggingFace repo.
 * Accepts: registry aliases, full HF repo names (org/model), or mlx-community/ prefixed names.
 */
function resolveModel(model) {
  if (!model) return null;
  if (MLX_MODEL_REGISTRY[model]) return MLX_MODEL_REGISTRY[model];
  // Validate model name to prevent arbitrary HF repo downloads
  if (!/^[\w\-.:\/]+$/.test(model)) return null;
  if (model.includes('/')) return model; // Full HF repo
  // Try with mlx-community prefix
  return `mlx-community/${model}`;
}

/**
 * Create an MLX local-model LLM provider.
 * Uses Apple's MLX framework via Python subprocess for GPU-accelerated inference.
 */
function createMLXProvider(config = {}) {
  return {
    type: 'mlx',

    async chat({ model, messages, tools, system, maxTokens, temperature, signal } = {}) {
      const mlxModel = resolveModel(model || config.model);
      if (!mlxModel) throw new Error('MLX: no model specified');

      // Build the prompt from messages
      const chatMessages = [];
      if (system) chatMessages.push({ role: 'system', content: system });
      if (messages) {
        for (const m of messages) {
          if (typeof m.content === 'string') {
            chatMessages.push({ role: m.role, content: m.content });
          } else if (Array.isArray(m.content)) {
            // Extract text from content blocks
            const text = m.content
              .filter(b => b.type === 'text')
              .map(b => b.text)
              .join('\n');
            if (text) chatMessages.push({ role: m.role, content: text });
            // Handle tool results
            const toolResults = m.content.filter(b => b.type === 'tool_result');
            for (const tr of toolResults) {
              const resultText = typeof tr.content === 'string' ? tr.content
                : Array.isArray(tr.content) ? tr.content.map(c => c.text || '').join('\n')
                : JSON.stringify(tr.content);
              chatMessages.push({ role: 'user', content: `[Tool result for ${tr.tool_use_id}]: ${resultText}` });
            }
          }
        }
      }

      // Build tool descriptions as system context (MLX doesn't natively support tool calling)
      let toolContext = '';
      if (tools && tools.length > 0) {
        toolContext = '\n\nYou have access to these tools. To use a tool, respond with a JSON block:\n```json\n{"tool": "tool_name", "input": {…}}\n```\n\nAvailable tools:\n' +
          tools.map(t => `- **${t.name}**: ${t.description}\n  Parameters: ${JSON.stringify(t.input_schema?.properties || {})}`).join('\n');
        if (chatMessages.length > 0 && chatMessages[0].role === 'system') {
          chatMessages[0].content += toolContext;
        } else {
          chatMessages.unshift({ role: 'system', content: toolContext.trim() });
        }
      }

      // Python script for MLX inference
      const pyScript = `
import json, sys, time
from mlx_lm import load, generate

messages = json.loads(sys.argv[1])
max_tokens = int(sys.argv[2])
model_name = sys.argv[3]

model, tokenizer = load(model_name)
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

t0 = time.time()
result = generate(model, tokenizer, prompt=prompt, max_tokens=max_tokens, verbose=False)
t1 = time.time()

output_tokens = len(tokenizer.encode(result))
prompt_tokens = len(tokenizer.encode(prompt))

print(json.dumps({
    "text": result,
    "output_tokens": output_tokens,
    "prompt_tokens": prompt_tokens,
    "latency_ms": round((t1 - t0) * 1000),
}))
`;

      const start = Date.now();
      const maxTok = maxTokens || 1024;

      try {
        const { stdout, stderr } = await execAsync(
          MLX_PYTHON,
          ['-c', pyScript, JSON.stringify(chatMessages), String(maxTok), mlxModel],
          { env: { ...process.env, TOKENIZERS_PARALLELISM: 'false' } }
        );

        // Find JSON line in output
        const lines = stdout.trim().split('\n');
        const jsonLine = lines.find(l => l.startsWith('{'));
        if (!jsonLine) throw new Error(`MLX: no JSON in output. stderr: ${stderr?.slice(0, 200)}`);

        const data = JSON.parse(jsonLine);
        const latencyMs = Date.now() - start;

        // Parse tool calls from output if tools were provided
        const content = [];
        let stopReason = 'end_turn';

        if (tools && tools.length > 0) {
          const toolMatch = data.text.match(/```json\s*\n?(\{[\s\S]*?\})\s*\n?```/);
          if (toolMatch) {
            try {
              const toolCall = JSON.parse(toolMatch[1]);
              if (toolCall.tool && toolCall.input) {
                content.push({
                  type: 'tool_use',
                  id: `mlx_${Date.now()}`,
                  name: toolCall.tool,
                  input: toolCall.input,
                });
                stopReason = 'tool_use';
              }
            } catch { /* not a valid tool call */ }
          }
        }

        if (content.length === 0) {
          content.push({ type: 'text', text: data.text });
        }

        return {
          content,
          stopReason,
          latencyMs,
          inputTokens: data.prompt_tokens,
          outputTokens: data.output_tokens,
          model: mlxModel,
          provider: 'mlx',
        };
      } catch (err) {
        throw new Error(`MLX inference failed: ${err.message}${err.stderr ? '\n' + err.stderr.slice(0, 300) : ''}`);
      }
    },

    async listModels() {
      // Return models from registry that are locally cached
      const models = [];
      const cacheDir = path.join(os.homedir(), '.cache', 'huggingface', 'hub');

      for (const [alias, repo] of Object.entries(MLX_MODEL_REGISTRY)) {
        // Check if model is cached locally
        const repoDir = `models--${repo.replace('/', '--')}`;
        const cached = fs.existsSync(path.join(cacheDir, repoDir));
        if (cached) {
          models.push({
            id: alias,
            name: `${alias} (MLX)`,
            repo,
            capabilities: ['code'],
            cached: true,
          });
        }
      }
      return models;
    },

    async isAvailable() {
      // Check if MLX Python env exists
      if (!fs.existsSync(MLX_PYTHON)) return false;
      try {
        const { stdout } = await execAsync(MLX_PYTHON, ['-c', 'import mlx_lm; print("ok")'], { timeout: 10000 });
        return stdout.trim() === 'ok';
      } catch {
        return false;
      }
    },
  };
}

module.exports = { createMLXProvider, resolveModel, MLX_MODEL_REGISTRY, MLX_PYTHON, MLX_VENV };
