import { env } from "cloudflare:workers";
import { withinSpendingLimit } from "../metering/quota";
import { numberVar } from "../limits";
import { all, first, run } from "../db";
import { completeJob, isCloudJob, runnerStub } from "./jobs";
import type { WorkflowJobRow, WorkflowRunRow } from "./jobs";
import { gateJob } from "./deployments";
import { forkRunsAllowed, hasForkEligibleRunner } from "./runner-policy";
import { recompileJob } from "./recompile";
import { CLOUD_RUNS_ON_VALUES, cloudRunsOnPlaceholders, runnerWeightSql } from "./runner-sizes";
import type { CompiledJob } from "./compile";
import type { Conclusion } from "~/lib/types";

export type JobState = {
  id: string;
  key: string;
  baseKey: string;
  needs: string[];
  runsOn: string[];
  status: string;
  conclusion: Conclusion | null;
  cloud: boolean;
  runIfOk: boolean;
  runIfFailed: boolean;
  row: WorkflowJobRow;
};

export type TickResult = { active: number; dispatched: number; deferred: number; gated: number; wakeAt: number | null };

const QUEUE_DEADLINE_MS = 24 * 60 * 60 * 1000;
const REVIEW_DEADLINE_MS = 72 * 60 * 60 * 1000;
const DAY_MS = 24 * 60 * 60 * 1000;
const MINUTE_MS = 60 * 1000;
const DEFAULT_JOBS_PER_OWNER = 4;
const DEFAULT_MINUTES_PER_DAY = 600;

type OwnerQuota = { ownerId: string; quotaExceeded: boolean; createdAt: number };

/**
 * Daily Container-minute budget per owner, weighted by runner size. Running and queued jobs are charged their
 * full timeout up front so an owner just under budget can't start a burst of long jobs.
 */
async function ownerQuota(runId: string): Promise<OwnerQuota | null> {
  const owner = await first<{ owner_id: string; created_at: number }>(
    "SELECT r.owner_id, w.created_at FROM workflow_runs w JOIN repos r ON r.id = w.repo_id WHERE w.id = ?",
    runId,
  );
  if (!owner) return null;
  const usage = await first<{ used_ms: number | null }>(
    `SELECT SUM(${runnerWeightSql("j.runs_on")} * CASE WHEN j.status = 'completed'
                     THEN COALESCE(j.completed_at - j.started_at, 0)
                     ELSE json_extract(j.spec, '$.timeoutMinutes') * ?1 END) AS used_ms
     FROM workflow_jobs j JOIN repos r ON r.id = j.repo_id
     WHERE r.owner_id = ?2 AND j.status IN ('completed', 'queued', 'in_progress')
       AND (j.status != 'completed' OR j.completed_at > ?3) AND j.runs_on IN (${cloudRunsOnPlaceholders(4)})`,
    MINUTE_MS,
    owner.owner_id,
    Date.now() - DAY_MS,
    ...CLOUD_RUNS_ON_VALUES,
  );
  const minutes = numberVar(env.CLOUD_MINUTES_PER_DAY, DEFAULT_MINUTES_PER_DAY);
  return {
    ownerId: owner.owner_id,
    quotaExceeded:
      (usage?.used_ms ?? 0) >= minutes * MINUTE_MS || !(await withinSpendingLimit(owner.owner_id, "compute.runner")),
    createdAt: owner.created_at,
  };
}

/** Atomically moves a cloud job to queued only while the owner is under its Container concurrency limit. */
async function claimCloudSlot(jobId: string, ownerId: string): Promise<boolean> {
  const result = await run(
    `UPDATE workflow_jobs SET status = 'queued', gate = NULL WHERE id = ?1 AND status = 'pending' AND (
       SELECT COUNT(*) FROM workflow_jobs j JOIN repos r ON r.id = j.repo_id
       WHERE r.owner_id = ?2 AND j.status IN ('queued', 'in_progress') AND j.runs_on IN (${cloudRunsOnPlaceholders(4)})
     ) < ?3`,
    jobId,
    ownerId,
    numberVar(env.CLOUD_JOBS_PER_OWNER, DEFAULT_JOBS_PER_OWNER),
    ...CLOUD_RUNS_ON_VALUES,
  );
  return result.meta.changes === 1;
}

function toState(row: WorkflowJobRow, spec: CompiledJob): JobState {
  return {
    id: row.id,
    key: row.job_key,
    baseKey: spec.baseKey,
    needs: spec.needs,
    runsOn: spec.runsOn,
    status: row.status,
    conclusion: row.conclusion,
    cloud: isCloudJob(spec.runsOn),
    runIfOk: spec.runIfOk,
    runIfFailed: spec.runIfFailed,
    row,
  };
}

export async function loadStates(runId: string): Promise<JobState[]> {
  const rows = await all<WorkflowJobRow>(
    "SELECT * FROM workflow_jobs WHERE run_id = ? ORDER BY created_at, job_key",
    runId,
  );
  return rows.map((row) => toState(row, JSON.parse(row.spec)));
}

