import { z } from "zod";
import { buildControllerStatus } from "./status";
import type { ControllerStore } from "./store";

const METRIC_NAMES = [
  "build_offload_queue_depth",
  "build_offload_queue_oldest_age_seconds",
  "build_offload_queue_age_p95_seconds",
  "build_offload_fallback_lease_expired",
  "build_offload_artifact_cas_mismatch_total",
  "build_offload_incident_alert_state",
  "build_offload_exit_total",
  "build_offload_host_running_jobs",
  "build_offload_host_slots_free",
  "build_offload_host_slots_total",
  "build_offload_host_dispatch_accept",
  "build_offload_host_load",
  "build_offload_host_cores",
  "build_offload_host_builds_24h",
  "build_offload_host_sessions",
  "build_offload_host_remote_work",
  "build_offload_host_mem_used_bytes",
  "build_offload_host_mem_total_bytes",
  "build_offload_host_swap_used_bytes",
  "build_offload_host_swap_total_bytes",
  "build_offload_host_net_rx_bytes_per_second",
  "build_offload_host_net_tx_bytes_per_second",
  "build_offload_host_disk_free_bytes",
  "build_offload_host_disk_size_bytes",
  "build_offload_host_core_load_percent",
  "build_offload_host_cpu_temp_celsius",
  "build_offload_host_memory_stall_percent",
  "build_offload_host_tmp_used_bytes",
  "build_offload_host_tmp_size_bytes",
  "build_offload_host_work_slice_memory_bytes",
  "build_offload_host_work_slice_oom_kills_total",
  "build_offload_host_work_slice_pids",
] as const;

/** PSI stall windows exported per host: the box's own no-forward-progress signal. */
export const STALL_WINDOWS = ["some60", "full60", "full300"] as const;
export type StallWindow = typeof STALL_WINDOWS[number];

export type MetricName = typeof METRIC_NAMES[number];

const LABELS_BY_NAME: Record<MetricName, readonly string[]> = {
  build_offload_queue_depth: [],
  build_offload_queue_oldest_age_seconds: [],
  build_offload_queue_age_p95_seconds: [],
  build_offload_fallback_lease_expired: [],
  build_offload_artifact_cas_mismatch_total: [],
  build_offload_incident_alert_state: ["key"],
  build_offload_exit_total: ["host", "command", "code"],
  build_offload_host_running_jobs: ["host"],
  build_offload_host_slots_free: ["host"],
  build_offload_host_slots_total: ["host"],
  build_offload_host_dispatch_accept: ["host"],
  build_offload_host_load: ["host"],
  build_offload_host_cores: ["host"],
  build_offload_host_builds_24h: ["host"],
  build_offload_host_sessions: ["host"],
  build_offload_host_remote_work: ["host", "kind"],
  build_offload_host_mem_used_bytes: ["host"],
  build_offload_host_mem_total_bytes: ["host"],
  build_offload_host_swap_used_bytes: ["host"],
  build_offload_host_swap_total_bytes: ["host"],
  build_offload_host_net_rx_bytes_per_second: ["host"],
  build_offload_host_net_tx_bytes_per_second: ["host"],
  build_offload_host_disk_free_bytes: ["host", "mountpoint"],
  build_offload_host_disk_size_bytes: ["host", "mountpoint"],
  build_offload_host_core_load_percent: ["host", "core"],
  build_offload_host_cpu_temp_celsius: ["host", "sensor"],
  build_offload_host_memory_stall_percent: ["host", "window"],
  build_offload_host_tmp_used_bytes: ["host"],
  build_offload_host_tmp_size_bytes: ["host"],
  build_offload_host_work_slice_memory_bytes: ["host", "slice"],
  build_offload_host_work_slice_oom_kills_total: ["host", "slice"],
  build_offload_host_work_slice_pids: ["host", "slice"],
};

