import { createHash } from "node:crypto";
import {
  artifacts,
  checkpoints,
  type Db,
  jobLeases,
  paperSequences,
  projectInstructionVersions,
  qualityGateResults,
  researchProjects,
  researchRuns,
  runAttempts,
  stageEvents,
  stageExecutions,
  type Tx,
} from "@autoscholar/db";
import { effectiveAutoPublish, QUALITY_POLICY_VERSION } from "@autoscholar/research-contract";
import {
  ACTIVE_RUN_STATUSES,
  type LaunchRunRequest,
  RUN_TRANSITIONS,
  type RunEventInput,
  RunManifest,
  type RunStatus,
} from "@autoscholar/schemas";
import { and, asc, desc, eq, gt, inArray, sql } from "drizzle-orm";
import { audit } from "./audit.js";
import { getEnv } from "./env.js";
import { AppError, conflict, invalid, notFound } from "./errors.js";

export type RunRow = typeof researchRuns.$inferSelect;

export function canTransition(from: RunStatus, to: RunStatus): boolean {
  return RUN_TRANSITIONS[from]?.includes(to) ?? false;
}

/** Appends an event with a per-run monotonically increasing seq. Must be called inside a transaction holding the run row. */
export async function appendEvent(
  tx: Tx,
  runId: string,
  ev: RunEventInput & { attemptId?: string | null; workerId?: string | null },
): Promise<number> {
  const [r] = await tx
    .update(researchRuns)
    .set({ lastEventSeq: sql`${researchRuns.lastEventSeq} + 1`, updatedAt: new Date() })
    .where(eq(researchRuns.id, runId))
    .returning({ seq: researchRuns.lastEventSeq });
  if (!r) throw notFound("run");
  await tx.insert(stageEvents).values({
    runId,
    attemptId: ev.attemptId ?? null,
    workerId: ev.workerId ?? null,
    seq: r.seq,
    stage: ev.stage ?? null,
    level: ev.level ?? "info",
    kind: ev.kind,
    message: ev.message.slice(0, 4000),
    data: ev.data ?? null,
    at: new Date(ev.at),
  });
  return r.seq;
}

/** Atomic, validated state transition with ledger entry. Locks the run row. */
export async function transition(
  tx: Tx,
  runId: string,
  to: RunStatus,
  opts: {
    reason?: string;
    actor?: string;
    data?: Record<string, unknown>;
    attemptId?: string | null;
    workerId?: string | null;
    expectFrom?: RunStatus[];
  } = {},
): Promise<RunRow> {
  const [run] = await tx.select().from(researchRuns).where(eq(researchRuns.id, runId)).for("update");
  if (!run) throw notFound("run");
  const from = run.status as RunStatus;
  if (opts.expectFrom && !opts.expectFrom.includes(from))
    throw conflict(`run is ${from}, expected ${opts.expectFrom.join("/")}`);
  if (from === to) return run;
  if (!canTransition(from, to)) throw conflict(`invalid transition ${from} -> ${to}`);
  const patch: Partial<typeof researchRuns.$inferInsert> = { status: to, updatedAt: new Date() };
  if (to === "queued") patch.queuedAt = new Date();
  if (to === "running" && !run.startedAt) patch.startedAt = new Date();
  if (["published", "cancelled", "failed"].includes(to)) patch.finishedAt = new Date();
  if (to === "failed" && opts.reason) patch.failureReason = opts.reason;
  const [updated] = await tx.update(researchRuns).set(patch).where(eq(researchRuns.id, runId)).returning();
  await appendEvent(tx, runId, {
    kind: "status.transition",
    level: to === "failed" ? "error" : "info",
    message: `${from} → ${to}${opts.reason ? `: ${opts.reason}` : ""}`,
    data: { from, to, actor: opts.actor ?? "system", ...(opts.data ?? {}) },
    at: new Date().toISOString(),
    attemptId: opts.attemptId ?? null,
    workerId: opts.workerId ?? null,
  });
  return updated!;
}

