// PRE Web GUI — tinygrad server lifecycle manager
// Manages starting, stopping, and health-checking the generic tinygrad GGUF
// inference server (tinygrad-server.py). Supports all architectures that
// tinygrad's Transformer.from_gguf() handles: Llama, Qwen2/3, Qwen3MoE,
// DeepSeek MLA, GLM-4MoE, etc. Integrates with gpu-detect for auto-start/stop
// on eGPU hot-plug events.

const { spawn, execSync } = require('child_process');
const http = require('http');
const fs = require('fs');
const path = require('path');
const os = require('os');
const { gpuDetector } = require('./gpu-detect');

const PRE_DIR = path.join(os.homedir(), '.pre');
const LOG_FILE = path.join(PRE_DIR, 'tinygrad-server.log');
const PID_FILE = path.join(PRE_DIR, 'tinygrad-server.pid');
const MAX_LOG_SIZE = 1024 * 1024; // 1MB
const SERVER_SCRIPT = path.join(__dirname, 'tinygrad-server.py');

// The subprocess is spawned detached with file-FD log inheritance so it survives
// PRE web-server restarts. This matters on macOS where tinygrad-NV is one-shot
// per boot — losing the subprocess means losing the GPU until reboot.
// _pid is the source of truth; _process is only set when *this* PRE instance
// spawned the running server (so we can wire .on('exit'); an adopted process
// has _pid set but _process null).
let _process = null;
let _pid = null;
let _startTime = null;
let _config = null;
let _autoStartBound = false;
let _killTimer = null;

/**
 * Find Python 3.11+ binary on this platform.
 * @returns {string|null}
 */
function getPythonPath() {
  for (const cmd of ['python3', 'python']) {
    try {
      const version = execSync(`${cmd} --version 2>&1`, { encoding: 'utf-8', timeout: 5000 }).trim();
      const match = version.match(/Python\s+(\d+)\.(\d+)/);
      if (match && (parseInt(match[1]) > 3 || (parseInt(match[1]) === 3 && parseInt(match[2]) >= 11))) {
        // Return full path
        const IS_WIN = process.platform === 'win32';
        const which = IS_WIN ? `where ${cmd}` : `which ${cmd}`;
        return execSync(which, { encoding: 'utf-8', timeout: 5000 }).trim().split('\n')[0].trim();
      }
    } catch {}
  }
  return null;
}

/**
 * Check if tinygrad is installed for the given interpreter.
 * @param {string} [python] - Python interpreter path; falls back to getPythonPath()
 * @param {string} [tinygradPath] - Optional PYTHONPATH entry for a source checkout
 * @returns {boolean}
 */
function isTinygradInstalled(python, tinygradPath) {
  const py = python || getPythonPath();
  if (!py) return false;
  try {
    const env = { ...process.env };
    if (tinygradPath) {
      env.PYTHONPATH = env.PYTHONPATH
        ? `${tinygradPath}${path.delimiter}${env.PYTHONPATH}`
        : tinygradPath;
    }
    execSync(`"${py}" -c "import tinygrad" 2>&1`, { timeout: 10000, env });
    return true;
  } catch {
    return false;
  }
}

// ── Detached subprocess support ──────────────────────────────────────────

function _isAlive(pid) {
  try { process.kill(pid, 0); return true; } catch { return false; }
}

// Verify the alive PID is actually a tinygrad-server.py process (defends
// against PID recycling after a PRE crash + long uptime).
function _isOurProcess(pid) {
  try {
    if (process.platform === 'win32') {
      const out = execSync(`wmic process where ProcessId=${pid} get CommandLine /value`,
        { timeout: 2000, encoding: 'utf-8' });
      return out.includes('tinygrad-server.py');
    }
    const cmd = execSync(`ps -p ${pid} -o command=`, { timeout: 2000, encoding: 'utf-8' });
    return cmd.includes('tinygrad-server.py');
  } catch {
    return false;
  }
}