export const MetricSeriesSchema = z
  .object({
    name: z.enum(METRIC_NAMES),
    labels: z.record(z.string()),
    value: z.number().finite(),
  })
  .strict()
  .superRefine((sample, context) => {
    const actual = Object.keys(sample.labels).sort();
    const expected = [...LABELS_BY_NAME[sample.name]].sort();
    if (actual.join("\0") !== expected.join("\0")) {
      context.addIssue({
        code: z.ZodIssueCode.custom,
        message: `${sample.name} labels must be ${expected.join(",")}`,
      });
    }
    for (const label of expected) {
      if (sample.labels[label] === "") {
        context.addIssue({ code: z.ZodIssueCode.custom, message: `${label} must be non-empty` });
      }
    }
    if (sample.name === "build_offload_host_cpu_temp_celsius") {
      if (!(["pkg", "max", "crit"] as const).includes(sample.labels.sensor as "pkg" | "max" | "crit")) {
        context.addIssue({ code: z.ZodIssueCode.custom, message: "invalid temperature sensor" });
      }
    }
    if (sample.name === "build_offload_host_memory_stall_percent") {
      if (!STALL_WINDOWS.includes(sample.labels.window as StallWindow)) {
        context.addIssue({ code: z.ZodIssueCode.custom, message: "invalid stall window" });
      }
    }
    if (sample.name === "build_offload_host_core_load_percent" && !/^\d+$/.test(sample.labels.core ?? "")) {
      context.addIssue({ code: z.ZodIssueCode.custom, message: "core must be an integer index" });
    }
    if (sample.name === "build_offload_exit_total" && !["126", "127"].includes(sample.labels.code ?? "")) {
      context.addIssue({ code: z.ZodIssueCode.custom, message: "exit code must be 126 or 127" });
    }
  });

export type MetricSeries = z.infer<typeof MetricSeriesSchema>;

/** Placement paths that put work on a box; the metric label `kind` carries these keys. */
export const REMOTE_WORK_KINDS = {
  agentSeats: "agent_seat",
  remoteBuildJobs: "remote_build_job",
  offloadShells: "offload_shell",
} as const;

export type RemoteWorkCounts = { [K in keyof typeof REMOTE_WORK_KINDS]: number };

/** No series at all when the box published no count, so a gap never scrapes as a zero. */
function remoteWorkSeries(
  labels: Record<string, string>,
  counts: RemoteWorkCounts | null | undefined,
): MetricSeries[] {
  if (!counts) return [];
  return (Object.entries(REMOTE_WORK_KINDS) as Array<[keyof RemoteWorkCounts, string]>).map(
    ([field, kind]) => series("build_offload_host_remote_work", { ...labels, kind }, counts[field]),
  );
}

export interface HostMachineMetrics {
  host: string;
  /** Live sessions the host's own ledger classifier reports; null when the host has no ledger to read. */
  sessions: number | null;
  /**
   * Work units the box observes executing on itself, one entry per placement path. Null when
   * the box publishes no such count — a box that did not look must not read as an idle box.
   */
  remoteWork: RemoteWorkCounts | null;
  load: number;
  cores: number;
  memUsedBytes: number;
  memTotalBytes: number;
  swapUsedBytes: number;
  swapTotalBytes: number;
  netRxBytesPerSecond: number;
  netTxBytesPerSecond: number;
  disks: Array<{ mountpoint: string; freeBytes: number; sizeBytes: number }>;
  coreLoadPercent: number[];
  temperatures: { pkg: number; max: number; crit: number };
  /**
   * Box-side equivalents of the laptop's mem-guard and tmpjail signals: kernel facts about
   * stall, tmpfs fill and the agent/build slices. Nothing here authorises a kill.
   */
  guard: HostGuardMetrics;
}

export interface HostGuardMetrics {
  memoryStallSome60: number;
  memoryStallFull60: number;
  memoryStallFull300: number;
  tmpUsedBytes: number;
  tmpSizeBytes: number;
  workSlices: Array<{
    slice: string;
    memoryBytes: number;
    oomKillTotal: number;
    pidsCurrent: number;
  }>;
}

export interface TelemetryGap {
  host: string;
  reason: string;
  // ISO timestamp of the failed read; null when the host has never been sampled.
  at: string | null;
}

export interface HostMetricsSample {
  metrics: HostMachineMetrics[];
  gaps: TelemetryGap[];
}

export interface HostMetricsSampler {
  sample(hosts: readonly string[]): HostMetricsSample;
}

export interface TemperatureObserver {
  observe(temperature: { host: string; pkg: number; crit: number }): void;
}

const EMPTY_SAMPLER: HostMetricsSampler = { sample: () => ({ metrics: [], gaps: [] }) };

interface PrometheusMatcher {
  label: string;
  operator: "=" | "!=" | "=~";
  value: string;
}

export class MetricsRegistry {
  private artifactCasMismatch = 0;
  private readonly exits = new Map<string, number>();
  private readonly builds24h = new Map<string, number>();
  private collectCache: { at: number; samples: MetricSeries[] } | null = null;
  private telemetryGaps: TelemetryGap[] = [];
  private readonly collectCacheMs: number;

