import { spawn } from 'node:child_process';
import type { ChildProcess } from 'node:child_process';
import { homedir } from 'node:os';
import http from 'node:http';
import type { LlmProvider, CompletionOptions, CompletionResult } from './types.js';
import { OpenAiProvider } from './openai.js';

export interface ManagedVllmConfig {
  /** Path to the Python venv containing vLLM (e.g. "~/vllm_env") */
  venvPath: string;
  /** Model to serve (e.g. "Qwen/Qwen2.5-7B-Instruct-AWQ") */
  model: string;
  /** Port for vLLM HTTP server (default: 8000) */
  port?: number;
  /** GPU memory utilization fraction (default: 0.75) */
  gpuMemoryUtilization?: number;
  /** Max model context length (default: 4096) */
  maxModelLen?: number;
  /** Minutes of inactivity before killing vLLM to free GPU (default: 15) */
  idleTimeoutMinutes?: number;
  /** Additional args passed to `vllm serve` */
  extraArgs?: string[];
  /** Override for testing — inject custom spawn function */
  spawnFn?: typeof spawn;
  /** Override for testing — inject custom health check */
  healthCheckFn?: (port: number) => Promise<boolean>;
}

export type ManagedVllmState = 'stopped' | 'starting' | 'running' | 'error';

export interface ManagedVllmStatus {
  state: ManagedVllmState;
  lastError: string | null;
  pid: number | null;
  uptime: number | null;
  /** Persistent disable. When true, complete()/ensureRunning() short-circuit. */
  disabled: boolean;
}

const POLL_INTERVAL_MS = 2000;
const STARTUP_TIMEOUT_MS = 120_000;
/** After entering error state, wait this long before retrying startup. */
const ERROR_COOLDOWN_MS = 60_000;

/**
 * Managed vLLM provider — spawns and manages a local vLLM process.
 *
 * Starts vLLM on first `complete()` call, stops it after configurable idle
 * timeout to free GPU memory. Delegates actual inference to an inner
 * OpenAiProvider pointed at the local vLLM endpoint.
 */
export class ManagedVllmProvider implements LlmProvider {
  readonly name = 'vllm-managed';

  private process: ChildProcess | null = null;
  private inner: OpenAiProvider | null = null;
  private state: ManagedVllmState = 'stopped';
  private lastError: string | null = null;
  private disabled = false;
  private lastUsed = 0;
  private startedAt = 0;
  private errorAt = 0;
  private idleTimer: ReturnType<typeof setInterval> | null = null;
  private startPromise: Promise<void> | null = null;

  private readonly venvPath: string;
  private readonly model: string;
  private readonly port: number;
  private readonly gpuMemoryUtilization: number;
  private readonly maxModelLen: number;
  private readonly idleTimeoutMs: number;
  private readonly extraArgs: string[];
  private readonly spawnFn: typeof spawn;
  private readonly healthCheckFn: (port: number) => Promise<boolean>;

  constructor(config: ManagedVllmConfig) {
    // Expand ~ in venvPath
    this.venvPath = config.venvPath.startsWith('~')
      ? config.venvPath.replace('~', homedir())
      : config.venvPath;
    this.model = config.model;
    this.port = config.port ?? 8000;
    this.gpuMemoryUtilization = config.gpuMemoryUtilization ?? 0.75;
    this.maxModelLen = config.maxModelLen ?? 4096;
    this.idleTimeoutMs = (config.idleTimeoutMinutes ?? 15) * 60 * 1000;
    this.extraArgs = config.extraArgs ?? [];
    this.spawnFn = config.spawnFn ?? spawn;
    this.healthCheckFn = config.healthCheckFn ?? defaultHealthCheck;
  }

  async complete(options: CompletionOptions): Promise<CompletionResult> {
    if (this.disabled) {
      throw new Error('vllm-managed provider is disabled — run `mcpctl provider enable <name>` to re-enable');
    }
    await this.ensureRunning();
    this.lastUsed = Date.now();
    this.resetIdleTimer();
    return this.inner!.complete(options);
  }

  async listModels(): Promise<string[]> {
    if (this.state === 'running' && this.inner) {
      return this.inner.listModels();
    }
    return [this.model];
  }

