import { env } from "cloudflare:workers";
import { ringJobDoorbells } from "./job-doorbell";
import { assertWithinQuota, QuotaExceededError } from "../metering/quota";
import { isScratchRepo } from "../sites/config";
import { numberVar } from "../limits";
import { repoDb } from "../scoped-db";
import { ownerRepoIds, repoOwnerId, runDb } from "./scope";
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 { holdJobMinutes, releaseJobHold } from "./job-hold";
import { CLOUD_RUNS_ON_VALUES, cloudRunnerSize, 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[];
  timeoutMinutes: number;
  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;

/** `usedMs`: finished work today; `heldMs`: full weighted timeouts of the owner's other queued/running jobs. */
type OwnerQuota = {
  ownerId: string;
  quotaExceeded: boolean;
  usedMs: number;
  heldMs: number;
  budgetMs: number;
  createdAt: number;
};

/**
 * Daily Container-minute budget per owner, weighted by runner size. Queued and running jobs are charged their full
 * timeout up front, so an owner just under budget can't start a burst of long jobs: once their charge fills the
 * budget, further jobs wait for them. A job with nothing else queued or running always starts, so an oversized
 * timeout can't wait forever.
 */
export function budgetBusy(quota: OwnerQuota): boolean {
  return quota.heldMs > 0 && quota.usedMs + quota.heldMs >= quota.budgetMs;
}

/** Charges a job just dispatched in this pass, so later jobs in the same pass see it. */
export function chargeDispatched(quota: OwnerQuota, runsOn: string[], timeoutMinutes: number): void {
  quota.heldMs += (cloudRunnerSize(runsOn)?.weight ?? 1) * timeoutMinutes * MINUTE_MS;
}

/** The job being decided is still `pending`, so it is never part of `heldMs`. */
export async function ownerQuota(runId: string): Promise<OwnerQuota | null> {
  const runRow = await (
    await runDb(runId)
  )?.first<{ repo_id: string; created_at: number }>(
    "SELECT repo_id, created_at FROM workflow_runs WHERE id = ?",
    runId,
  );
  const ownerId = runRow ? await repoOwnerId(runRow.repo_id) : null;
  if (!runRow || !ownerId) return null;
  const owner = { owner_id: ownerId, created_at: runRow.created_at };
  const perRepo = await Promise.all(
    (await ownerRepoIds(ownerId)).map((repoId) =>
      repoDb(repoId).first<{ used_ms: number | null; held_ms: number | null }>(
        `SELECT SUM(${runnerWeightSql("j.runs_on")} * CASE WHEN j.status = 'completed'
                         THEN COALESCE(j.completed_at - j.started_at, 0) ELSE 0 END) AS used_ms,
                SUM(${runnerWeightSql("j.runs_on")} * CASE WHEN j.status IN ('queued', 'in_progress')
                         THEN json_extract(j.spec, '$.timeoutMinutes') * ?1 ELSE 0 END) AS held_ms
         FROM workflow_jobs j
         WHERE j.status IN ('completed', 'queued', 'in_progress')
           AND (j.status != 'completed' OR j.completed_at > ?2) AND j.runs_on IN (${cloudRunsOnPlaceholders(3)})`,
        MINUTE_MS,
        Date.now() - DAY_MS,
        ...CLOUD_RUNS_ON_VALUES,
      ),
    ),
  );
  const usage = {
    used_ms: perRepo.reduce((sum, row) => sum + (row?.used_ms ?? 0), 0),
    held_ms: perRepo.reduce((sum, row) => sum + (row?.held_ms ?? 0), 0),
  };
  const budgetMs = numberVar(env.CLOUD_MINUTES_PER_DAY, DEFAULT_MINUTES_PER_DAY) * MINUTE_MS;
  const usedMs = usage?.used_ms ?? 0;
  return {
    ownerId: owner.owner_id,
    quotaExceeded: isScratchRepo(owner) || usedMs >= budgetMs || !(await runnerQuotaLeft(owner.owner_id)),
    usedMs,
    heldMs: usage?.held_ms ?? 0,
    budgetMs,
    createdAt: owner.created_at,
  };
}

/** Cloud minutes left this month: the plan's hard cap (free) or the spending limit (paid). */
async function runnerQuotaLeft(ownerId: string): Promise<boolean> {
  try {
    await assertWithinQuota(ownerId, "compute.runner");
    return true;
  } catch (error) {
    if (error instanceof QuotaExceededError) return false;
    throw error;
  }
}

/** Queued and running cloud jobs of an owner, across its repositories (optionally leaving one out). */
async function ownerActiveCloudJobs(ownerId: string, exceptRepoId: string | null = null): Promise<number> {
  const counts = await Promise.all(
    (await ownerRepoIds(ownerId))
      .filter((repoId) => repoId !== exceptRepoId)
      .map((repoId) =>
        repoDb(repoId).first<{ count: number }>(
          `SELECT COUNT(*) AS count FROM workflow_jobs
           WHERE status IN ('queued', 'in_progress') AND runs_on IN (${cloudRunsOnPlaceholders(1)})`,
          ...CLOUD_RUNS_ON_VALUES,
        ),
      ),
  );
  return counts.reduce((sum, row) => sum + (row?.count ?? 0), 0);
}

/**
 * Moves a cloud job to queued only while the owner is under its Container concurrency limit. Atomic within the
 * job's repository; the owner's other repositories are counted just before.
 */
async function claimCloudSlot(repoId: string, jobId: string, ownerId: string): Promise<boolean> {
  const elsewhere = await ownerActiveCloudJobs(ownerId, repoId);
  const result = await repoDb(repoId).run(
    `UPDATE workflow_jobs SET status = 'queued', gate = NULL WHERE id = ?1 AND status = 'pending' AND (
       SELECT COUNT(*) FROM workflow_jobs j
       WHERE j.status IN ('queued', 'in_progress') AND j.runs_on IN (${cloudRunsOnPlaceholders(4)})
     ) + ?2 < ?3`,
    jobId,
    elsewhere,
    numberVar(env.CLOUD_JOBS_PER_OWNER, DEFAULT_JOBS_PER_OWNER),
    ...CLOUD_RUNS_ON_VALUES,
  );
  return result.meta.changes === 1;
}

async function ownerHasActiveCloudJobs(ownerId: string): Promise<boolean> {
  return (await ownerActiveCloudJobs(ownerId)) > 0;
}

function toState(row: WorkflowJobRow, spec: CompiledJob): JobState {
  return {
    id: row.id,
    key: row.job_key,
    baseKey: spec.baseKey,
    needs: spec.needs,
    runsOn: spec.runsOn,
    timeoutMinutes: spec.timeoutMinutes,
    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 db = await runDb(runId);
  if (!db) return [];
  const rows = await db.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 db = await runDb(runId);
  const runRow = (await db?.first<WorkflowRunRow>("SELECT * FROM workflow_runs WHERE id = ?", runId)) ?? null;
  const jobs = await loadStates(runId);
  const result: TickResult = { active: 0, dispatched: 0, deferred: 0, gated: 0, wakeAt: null };
  if (!runRow || !db) 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;
  let queuedSelfHosted = false;
  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 db.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 && budgetBusy(quota)) {
      // Other jobs' up-front timeout charges are in the way; they shrink to real usage when those jobs finish.
      result.deferred++;
      continue;
    }
    if (job.cloud && quota) {
      const hold = { jobId: job.id, ownerId: quota.ownerId, runsOn: job.runsOn, timeoutMinutes: job.timeoutMinutes };
      if (!(await holdJobMinutes(hold))) {
        // Other jobs' holds may be what's in the way; they shrink to real usage when those jobs finish.
        if (await ownerHasActiveCloudJobs(quota.ownerId)) {
          result.deferred++;
        } else {
          await settle(pending, "cancelled");
        }
        continue;
      }
      let claimed = false;
      try {
        claimed = await claimCloudSlot(runRow.repo_id, job.id, quota.ownerId);
      } finally {
        if (!claimed) await releaseJobHold(job.id, quota.ownerId);
      }
      if (!claimed) {
        result.deferred++;
        continue;
      }
      chargeDispatched(quota, job.runsOn, job.timeoutMinutes);
    } else {
      await db.run(
        "UPDATE workflow_jobs SET status = 'queued', gate = NULL WHERE id = ? AND status = 'pending'",
        job.id,
      );
      queuedSelfHosted = true;
    }
    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++;
  }
  if (queuedSelfHosted) {
    const ownerId = await repoOwnerId(runRow.repo_id);
    if (ownerId) await ringJobDoorbells(runRow.repo_id, ownerId);
  }
  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 db = await runDb(runId);
  if (!db) return;
  const stale = await db.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");
}