  constructor(
    private readonly store: ControllerStore,
    private readonly sampler: HostMetricsSampler = EMPTY_SAMPLER,
    private readonly now: () => number = () => Date.now(),
    private readonly temperatureObserver?: TemperatureObserver,
    collectCacheMs = 2_000,
  ) {
    this.collectCacheMs = collectCacheMs;
  }

  // Why a host has no machine numbers right now. Empty when every host reported.
  listTelemetryGaps(): readonly TelemetryGap[] {
    this.collect();
    return this.telemetryGaps;
  }

  incrementArtifactCasMismatch(count = 1): void {
    this.artifactCasMismatch += count;
  }

  recordExit(host: string, command: string, code: 126 | 127, count = 1): void {
    const key = JSON.stringify([host, command, String(code)]);
    this.exits.set(key, (this.exits.get(key) ?? 0) + count);
  }

  setHostBuilds24h(host: string, count: number): void {
    this.builds24h.set(host, count);
  }

  collect(): MetricSeries[] {
    const now = this.now();
    if (this.collectCache && now - this.collectCache.at < this.collectCacheMs) {
      return this.collectCache.samples;
    }
    const samples = this.collectUncached();
    this.collectCache = { at: now, samples };
    return samples;
  }

  private collectUncached(): MetricSeries[] {
    const hosts = this.store.listHosts().map(({ hostname }) => hostname);
    const members = new Set(hosts);
    const { metrics: machineMetrics, gaps: telemetryGaps } = this.sampler.sample(hosts);
    this.telemetryGaps = telemetryGaps;
    for (const host of machineMetrics) {
      if (!members.has(host.host)) continue;
      this.temperatureObserver?.observe({
        host: host.host,
        pkg: host.temperatures.pkg,
        crit: host.temperatures.crit,
      });
    }
    const status = buildControllerStatus(this.store, undefined, this.now);
    const samples: MetricSeries[] = [
      series("build_offload_queue_depth", {}, status.queue.depth),
      series("build_offload_queue_oldest_age_seconds", {}, status.queue.oldestAgeSeconds),
      series("build_offload_queue_age_p95_seconds", {}, status.queue.p95AgeSeconds),
      series("build_offload_fallback_lease_expired", {}, status.lease.expired ? 1 : 0),
      series("build_offload_artifact_cas_mismatch_total", {}, this.artifactCasMismatch),
    ];

    for (const incident of this.store.listIncidents()) {
      samples.push(series(
        "build_offload_incident_alert_state",
        { key: incident.key },
        incident.state === "open" ? 1 : 0,
      ));
    }

    const statusHosts = Object.keys(status.hosts);
    for (const hostname of statusHosts) {
      const host = this.store.getHost(hostname);
      if (!host) continue;
      const labels = { host: hostname };
      samples.push(
        series("build_offload_host_running_jobs", labels, host.runningJobs),
        series("build_offload_host_slots_free", labels, Math.max(0, host.slotsTotal - host.slotsUsed)),
        series("build_offload_host_slots_total", labels, host.slotsTotal),
        series(
          "build_offload_host_dispatch_accept",
          labels,
          host.state === "available" && host.capabilityOk && !host.dispatchPaused ? 1 : 0,
        ),
      );
      const builds24h = this.builds24h.get(hostname);
      if (builds24h !== undefined) {
        samples.push(series("build_offload_host_builds_24h", labels, builds24h));
      }
    }

    for (const host of machineMetrics) {
      if (!members.has(host.host)) continue;
      const labels = { host: host.host };
      samples.push(
        ...(host.sessions === null ? [] : [series("build_offload_host_sessions", labels, host.sessions)]),
        ...remoteWorkSeries(labels, host.remoteWork),
        series("build_offload_host_load", labels, host.load),
        series("build_offload_host_cores", labels, host.cores),
        series("build_offload_host_mem_used_bytes", labels, host.memUsedBytes),
        series("build_offload_host_mem_total_bytes", labels, host.memTotalBytes),
        series("build_offload_host_swap_used_bytes", labels, host.swapUsedBytes),
        series("build_offload_host_swap_total_bytes", labels, host.swapTotalBytes),
        series("build_offload_host_net_rx_bytes_per_second", labels, host.netRxBytesPerSecond),
        series("build_offload_host_net_tx_bytes_per_second", labels, host.netTxBytesPerSecond),
      );
      for (const disk of host.disks) {
        const diskLabels = { host: host.host, mountpoint: disk.mountpoint };
        samples.push(
          series("build_offload_host_disk_free_bytes", diskLabels, disk.freeBytes),
          series("build_offload_host_disk_size_bytes", diskLabels, disk.sizeBytes),
        );
      }
      host.coreLoadPercent.forEach((value, core) => {
        samples.push(series("build_offload_host_core_load_percent", { host: host.host, core: String(core) }, value));
      });
      for (const sensor of ["pkg", "max", "crit"] as const) {
        samples.push(series("build_offload_host_cpu_temp_celsius", { host: host.host, sensor }, host.temperatures[sensor]));
      }
      const stallByWindow: Record<StallWindow, number> = {
        some60: host.guard.memoryStallSome60,
        full60: host.guard.memoryStallFull60,
        full300: host.guard.memoryStallFull300,
      };
      for (const window of STALL_WINDOWS) {
        samples.push(series("build_offload_host_memory_stall_percent", { host: host.host, window }, stallByWindow[window]));
      }
      samples.push(
        series("build_offload_host_tmp_used_bytes", labels, host.guard.tmpUsedBytes),
        series("build_offload_host_tmp_size_bytes", labels, host.guard.tmpSizeBytes),
      );
      for (const slice of host.guard.workSlices) {
        const sliceLabels = { host: host.host, slice: slice.slice };
        samples.push(
          series("build_offload_host_work_slice_memory_bytes", sliceLabels, slice.memoryBytes),
          series("build_offload_host_work_slice_oom_kills_total", sliceLabels, slice.oomKillTotal),
          series("build_offload_host_work_slice_pids", sliceLabels, slice.pidsCurrent),
        );
      }
    }

    for (const [key, value] of this.exits) {
      const [host, command, code] = JSON.parse(key) as [string, string, string];
      if (members.has(host)) {
        samples.push(series("build_offload_exit_total", { host, command, code }, value));
      }
    }
    return samples;
  }