function _readPidFile() {
  try {
    const data = JSON.parse(fs.readFileSync(PID_FILE, 'utf-8'));
    if (data && typeof data.pid === 'number') return data;
  } catch {}
  return null;
}

function _writePidFile(info) {
  try {
    fs.writeFileSync(PID_FILE, JSON.stringify(info, null, 2));
  } catch (err) {
    console.warn(`[tinygrad] Could not write PID file: ${err.message}`);
  }
}

function _clearPidFile() {
  try { fs.unlinkSync(PID_FILE); } catch {}
}

// If a prior PRE run left a tinygrad subprocess behind, pick it up.
function _tryAdopt() {
  if (_pid) return false;
  const info = _readPidFile();
  if (!info) return false;
  if (!_isAlive(info.pid) || !_isOurProcess(info.pid)) {
    _clearPidFile();
    return false;
  }
  _pid = info.pid;
  _startTime = info.startTime || Date.now();
  _config = info.config || {};
  console.log(`[tinygrad] Adopted existing server (pid ${_pid}, port ${_config.port || 8091})`);
  return true;
}

/**
 * Start the tinygrad inference server.
 * @param {Object} config - Server configuration
 * @param {string} config.model_path - Path to model file (safetensors/GGUF)
 * @param {number} [config.port=8091] - Server port
 * @param {string} [config.quantization='int8'] - Quantization: int8|nf4|fp8|none
 * @param {number} [config.context_size=8192] - Context length
 * @param {number} [config.startup_timeout=180000] - Max wait for server ready
 * @returns {Promise<{pid: number, port: number}>}
 */