  /**
   * A managed provider is "available" unless in a permanent error state.
   * When stopped, it can be auto-started on demand. Disabled providers
   * report unavailable so health probes don't try to wake them.
   */
  async isAvailable(): Promise<boolean> {
    if (this.disabled) return false;
    return this.state !== 'error';
  }

  getStatus(): ManagedVllmStatus {
    return {
      state: this.state,
      lastError: this.lastError,
      pid: this.process?.pid ?? null,
      uptime: this.state === 'running' && this.startedAt > 0
        ? Math.floor((Date.now() - this.startedAt) / 1000)
        : null,
      disabled: this.disabled,
    };
  }

  /**
   * Toggle persistent disable. When set true, dispose any running process
   * immediately and gate complete()/ensureRunning() so the next chat fails
   * fast rather than spawning a new vLLM. The caller (HTTP endpoint) is
   * responsible for persisting the flag to ~/.mcpctl/config.json.
   */
  setDisabled(value: boolean): void {
    if (this.disabled === value) return;
    this.disabled = value;
    if (value) {
      this.killProcess();
      this.clearIdleTimer();
      // Reset the error cooldown so re-enable starts clean.
      if (this.state === 'error') {
        this.state = 'stopped';
        this.lastError = null;
      }
    }
  }

  isDisabled(): boolean {
    return this.disabled;
  }

  /** Eagerly start vLLM so it's ready when the first complete() call arrives. */
  warmup(): void {
    if (this.state === 'stopped') {
      this.ensureRunning().catch((err) => {
        process.stderr.write(`[vllm-managed] warmup failed: ${(err as Error).message}\n`);
      });
    }
  }

  dispose(): void {
    this.killProcess();
    this.clearIdleTimer();
  }

  // --- Internal ---

  async ensureRunning(): Promise<void> {
    if (this.disabled) {
      throw new Error('vllm-managed provider is disabled — run `mcpctl provider enable <name>` to re-enable');
    }
    if (this.state === 'running' && this.process && !this.process.killed) {
      return;
    }

    if (this.state === 'starting' && this.startPromise) {
      return this.startPromise;
    }

    // Fast-fail if we recently errored — don't retry startup on every call
    if (this.state === 'error' && (Date.now() - this.errorAt) < ERROR_COOLDOWN_MS) {
      throw new Error(this.lastError ?? 'vLLM in error state (cooldown active)');
    }

    this.startPromise = this.doStart();
    try {
      await this.startPromise;
    } finally {
      this.startPromise = null;
    }
  }