  exposition(): string {
    return `${this.collect().map(formatSample).join("\n")}\n`;
  }

  query(expression: string): {
    status: "success";
    data: {
      resultType: "vector";
      result: Array<{ metric: Record<string, string>; value: [number, string] }>;
    };
  } {
    const { name, matchers } = parseSelector(expression);
    const timestamp = this.now() / 1000;
    const result = this.collect()
      .filter((sample) => sample.name === name && matchers.every((matcher) => matches(sample.labels, matcher)))
      .map((sample) => ({
        metric: { __name__: sample.name, ...sample.labels },
        value: [timestamp, String(sample.value)] as [number, string],
      }));
    return { status: "success", data: { resultType: "vector", result } };
  }
}

function series(name: MetricName, labels: Record<string, string>, value: number): MetricSeries {
  return MetricSeriesSchema.parse({ name, labels, value });
}

function formatSample(sample: MetricSeries): string {
  const labels = Object.entries(sample.labels)
    .map(([key, value]) => `${key}="${escapeLabel(value)}"`)
    .join(",");
  return `${sample.name}${labels === "" ? "" : `{${labels}}`} ${sample.value}`;
}

function escapeLabel(value: string): string {
  return value.replaceAll("\\", "\\\\").replaceAll("\n", "\\n").replaceAll('"', '\\"');
}

function parseSelector(expression: string): { name: MetricName; matchers: PrometheusMatcher[] } {
  const parsed = /^([a-zA-Z_:][a-zA-Z0-9_:]*)(?:\{(.*)\})?$/.exec(expression.trim());
  if (!parsed || !METRIC_NAMES.includes(parsed[1] as MetricName)) {
    throw new Error("unsupported Prometheus query");
  }
  const matcherSource = parsed[2]?.trim();
  if (!matcherSource) return { name: parsed[1] as MetricName, matchers: [] };
  const matchers = matcherSource.split(",").map((part) => {
    const matcher = /^([a-zA-Z_][a-zA-Z0-9_]*)\s*(=~|!=|=)\s*("(?:\\.|[^"])*")$/.exec(part.trim());
    if (!matcher) throw new Error("unsupported Prometheus matcher");
    return {
      label: matcher[1] ?? "",
      operator: matcher[2] as PrometheusMatcher["operator"],
      value: JSON.parse(matcher[3] ?? '""') as string,
    };
  });
  return { name: parsed[1] as MetricName, matchers };
}

function matches(labels: Record<string, string>, matcher: PrometheusMatcher): boolean {
  const actual = labels[matcher.label] ?? "";
  if (matcher.operator === "=") return actual === matcher.value;
  if (matcher.operator === "!=") return actual !== matcher.value;
  return new RegExp(`^(?:${matcher.value})$`).test(actual);
}