async function start(config) {
  // Pick up a prior PRE run's subprocess before deciding to spawn.
  _tryAdopt();
  if (_pid) {
    const adoptedPort = _config.port || 8091;
    if (await _healthPing(adoptedPort)) {
      console.log(`[tinygrad] Already running (pid ${_pid}, port ${adoptedPort}) — no spawn needed`);
      return { pid: _pid, port: adoptedPort, adopted: true };
    }
    console.log(`[tinygrad] Adopted pid ${_pid} but port ${adoptedPort} is unresponsive — clearing`);
    _pid = null;
    _startTime = null;
    _config = null;
    _clearPidFile();
  }

  if (!config.model_path) {
    throw new Error('model_path is required');
  }

  if (!fs.existsSync(config.model_path)) {
    throw new Error(`Model file not found: ${config.model_path}`);
  }

  const python = config.python_path || getPythonPath();
  if (!python) {
    throw new Error('Python 3.11+ not found. Install via: brew install python@3.12 (macOS) or apt install python3 (Linux)');
  }

  const port = config.port || 8091;

  // Build command args for generic tinygrad GGUF server
  const args = [
    SERVER_SCRIPT,
    '--model', config.model_path,
    '--port', String(port),
    '--host', '127.0.0.1',
    '--max-context', String(config.context_size || 8192),
  ];

  // Tokenizer: HuggingFace name/path or 'auto' to detect from GGUF arch
  if (config.tokenizer_name && config.tokenizer_name !== 'auto') {
    args.push('--tokenizer', config.tokenizer_name);
  }

  // Route computation to the best available GPU.
  // DEV=NV tells tinygrad to use the NVIDIA GPU:
  //   - Linux/Windows: via /dev/nvidiactl (native driver)
  //   - macOS: via TinyGPU DriverKit extension (PCIe BAR1 over USB4/Thunderbolt)
  // Note: macOS + Blackwell (RTX 50xx/6000) requires tinygrad GSP firmware >=575.
  // Ampere (30xx) and Ada (40xx) work today. See tinygrad/tinygrad#14338.
  const gpuState = gpuDetector.getState();
  const devOverride = (gpuState.available && gpuState.vendor === 'nvidia') ? 'NV' : undefined;
  const env = { ...process.env, PYTHONUNBUFFERED: '1' };
  if (devOverride) env.DEV = devOverride;
  if (config.tinygrad_path) {
    env.PYTHONPATH = env.PYTHONPATH
      ? `${config.tinygrad_path}${path.delimiter}${env.PYTHONPATH}`
      : config.tinygrad_path;
  }
  // Ensure ~/.local/bin is on PATH so tinygrad's CUDA JIT can find nvcc
  // (or a Docker-backed shim) when the server is launched from a minimal env.
  const localBin = path.join(os.homedir(), '.local', 'bin');
  env.PATH = env.PATH ? `${localBin}${path.delimiter}${env.PATH}` : localBin;

  console.log(`[tinygrad] Starting: ${python} ${args.join(' ')}`);

  // Rotate log if too large
  try {
    const stat = fs.statSync(LOG_FILE);
    if (stat.size > MAX_LOG_SIZE) fs.unlinkSync(LOG_FILE);
  } catch {}

  // Open the log file as raw FDs and hand them to the child. Using file FDs
  // (not parent pipes) means the child writes directly to disk and survives
  // parent exit — no SIGPIPE when PRE goes away.
  const logFd = fs.openSync(LOG_FILE, 'a');

  _process = spawn(python, args, {
    stdio: ['ignore', logFd, logFd],
    env,
    detached: true,
  });

  // Drop the parent-side handle so PRE's event loop doesn't keep the child
  // tethered, and close our copy of the log FD (the child kept its inherited copy).
  _process.unref();
  try { fs.closeSync(logFd); } catch {}

  _pid = _process.pid;
  _config = config;
  _startTime = Date.now();
  _writePidFile({ pid: _pid, port, startTime: _startTime, config });

  _process.on('exit', (code, signal) => {
    console.log(`[tinygrad] Process exited: code=${code} signal=${signal}`);
    _process = null;
    _pid = null;
    _startTime = null;
    _clearPidFile();
  });

  _process.on('error', (err) => {
    console.log(`[tinygrad] Process error: ${err.message}`);
    _process = null;
    _pid = null;
    _startTime = null;
    _clearPidFile();
  });

  // Wait for server to become healthy (model loading can take 30-120s)
  const timeout = config.startup_timeout || 180000;
  const startMs = Date.now();
  while (Date.now() - startMs < timeout) {
    if (!_pid || !_isAlive(_pid)) {
      _clearPidFile();
      throw new Error('tinygrad server process died during startup. Check ' + LOG_FILE);
    }
    const healthy = await _healthPing(port);
    if (healthy) {
      console.log(`[tinygrad] Server ready on port ${port} (${((Date.now() - startMs) / 1000).toFixed(1)}s)`);
      return { pid: _pid, port };
    }
    await new Promise(r => setTimeout(r, 1000));
  }

  stop();
  throw new Error(`tinygrad server failed to start within ${timeout / 1000}s. Check ${LOG_FILE}`);
}

/**
 * Stop the tinygrad server process.
 */
function stop() {
  if (_killTimer) { clearTimeout(_killTimer); _killTimer = null; }
  _tryAdopt();
  if (!_pid) return;
  const pid = _pid;
  console.log(`[tinygrad] Stopping server (pid ${pid}) via SIGTERM`);
  try {
    if (process.platform === 'win32') {
      execSync(`taskkill /PID ${pid} /T`, { timeout: 5000 });
    } else {
      // SIGTERM only. Do NOT fall back to SIGKILL: per memory, SIGKILL on a
      // tinygrad-NV process mid-IOKit wedges Thunderbolt and panics the kernel.
      // If the process won't exit on SIGTERM, the right move is reboot.
      process.kill(pid, 'SIGTERM');
      const deadline = Date.now() + 30000;
      const poll = () => {
        if (!_isAlive(pid)) { _killTimer = null; return; }
        if (Date.now() < deadline) { _killTimer = setTimeout(poll, 1000); return; }
        console.warn(`[tinygrad] pid ${pid} did not exit 30s after SIGTERM. Refusing ` +
                     `SIGKILL (would wedge eGPU). Reboot to fully reclaim NV state.`);
        _killTimer = null;
      };
      _killTimer = setTimeout(poll, 1000);
    }
  } catch (err) {
    console.warn(`[tinygrad] stop() signal error: ${err.message}`);
  }
  _process = null;
  _pid = null;
  _startTime = null;
  _config = null;
  _clearPidFile();
}