/**
 * Idempotent launch: same idempotencyKey → same run. Builds and freezes the immutable manifest server-side.
 * draft → validated → queued in one transaction.
 */
export async function launchRun(
  db: Db,
  req: LaunchRunRequest,
  actor: { id: string; ip?: string | null },
): Promise<{ run: RunRow; created: boolean }> {
  const env = getEnv();
  const existing = await db
    .select()
    .from(researchRuns)
    .where(eq(researchRuns.idempotencyKey, req.idempotencyKey))
    .limit(1);
  if (existing[0]) return { run: existing[0], created: false };
  const [iv] = await db
    .select({ iv: projectInstructionVersions, p: researchProjects })
    .from(projectInstructionVersions)
    .innerJoin(researchProjects, eq(projectInstructionVersions.projectId, researchProjects.id))
    .where(
      and(
        eq(projectInstructionVersions.id, req.instructionVersionId),
        eq(projectInstructionVersions.projectId, req.projectId),
      ),
    )
    .limit(1);
  if (!iv) throw notFound("instruction version for this project");
  const autoPub = effectiveAutoPublish(req.researchMode, req.discipline, req.autoPublish);
  const maxBudgetUsd = req.maxBudgetUsd ?? env.DEFAULT_MAX_BUDGET_USD;
  const maxRuntimeMinutes = req.maxRuntimeMinutes ?? env.DEFAULT_MAX_RUNTIME_MINUTES;
  if (maxBudgetUsd > 2000) throw invalid("budget cap exceeds platform maximum");
  const runId = crypto.randomUUID();
  const manifest = RunManifest.parse({
    schemaVersion: 1,
    runId,
    projectId: req.projectId,
    instructionVersionId: iv.iv.id,
    instructionVersion: iv.iv.version,
    instructionSha256: iv.iv.sha256,
    researchMode: req.researchMode,
    secondaryModes: req.secondaryModes,
    discipline: req.discipline,
    computeProfile: req.computeProfile,
    provider: "anthropic",
    model: req.model ?? env.ANTHROPIC_MODEL,
    reviewerModel: env.ANTHROPIC_REVIEWER_MODEL,
    maxBudgetUsd,
    maxRuntimeMinutes,
    maxAttempts: req.maxAttempts,
    autoPublish: autoPub.allowed,
    qualityPolicyVersion: QUALITY_POLICY_VERSION,
    createdAt: new Date().toISOString(),
    createdBy: actor.id,
    title: req.title,
  });
  try {
    const run = await db.transaction(async (tx) => {
      const year = new Date().getUTCFullYear();
      await tx.insert(paperSequences).values({ year, last: 0 }).onConflictDoNothing();
      const [seq] = await tx
        .update(paperSequences)
        .set({ last: sql`${paperSequences.last} + 1` })
        .where(eq(paperSequences.year, year))
        .returning({ last: paperSequences.last });
      manifest.paperNumber = `AS-WP-${year}-${String(seq!.last).padStart(4, "0")}`;
      const manifestSha256 = createHash("sha256").update(JSON.stringify(manifest)).digest("hex");
      const [row] = await tx
        .insert(researchRuns)
        .values({
          id: runId,
          projectId: req.projectId,
          instructionVersionId: iv.iv.id,
          idempotencyKey: req.idempotencyKey,
          title: req.title,
          status: "draft",
          researchMode: req.researchMode,
          secondaryModes: req.secondaryModes,
          discipline: req.discipline,
          computeProfile: req.computeProfile,
          model: manifest.model,
          reviewerModel: manifest.reviewerModel,
          maxBudgetUsd: String(maxBudgetUsd),
          maxRuntimeMinutes,
          maxAttempts: req.maxAttempts,
          autoPublish: autoPub.allowed,
          qualityPolicyVersion: QUALITY_POLICY_VERSION,
          manifest,
          manifestSha256,
          createdBy: actor.id,
        })
        .returning();
      await appendEvent(tx, runId, {
        kind: "run.created",
        message: `Run created from instruction v${iv.iv.version} (${iv.iv.sha256.slice(0, 12)})`,
        data: { manifestSha256, autoPublishReason: autoPub.reason ?? null },
        at: new Date().toISOString(),
      });
      await transition(tx, runId, "validated", { actor: `admin:${actor.id}` });
      await transition(tx, runId, "queued", { actor: `admin:${actor.id}` });
      await audit(
        tx,
        { type: "admin", id: actor.id, ip: actor.ip },
        "run.launch",
        { type: "run", id: runId },
        { idempotencyKey: req.idempotencyKey, manifestSha256 },
      );
      return row!;
    });
    return { run: { ...run, status: "queued" }, created: true };
  } catch (e) {
    const again = await db
      .select()
      .from(researchRuns)
      .where(eq(researchRuns.idempotencyKey, req.idempotencyKey))
      .limit(1);
    if (again[0]) return { run: again[0], created: false };
    throw e;
  }
}