  private async doStart(): Promise<void> {
    this.state = 'starting';
    this.lastError = null;

    const vllmBin = `${this.venvPath}/bin/vllm`;

    const args = [
      'serve', this.model,
      '--dtype', 'auto',
      '--max-model-len', String(this.maxModelLen),
      '--gpu-memory-utilization', String(this.gpuMemoryUtilization),
      '--port', String(this.port),
      ...this.extraArgs,
    ];

    const env: Record<string, string> = { ...process.env as Record<string, string> };
    // Ensure NVIDIA libraries are on the linker path
    const existingLd = env['LD_LIBRARY_PATH'] ?? '';
    env['LD_LIBRARY_PATH'] = existingLd
      ? `/usr/lib64/nvidia:${existingLd}`
      : '/usr/lib64/nvidia';
    env['VIRTUAL_ENV'] = this.venvPath;
    // Pin to NVIDIA GPU only — prevent vLLM from seeing AMD GPUs via ROCm/HIP
    env['CUDA_VISIBLE_DEVICES'] = '0';
    env['HIP_VISIBLE_DEVICES'] = '';
    env['ROCR_VISIBLE_DEVICES'] = '';

    try {
      const child = this.spawnFn(vllmBin, args, {
        env,
        stdio: ['ignore', 'pipe', 'pipe'],
        detached: false,
      });

      this.process = child;

      // Capture stderr for error reporting
      let stderrBuf = '';
      child.stderr?.on('data', (chunk: Buffer) => {
        stderrBuf += chunk.toString();
        // Keep only last 2KB
        if (stderrBuf.length > 2048) {
          stderrBuf = stderrBuf.slice(-2048);
        }
      });

      // Handle early exit
      const exitPromise = new Promise<number | null>((resolve) => {
        child.on('exit', (code) => resolve(code));
        child.on('error', (err) => {
          this.lastError = err.message;
          resolve(null);
        });
      });

      // Poll /v1/models until ready or timeout
      const ready = await this.waitForReady(exitPromise);

      if (!ready) {
        const exitCode = child.exitCode;
        if (exitCode !== null) {
          this.lastError = `vLLM exited with code ${exitCode}: ${stderrBuf.trim().slice(-200)}`;
        } else if (!this.lastError) {
          this.lastError = `vLLM startup timed out after ${STARTUP_TIMEOUT_MS / 1000}s`;
        }
        this.killProcess();
        this.state = 'error';
        this.errorAt = Date.now();
        throw new Error(this.lastError);
      }

      this.state = 'running';
      this.startedAt = Date.now();
      this.lastUsed = Date.now();

      // Create inner OpenAI-compatible provider pointed at local vLLM
      this.inner = new OpenAiProvider({
        apiKey: 'unused',
        baseUrl: `http://localhost:${this.port}`,
        defaultModel: this.model,
      });

      // Watch for unexpected exit
      child.on('exit', () => {
        if (this.state === 'running') {
          this.state = 'stopped';
          this.inner = null;
          this.process = null;
          this.startedAt = 0;
        }
      });

      this.resetIdleTimer();
    } catch (err) {
      if (this.state === 'starting') {
        this.state = 'error';
        this.errorAt = Date.now();
        this.lastError = (err as Error).message;
      }
      throw err;
    }
  }

  private async waitForReady(exitPromise: Promise<number | null>): Promise<boolean> {
    const deadline = Date.now() + STARTUP_TIMEOUT_MS;

    while (Date.now() < deadline) {
      // Check if process already exited
      const raceResult = await Promise.race([
        this.sleep(POLL_INTERVAL_MS).then(() => 'poll' as const),
        exitPromise.then(() => 'exited' as const),
      ]);

      if (raceResult === 'exited' && this.process?.exitCode !== null) {
        return false;
      }

      try {
        const ok = await this.healthCheckFn(this.port);
        if (ok) return true;
      } catch {
        // Not ready yet
      }
    }

    return false;
  }

  private resetIdleTimer(): void {
    this.clearIdleTimer();
    this.idleTimer = setInterval(() => {
      if (this.state === 'running' && Date.now() - this.lastUsed > this.idleTimeoutMs) {
        process.stderr.write(
          `[vllm-managed] vLLM stopped after ${Math.round(this.idleTimeoutMs / 60000)}min inactivity (GPU memory freed)\n`,
        );
        this.killProcess();
        this.state = 'stopped';
        this.inner = null;
        this.startedAt = 0;
        this.clearIdleTimer();
      }
    }, 30_000); // Check every 30 seconds
    // Unref so it doesn't keep the process alive
    if (this.idleTimer && typeof this.idleTimer === 'object' && 'unref' in this.idleTimer) {
      this.idleTimer.unref();
    }
  }

  private clearIdleTimer(): void {
    if (this.idleTimer) {
      clearInterval(this.idleTimer);
      this.idleTimer = null;
    }
  }

  private killProcess(): void {
    if (this.process && !this.process.killed) {
      this.process.kill('SIGTERM');
      // Force kill after 5s if still alive
      const p = this.process;
      setTimeout(() => {
        if (!p.killed) p.kill('SIGKILL');
      }, 5000).unref();
    }
    this.process = null;
    this.inner = null;
    this.startedAt = 0;
  }

  private sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
  }
}

/** Default health check: GET /v1/models on localhost:{port} */
function defaultHealthCheck(port: number): Promise<boolean> {
  return new Promise((resolve) => {
    const req = http.get(`http://localhost:${port}/v1/models`, { timeout: 3000 }, (res) => {
      res.resume();
      resolve(res.statusCode !== undefined && res.statusCode >= 200 && res.statusCode < 400);
    });
    req.on('error', () => resolve(false));
    req.on('timeout', () => { req.destroy(); resolve(false); });
  });
}