/**
 * Ping the health endpoint (/v1/models).
 * @param {number} port
 * @returns {Promise<boolean>}
 */
function _healthPing(port) {
  return new Promise((resolve) => {
    const req = http.request({
      hostname: '127.0.0.1',
      port,
      path: '/v1/models',
      method: 'GET',
      timeout: 2000,
    }, (res) => {
      let data = '';
      res.on('data', chunk => { data += chunk; });
      res.on('end', () => {
        resolve(res.statusCode === 200);
      });
    });
    req.on('error', () => resolve(false));
    req.on('timeout', () => { req.destroy(); resolve(false); });
    req.end();
  });
}

/**
 * Check if the server is running and healthy.
 * @returns {Promise<boolean>}
 */
async function isRunning() {
  _tryAdopt();
  if (!_pid || !_isAlive(_pid)) return false;
  return _healthPing(_config?.port || 8091);
}

/**
 * Get server status.
 * @returns {Promise<Object>}
 */
async function getStatus() {
  const running = await isRunning();
  const python = _config?.python_path || getPythonPath();
  const tinygradPath = _config?.tinygrad_path;
  return {
    running,
    pid: _pid || null,
    port: _config?.port || 8091,
    model: _config?.model_path || null,
    tokenizer: _config?.tokenizer_name || 'auto',
    uptime: _startTime ? Math.round((Date.now() - _startTime) / 1000) : 0,
    gpu: gpuDetector.getState(),
    pythonPath: python,
    tinygradInstalled: isTinygradInstalled(python, tinygradPath),
  };
}

// ── Auto-start/stop on GPU hot-plug ──────────────────────────────────────

/**
 * Enable auto-start/stop based on GPU availability.
 * When an eGPU is connected and tinygrad is configured with auto_start=true,
 * the server starts automatically. When the eGPU is disconnected, it stops.
 * @param {Function} [getProviderConfig] - Returns current tinygrad provider config
 */
function enableAutoStart(getProviderConfig) {
  if (_autoStartBound) return;
  _autoStartBound = true;

  gpuDetector.on('gpu-available', async (gpuState) => {
    _tryAdopt();
    if (_pid) return; // Already running (this run, or adopted from a prior run)
    const config = getProviderConfig?.();
    if (!config || config.type !== 'tinygrad' || !config.auto_start) return;
    console.log(`[tinygrad] GPU detected (${gpuState.name}) — auto-starting server`);
    try {
      await start(config);
    } catch (err) {
      console.error(`[tinygrad] Auto-start failed: ${err.message}`);
    }
  });

  gpuDetector.on('gpu-unavailable', () => {
    _tryAdopt();
    if (!_pid) return;
    // eGPU is gone; the NV process is already dead or about to die.
    // Clear our state so a future hot-plug can re-adopt cleanly.
    console.log('[tinygrad] GPU disconnected — stopping server');
    stop();
  });
}

/**
 * Graceful shutdown — call on PRE server exit.
 *
 * Intentionally does NOT stop the tinygrad subprocess. It runs detached so it
 * survives PRE web-server restarts; on macOS this is essential because
 * tinygrad-NV is one-shot per boot. Use stop() explicitly when the user
 * disables tinygrad or the eGPU is physically disconnected.
 */
function shutdown() {
  if (_killTimer) { clearTimeout(_killTimer); _killTimer = null; }
  _autoStartBound = false;
  gpuDetector.removeAllListeners('gpu-available');
  gpuDetector.removeAllListeners('gpu-unavailable');
  if (_pid) {
    console.log(`[tinygrad] Leaving server running (pid ${_pid}) on PRE shutdown`);
  }
}

module.exports = {
  start,
  stop,
  isRunning,
  getStatus,
  getPythonPath,
  isTinygradInstalled,
  enableAutoStart,
  shutdown,
  // Exported for testing
  _healthPing,
};