export async function getRun(db: Db, runId: string) {
  const [run] = await db.select().from(researchRuns).where(eq(researchRuns.id, runId)).limit(1);
  if (!run) throw notFound("run");
  const attempts = await db
    .select()
    .from(runAttempts)
    .where(eq(runAttempts.runId, runId))
    .orderBy(asc(runAttempts.attemptNo));
  const stages = await db
    .select()
    .from(stageExecutions)
    .where(eq(stageExecutions.runId, runId))
    .orderBy(asc(stageExecutions.startedAt));
  const arts = await db
    .select({
      id: artifacts.id,
      path: artifacts.path,
      kind: artifacts.kind,
      sha256: artifacts.sha256,
      bytes: artifacts.bytes,
      contentType: artifacts.contentType,
      stage: artifacts.stage,
      version: artifacts.version,
      synthetic: artifacts.synthetic,
      createdAt: artifacts.createdAt,
      attemptId: artifacts.attemptId,
    })
    .from(artifacts)
    .where(eq(artifacts.runId, runId))
    .orderBy(asc(artifacts.path), desc(artifacts.version));
  const [lease] = await db.select().from(jobLeases).where(eq(jobLeases.runId, runId)).limit(1);
  const gates = await db
    .select()
    .from(qualityGateResults)
    .where(eq(qualityGateResults.runId, runId))
    .orderBy(desc(qualityGateResults.createdAt));
  const latestEval = gates[0]?.evaluationId;
  return {
    run,
    attempts,
    stages,
    artifacts: arts,
    lease: lease ?? null,
    gates: gates.filter((g) => g.evaluationId === latestEval),
  };
}

export async function listRuns(
  db: Db,
  opts: { status?: RunStatus[]; limit?: number; projectId?: string } = {},
) {
  const conds = [];
  if (opts.status?.length) conds.push(inArray(researchRuns.status, opts.status));
  if (opts.projectId) conds.push(eq(researchRuns.projectId, opts.projectId));
  return db
    .select()
    .from(researchRuns)
    .where(conds.length ? and(...conds) : undefined)
    .orderBy(desc(researchRuns.createdAt))
    .limit(opts.limit ?? 100);
}

export async function listEvents(db: Db, runId: string, afterSeq = 0, limit = 500) {
  return db
    .select()
    .from(stageEvents)
    .where(and(eq(stageEvents.runId, runId), gt(stageEvents.seq, afterSeq)))
    .orderBy(asc(stageEvents.seq))
    .limit(limit);
}