function dependencyOutcome(job: JobState, jobs: JobState[]): "waiting" | "ok" | "failed" | "skipped" {
  const dependencies = jobs.filter((candidate) => job.needs.includes(candidate.baseKey));
  if (dependencies.some((dependency) => dependency.status !== "completed")) return "waiting";
  if (dependencies.some((dependency) => dependency.conclusion !== "success" && dependency.conclusion !== "skipped"))
    return "failed";
  return dependencies.some((dependency) => dependency.conclusion === "skipped") ? "skipped" : "ok";
}

async function settle(job: JobState, conclusion: Conclusion): Promise<void> {
  await completeJob(job.id, conclusion);
  job.status = "completed";
  job.conclusion = conclusion;
}

/** Resolves late-bound expressions and environment protection; returns null once the job was settled. */
async function prepare(runRow: WorkflowRunRow, job: JobState): Promise<JobState | null> {
  let prepared = job;
  if (job.needs.length > 0 || job.row.environment) {
    try {
      prepared = toState(job.row, await recompileJob(runRow, job.row));
    } catch (error) {
      console.error("Failed to compile job", job.id, error);
      await settle(job, "failure");
      return null;
    }
  }
  return prepared;
}

/** One scheduling pass: dispatches every job whose dependencies resolved, within owner quotas and gates. */
export async function scheduleTick(runId: string): Promise<TickResult> {
  const runRow = await first<WorkflowRunRow>("SELECT * FROM workflow_runs WHERE id = ?", runId);
  const jobs = await loadStates(runId);
  const result: TickResult = { active: 0, dispatched: 0, deferred: 0, gated: 0, wakeAt: null };
  if (!runRow) return result;
  const quota = jobs.some((job) => job.cloud && job.status === "pending") ? await ownerQuota(runId) : null;
  const pastDeadline = Date.now() - runRow.created_at > QUEUE_DEADLINE_MS;
  for (const pending of jobs.filter((candidate) => candidate.status === "pending")) {
    const outcome = dependencyOutcome(pending, jobs);
    if (outcome === "waiting") continue;
    const job = await prepare(runRow, pending);
    if (!job) continue;
    // Like GitHub, a skipped dependency skips its dependents unless their condition holds either way (`always()`).
    const shouldRun =
      outcome === "ok" ? job.runIfOk : outcome === "failed" ? job.runIfFailed : job.runIfOk && job.runIfFailed;
    if (!shouldRun) {
      await settle(pending, "skipped");
      continue;
    }
    if (
      !job.cloud &&
      (runRow.source_repo_id !== runRow.repo_id || runRow.untrusted === 1) &&
      !((await forkRunsAllowed(runRow.repo_id)) && (await hasForkEligibleRunner(runRow.repo_id, job.runsOn)))
    ) {
      // Self-hosted runners never execute fork or untrusted code unless an admin opted in and an eligible runner
      // exists, so the job fails instead of waiting forever.
      console.warn("Refusing self-hosted job for an untrusted run", job.id);
      await settle(pending, "failure");
      continue;
    }
    const gate = await gateJob(runRow, job.row);
    if (gate.state === "denied") {
      console.warn("Job gated out", job.id, gate.reason);
      await settle(pending, "failure");
      continue;
    }
    if (gate.state === "waiting") {
      if (Date.now() - runRow.created_at > REVIEW_DEADLINE_MS) {
        await settle(pending, "cancelled");
        continue;
      }
      if (job.row.gate !== "waiting") await run("UPDATE workflow_jobs SET gate = 'waiting' WHERE id = ?", job.id);
      result.gated++;
      if (gate.until) result.wakeAt = Math.min(result.wakeAt ?? gate.until, gate.until);
      continue;
    }
    if (job.cloud && (quota === null || quota.quotaExceeded || pastDeadline)) {
      await settle(pending, "cancelled");
      continue;
    }
    if (job.cloud && quota) {
      if (!(await claimCloudSlot(job.id, quota.ownerId))) {
        result.deferred++;
        continue;
      }
    } else {
      await run("UPDATE workflow_jobs SET status = 'queued', gate = NULL WHERE id = ? AND status = 'pending'", job.id);
    }
    pending.status = "queued";
    if (job.cloud) {
      try {
        await runnerStub(job.id, job.runsOn).startJob(job.id);
      } catch (error) {
        console.error("Failed to start cloud job", job.id, error);
        await completeJob(job.id, "failure");
      }
    }
    result.dispatched++;
  }
  const refreshed = await loadStates(runId);
  const active = refreshed.filter((job) => job.status !== "completed");
  const running = active.filter((job) => job.status === "queued" || job.status === "in_progress");
  result.active = active.length;
  if (active.length > 0 && running.length === 0 && result.deferred === 0 && result.gated === 0) {
    for (const job of active) await completeJob(job.id, "cancelled");
    result.active = 0;
  }
  return result;
}

export async function expireStaleJobs(runId: string): Promise<void> {
  const stale = await all<{ id: string }>(
    "SELECT id FROM workflow_jobs WHERE run_id = ? AND status = 'queued' AND created_at < ?",
    runId,
    Date.now() - QUEUE_DEADLINE_MS,
  );
  for (const job of stale) await completeJob(job.id, "cancelled");
}