import { existsSync } from "node:fs";
import { arch, cpus, freemem, platform, totalmem } from "node:os";
import { Effect } from "effect";
import type { GpuInfo, RuntimeGpuMonitoringTool } from "../../models/types";
import { runCommandAsyncEffect } from "../../../core/command";
import { getGpuInfoFromAmdSmi, getGpuInfoFromRocmSmi } from "./amd-gpu";
import { getGpuInfoFromIntelSysfs } from "./intel-gpu";
import { resolveRocmSmiTool } from "./rocm-info";
import {
  resolveAmdSmiBinary,
  resolveForcedGpuMonitoringTool,
  resolveNvidiaSmiBinary,
  resolveRocmSmiBinary,
} from "./smi-tools";

const NVIDIA_SMI_GPU_FIELDS = [
  "uuid",
  "pci.bus_id",
  "name",
  "memory.total",
  "memory.used",
  "memory.free",
  "utilization.gpu",
  "temperature.gpu",
  "power.draw",
  "power.limit",
] as const;

const NVIDIA_SMI_SNAPSHOT_QUERY = [...NVIDIA_SMI_GPU_FIELDS, "driver_version"].join(",");
const NVIDIA_SMI_ARGS = [
  `--query-gpu=${NVIDIA_SMI_SNAPSHOT_QUERY}`,
  "--format=csv,noheader,nounits",
];
const NVIDIA_SMI_TIMEOUT_MS = 5_000;

const parseNvidiaSmiGpuLine = (line: string, index: number): GpuInfo => {
  const parts = line.split(",").map((value) => value.trim());
  const [
    rawUuid,
    rawPciBusId,
    rawName,
    memoryTotal,
    memoryUsed,
    memoryFree,
    utilization,
    temperature,
    powerDraw,
    powerLimit,
  ] = parts;
  const name = rawName ?? "Unknown";
  const identity = (value: string | undefined): string | undefined => {
    if (!value || /^(?:N\/A|\[Not Supported\])$/i.test(value)) return undefined;
    return value;
  };
  const uuid = identity(rawUuid);
  const pciBusId = identity(rawPciBusId);
  const toFiniteNumber = (value: string | undefined): number => {
    const parsed = Number(value ?? 0);
    return Number.isFinite(parsed) ? parsed : 0;
  };
  const toMb = (megabytes: string | undefined): number =>
    Math.max(0, Math.round(toFiniteNumber(megabytes)));
  const reportedTotalMb = toMb(memoryTotal);
  const isUnifiedMemoryNvidia = reportedTotalMb === 0 && /\b(?:GB10|Grace)\b/i.test(name);
  const fallbackTotalMb = isUnifiedMemoryNvidia ? Math.round(totalmem() / 1024 / 1024) : 0;
  const fallbackFreeMb = isUnifiedMemoryNvidia ? Math.round(freemem() / 1024 / 1024) : 0;
  const fallbackUsedMb = Math.max(0, fallbackTotalMb - fallbackFreeMb);
  const memoryTotalMb = reportedTotalMb || fallbackTotalMb;
  const memoryUsedMb = toMb(memoryUsed) || fallbackUsedMb;
  const memoryFreeMb = toMb(memoryFree) || fallbackFreeMb;
  return {
    ...(uuid ? { uuid } : {}),
    ...(pciBusId ? { pci_bus_id: pciBusId } : {}),
    index,
    name,
    memory_total_mb: memoryTotalMb,
    memory_used_mb: memoryUsedMb,
    memory_free_mb: memoryFreeMb,
    utilization_pct: toFiniteNumber(utilization),
    temp_c: toFiniteNumber(temperature),
    power_draw: toFiniteNumber(powerDraw),
    power_limit: toFiniteNumber(powerLimit),
  };
};

const splitSmiLines = (stdout: string): string[] =>
  stdout
    .trim()
    .split("\n")
    .map((line) => line.trim())
    .filter(Boolean);

const parseNvidiaSmiGpuOutput = (stdout: string): GpuInfo[] =>
  splitSmiLines(stdout).map(parseNvidiaSmiGpuLine);

const parseNvidiaSmiDriverVersion = (stdout: string): string | null => {
  const firstLine = splitSmiLines(stdout)[0];
  if (!firstLine) return null;
  const driver = firstLine.split(",")[NVIDIA_SMI_GPU_FIELDS.length]?.trim();
  return driver || null;
};

export type NvidiaSmiSnapshot = {
  available: boolean;
  gpus: GpuInfo[];
  driverVersion: string | null;
};