export async function requestCancel(db: Db, runId: string, actor: { id: string; ip?: string | null }) {
  return db.transaction(async (tx) => {
    const [run] = await tx.select().from(researchRuns).where(eq(researchRuns.id, runId)).for("update");
    if (!run) throw notFound("run");
    const st = run.status as RunStatus;
    if (["published", "cancelled", "failed"].includes(st)) throw conflict(`run already ${st}`);
    if (
      [
        "queued",
        "validated",
        "draft",
        "quality_failed",
        "ready_to_publish",
        "awaiting_approval",
        "retrying",
        "revising",
      ].includes(st)
    ) {
      const updated = await transition(tx, runId, "cancelled", {
        actor: `admin:${actor.id}`,
        reason: "cancelled by admin",
      });
      await tx.delete(jobLeases).where(eq(jobLeases.runId, runId));
      await audit(tx, { type: "admin", id: actor.id, ip: actor.ip }, "run.cancel", {
        type: "run",
        id: runId,
      });
      return updated;
    }
    // running on a worker: cooperative cancel, worker acknowledges via complete(cancelled); lease expiry is the fallback
    await tx
      .update(researchRuns)
      .set({ cancelRequested: true, updatedAt: new Date() })
      .where(eq(researchRuns.id, runId));
    const updated = await transition(tx, runId, "cancelling", {
      actor: `admin:${actor.id}`,
      reason: "cancellation requested",
    });
    await audit(tx, { type: "admin", id: actor.id, ip: actor.ip }, "run.cancel", { type: "run", id: runId });
    return updated;
  });
}

/** Retry from the last checkpoint with a new attempt; preserves frozen input and previous artifacts. */
export async function retryRun(db: Db, runId: string, actor: { id: string; ip?: string | null }) {
  return db.transaction(async (tx) => {
    const [run] = await tx.select().from(researchRuns).where(eq(researchRuns.id, runId)).for("update");
    if (!run) throw notFound("run");
    const st = run.status as RunStatus;
    if (!["failed", "quality_failed", "cancelled"].includes(st))
      throw conflict(`cannot retry a run in status ${st}`);
    if (st === "cancelled")
      throw conflict("cancelled runs cannot be retried; duplicate the project run instead");
    if (run.currentAttemptNo >= run.maxAttempts) {
      await tx
        .update(researchRuns)
        .set({ maxAttempts: run.currentAttemptNo + 1 })
        .where(eq(researchRuns.id, runId));
    }
    await tx.delete(jobLeases).where(eq(jobLeases.runId, runId));
    await tx
      .update(researchRuns)
      .set({ cancelRequested: false, failureReason: null, finishedAt: null })
      .where(eq(researchRuns.id, runId));
    if (st === "quality_failed") await transition(tx, runId, "revising", { actor: `admin:${actor.id}` });
    const updated = await transition(tx, runId, "queued", {
      actor: `admin:${actor.id}`,
      reason: "retry requested",
    });
    await audit(
      tx,
      { type: "admin", id: actor.id, ip: actor.ip },
      "run.retry",
      { type: "run", id: runId },
      { fromStatus: st },
    );
    return updated;
  });
}

export async function latestCheckpoint(db: Db | Tx, runId: string) {
  const [c] = await db
    .select()
    .from(checkpoints)
    .where(eq(checkpoints.runId, runId))
    .orderBy(desc(checkpoints.createdAt))
    .limit(1);
  return c ?? null;
}

export async function overviewCounts(db: Db) {
  const rows = await db
    .select({ status: researchRuns.status, n: sql<number>`count(*)::int` })
    .from(researchRuns)
    .groupBy(researchRuns.status);
  const byStatus: Record<string, number> = {};
  for (const r of rows) byStatus[r.status] = r.n;
  const [cost] = await db
    .select({ total: sql<string>`coalesce(sum(${researchRuns.totalCostUsd}), 0)::text` })
    .from(researchRuns);
  return {
    byStatus,
    active: ACTIVE_RUN_STATUSES.reduce((n, s) => n + (byStatus[s] ?? 0), 0),
    totalCostUsd: Number(cost?.total ?? 0),
  };
}

export { AppError };