export const queryNvidiaSmiSnapshot = (): Effect.Effect<NvidiaSmiSnapshot | null> => {
  const nvidiaSmi = resolveNvidiaSmiBinary();
  if (!nvidiaSmi) return Effect.succeed(null);
  return runCommandAsyncEffect(nvidiaSmi, NVIDIA_SMI_ARGS, {
    timeoutMs: NVIDIA_SMI_TIMEOUT_MS,
  }).pipe(
    Effect.map((result) => {
      if (result.status !== 0 || !result.stdout) {
        return { available: result.status === 0, gpus: [], driverVersion: null };
      }
      return {
        available: true,
        gpus: parseNvidiaSmiGpuOutput(result.stdout),
        driverVersion: parseNvidiaSmiDriverVersion(result.stdout),
      };
    }),
    Effect.catch(() => Effect.succeed({ available: false, gpus: [], driverVersion: null })),
  );
};

export const getGpuInfoFromNvidiaSmi = (): Effect.Effect<GpuInfo[]> =>
  queryNvidiaSmiSnapshot().pipe(Effect.map((snapshot) => snapshot?.gpus ?? []));

export const detectGpuMonitoringTool = (): Effect.Effect<RuntimeGpuMonitoringTool | null> =>
  Effect.gen(function* () {
    const forced = resolveForcedGpuMonitoringTool();
    if (forced) return forced;
    if (resolveNvidiaSmiBinary()) return "nvidia-smi";
    const rocmTool = resolveRocmSmiTool();
    if (rocmTool) return rocmTool;
    if ((yield* getGpuInfoFromIntelSysfs()).length > 0) return "intel-sysfs";
    return null;
  });

let warnedNoGpuTooling = false;

const warnNoGpuToolingOnce = (): void => {
  if (warnedNoGpuTooling) return;
  warnedNoGpuTooling = true;
  const attempted = [
    `nvidia-smi=${resolveNvidiaSmiBinary() ? "found" : "not found"}`,
    `amd-smi=${resolveAmdSmiBinary() ? "found" : "not found"}`,
    `rocm-smi=${resolveRocmSmiBinary() ? "found" : "not found"}`,
    `intel-sysfs=${existsSync("/sys/bus/pci/devices") ? "no compute GPUs" : "unavailable"}`,
  ].join(" ");
  console.warn(`No GPUs reported by any monitoring tool; attempted: ${attempted}`);
};

const collectGpuInfo = (): Effect.Effect<GpuInfo[]> =>
  Effect.gen(function* () {
    const forced = resolveForcedGpuMonitoringTool();
    if (forced === "nvidia-smi") {
      return yield* getGpuInfoFromNvidiaSmi();
    }
    if (forced === "amd-smi") {
      return yield* getGpuInfoFromAmdSmi();
    }
    if (forced === "rocm-smi") {
      return yield* getGpuInfoFromRocmSmi();
    }
    if (forced === "intel-sysfs") {
      return yield* getGpuInfoFromIntelSysfs();
    }

    const nvidia = yield* getGpuInfoFromNvidiaSmi();
    if (nvidia.length > 0) {
      return nvidia;
    }

    const rocmTool = resolveRocmSmiTool();
    if (rocmTool === "amd-smi") {
      const amd = yield* getGpuInfoFromAmdSmi();
      if (amd.length > 0) return amd;
      return yield* getGpuInfoFromRocmSmi();
    }
    if (rocmTool === "rocm-smi") {
      const rocm = yield* getGpuInfoFromRocmSmi();
      if (rocm.length > 0) return rocm;
      return yield* getGpuInfoFromAmdSmi();
    }

    const intel = yield* getGpuInfoFromIntelSysfs();
    if (intel.length > 0) {
      return intel;
    }

    if (platform() === "darwin" && arch() === "arm64") {
      const cpuName = cpus()[0]?.model?.trim() || "Apple Silicon";
      const memoryTotalMb = Math.round(totalmem() / 1024 / 1024);
      return [
        {
          id: "apple-metal-0",
          index: 0,
          name: `${cpuName} GPU`,
          memory_total_mb: memoryTotalMb,
          memory_used_mb: 0,
          memory_free_mb: memoryTotalMb,
          utilization_pct: 0,
          temp_c: 0,
          power_draw: 0,
          power_limit: 0,
          memory_shared: true,
          memory_usage_available: false,
          utilization_available: false,
          temperature_available: false,
          power_available: false,
        },
      ];
    }

    return [];
  });

export const getGpuInfo = (): Effect.Effect<GpuInfo[]> =>
  Effect.gen(function* () {
    const gpus = yield* collectGpuInfo();
    if (gpus.length === 0) {
      yield* Effect.sync(warnNoGpuToolingOnce);
    }
    return gpus;
  });