import { FakeHarnessLive } from "@code/primitives/harness-runtime";
import type { FakeScenario } from "@code/primitives/harness-runtime";
import { defaultCodingKitV0, resolveOutcome } from "@code/primitives/resolver";
import type { WorkEventKind } from "@code/primitives/work";
import { makeFunctionReference } from "convex/server";
import { ConvexError, v } from "convex/values";
import { Effect } from "effect";

import type { Doc, Id } from "./_generated/dataModel";
import {
  internalAction,
  internalMutation,
  mutation,
  query,
} from "./_generated/server";
import type { MutationCtx } from "./_generated/server";
import { requireProjectMember } from "./authz";

const executeRef = makeFunctionReference<
  "action",
  { attemptId: Id<"workAttempts">; scenario: FakeScenario }
>("workExecution:executeFakeAttempt");
const claimRef = makeFunctionReference<
  "mutation",
  { attemptId: Id<"workAttempts">; leaseMs: number; owner: string },
  { number: number } | null
>("workExecution:claimAttempt");
const checkpointRef = makeFunctionReference<
  "mutation",
  {
    attemptId: Id<"workAttempts">;
    kind: string;
    message: string;
    metadataJson: string;
    owner: string;
    sequence: number;
  },
  unknown
>("workExecution:checkpointAttempt");
const finishRef = makeFunctionReference<
  "mutation",
  {
    attemptId: Id<"workAttempts">;
    classification: string;
    owner: string;
    retryable: boolean;
    summary: string;
  },
  unknown
>("workExecution:finishAttempt");
const now = () => Date.now();
const attemptClassification = v.union(
  v.literal("Succeeded"),
  v.literal("RetryableFailure"),
  v.literal("NeedsInput"),
  v.literal("Blocked"),
  v.literal("VerificationFailed"),
  v.literal("BudgetExhausted"),
  v.literal("Cancelled"),
  v.literal("PermanentFailure")
);

const requireWorkForMember = async (
  ctx: MutationCtx,
  workId: Id<"works">
): Promise<Doc<"works">> => {
  const work = await ctx.db.get(workId);
  if (!work) {
    throw new ConvexError("Work not found");
  }
  await requireProjectMember(ctx, work.projectId);
  return work;
};

const appendWorkEvent = async (
  ctx: MutationCtx,
  workId: Id<"works">,
  kind: WorkEventKind,
  idempotencyKey: string,
  referenceId?: string,
  payloadJson?: string
) => {
  const existing = await ctx.db
    .query("workEvents")
    .withIndex("by_work_and_idempotencyKey", (q) =>
      q.eq("workId", workId).eq("idempotencyKey", idempotencyKey)
    )
    .unique();
  if (existing) {
    return;
  }
  await ctx.db.insert("workEvents", {
    createdAt: now(),
    idempotencyKey,
    kind,
    workId,
    ...(referenceId ? { referenceId } : {}),
    ...(payloadJson ? { payloadJson } : {}),
  });
};

// Slice the Run targets for the currently approved Design. `sliceId` is
// optional only to let the resolver default to the first ready slice.
const resolveRunSlice = async (
  ctx: MutationCtx,
  work: Doc<"works">,
  sliceId?: string
): Promise<Doc<"workSlices">> => {
  const currentDesignVersion = work.designVersion;
  if (currentDesignVersion === undefined) {
    throw new ConvexError("Work has no approved Design to execute");
  }
  const slices = await ctx.db
    .query("workSlices")
    .withIndex("by_workId_and_designVersion", (q) =>
      q.eq("workId", work._id).eq("designVersion", currentDesignVersion)
    )
    .collect();
  if (slices.length === 0) {
    throw new ConvexError("No slices found for the current approved Design");
  }
  const slice = sliceId
    ? slices.find((row) => row.sliceId === sliceId)
    : slices
        .sort((a, b) => a.ordinal - b.ordinal)
        .find((row) => row.status === "ready");
  if (!slice) {
    throw new ConvexError(
      "Requested slice does not belong to the current approved Design"
    );
  }
  if (slice.status !== "ready") {
    throw new ConvexError("Only the next ready slice can be executed");
  }
  return slice;
};

const getRunSlice = async (
  ctx: MutationCtx,
  run: Doc<"workRuns">
): Promise<Doc<"workSlices"> | null> => {
  if (run.sliceRowId) {
    return await ctx.db.get(run.sliceRowId);
  }
  if (run.designVersion === undefined || run.sliceId === undefined) {
    return null;
  }
  return await ctx.db
    .query("workSlices")
    .withIndex("by_workId_and_designVersion_and_sliceId", (q) =>
      q
        .eq("workId", run.workId)
        .eq("designVersion", run.designVersion!)
        .eq("sliceId", run.sliceId!)
    )
    .unique();
};

const RETRYABLE_WORK_STATUSES = new Set(["failed", "needs-input", "blocked"]);

export const startSimulatedExecution = mutation({
  args: {
    scenario: v.union(
      v.literal("success"),
      v.literal("transient-failure-then-success"),
      v.literal("needs-input"),
      v.literal("permanent-failure"),
      v.literal("cancelled")
    ),
    sliceId: v.optional(v.string()),
    workId: v.id("works"),
  },
  handler: async (ctx, args) => {
    const work = await requireWorkForMember(ctx, args.workId);
    if (work.status !== "ready") {
      throw new ConvexError("Work must be Ready before simulated execution");
    }
    if (
      work.definitionApprovalVersion !== work.definitionVersion ||
      work.designApprovalVersion !== work.designVersion
    ) {
      throw new ConvexError(
        "Execution requires exact approved Definition and Design versions"
      );
    }
    const slice = await resolveRunSlice(ctx, work, args.sliceId);
    const createdAt = now();
    const runId = await ctx.db.insert("workRuns", {
      createdAt,
      designVersion: slice.designVersion,
      kitId: defaultCodingKitV0.id,
      kitVersion: defaultCodingKitV0.version,
      scenario: args.scenario,
      sliceId: slice.sliceId,
      sliceRowId: slice._id,
      status: "ready",
      workId: work._id,
    });
    const attemptId = await ctx.db.insert("workAttempts", {
      number: 1,
      runId,
      status: "queued",
      workId: work._id,
    });
    await ctx.db.patch(work._id, {
      status: "executing",
      updatedAt: createdAt,
    });
    await ctx.db.patch(slice._id, { status: "running" });
    await ctx.db.patch(runId, { startedAt: createdAt, status: "running" });
    await appendWorkEvent(
      ctx,
      work._id,
      "run.started",
      `run-started:${runId}`,
      String(runId)
    );
    await appendWorkEvent(
      ctx,
      work._id,
      "slice.started",
      `slice-started:${runId}:${slice.sliceId}`,
      String(slice._id)
    );
    await ctx.scheduler.runAfter(0, executeRef, {
      attemptId,
      scenario: args.scenario,
    });
    return { attemptId, runId };
  },
});

export const cancelSimulatedExecution = mutation({
  args: { runId: v.id("workRuns") },
  handler: async (ctx, args) => {
    const run = await ctx.db.get(args.runId);
    if (!run) {
      throw new ConvexError("Run not found");
    }
    const work = await requireWorkForMember(ctx, run.workId);
    if (run.status === "terminal" || run.status === "cancelled") {
      return { cancelled: false };
    }
    await ctx.db.patch(run._id, {
      endedAt: now(),
      status: "cancelled",
      terminalClassification: "Cancelled",
      terminalSummary: "Simulation cancelled",
    });
    const attempts = await ctx.db
      .query("workAttempts")
      .withIndex("by_runId_and_number", (q) => q.eq("runId", run._id))
      .collect();
    for (const attempt of attempts) {
      if (attempt.status !== "terminal") {
        await ctx.db.patch(attempt._id, {
          classification: "Cancelled",
          endedAt: now(),
          status: "terminal",
          summary: "Simulation cancelled",
        });
        await ctx.db.insert("resolverDecisions", {
          attemptId: attempt._id,
          attemptNumber: attempt.number,
          classification: "Cancelled",
          createdAt: now(),
          decision: "terminal",
          resultingWorkStatus: "ready",
          runId: run._id,
          summary: "Simulation cancelled",
          workId: work._id,
        });
        await appendWorkEvent(
          ctx,
          work._id,
          "resolver.decided",
          `resolver-decided:${attempt._id}`,
          String(attempt._id),
          JSON.stringify({ classification: "Cancelled", decision: "terminal" })
        );
      }
    }
    const slice = await getRunSlice(ctx, run);
    if (slice) {
      await ctx.db.patch(slice._id, { status: "ready" });
    }
    await ctx.db.patch(work._id, { status: "ready", updatedAt: now() });
    await appendWorkEvent(
      ctx,
      work._id,
      "run.cancelled",
      `run-cancelled:${run._id}`,
      String(run._id)
    );
    return { cancelled: true };
  },
});

// User-initiated restart of a terminal Run. Unlike the automatic resolver
// retry, this is an explicit decision and may exceed the kit budget.
export const retrySimulatedExecution = mutation({
  args: { runId: v.id("workRuns") },
  handler: async (ctx, args) => {
    const run = await ctx.db.get(args.runId);
    if (!run) {
      throw new ConvexError("Run not found");
    }
    const work = await requireWorkForMember(ctx, run.workId);
    if (run.status !== "terminal") {
      throw new ConvexError("Only terminal Runs can be retried");
    }
    if (!RETRYABLE_WORK_STATUSES.has(work.status)) {
      throw new ConvexError("Work is not in a retryable state");
    }
    const attempts = await ctx.db
      .query("workAttempts")
      .withIndex("by_runId_and_number", (q) => q.eq("runId", run._id))
      .collect();
    const number = attempts.length + 1;
    const attemptId = await ctx.db.insert("workAttempts", {
      number,
      runId: run._id,
      status: "queued",
      workId: work._id,
    });
    await ctx.db.patch(run._id, {
      endedAt: undefined,
      startedAt: now(),
      status: "running",
      terminalClassification: undefined,
      terminalSummary: undefined,
    });
    await ctx.db.patch(work._id, { status: "executing", updatedAt: now() });
    const slice = await getRunSlice(ctx, run);
    if (!slice) {
      throw new ConvexError("Run slice not found");
    }
    await ctx.db.patch(slice._id, { status: "running" });
    await ctx.scheduler.runAfter(0, executeRef, {
      attemptId,
      scenario: run.scenario,
    });
    return { attemptId, number };
  },
});

export const claimAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    leaseMs: v.number(),
    owner: v.string(),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt || attempt.status !== "queued") {
      return null;
    }
    const run = await ctx.db.get(attempt.runId);
    if (!run || run.status !== "running") {
      return null;
    }
    const claimedAt = now();
    await ctx.db.patch(attempt._id, {
      leaseExpiresAt: claimedAt + args.leaseMs,
      leaseOwner: args.owner,
      startedAt: attempt.startedAt ?? claimedAt,
      status: "claimed",
    });
    await ctx.db.patch(attempt.runId, {
      startedAt: claimedAt,
      status: "running",
    });
    return { ...attempt, leaseOwner: args.owner, status: "claimed" as const };
  },
});

export const checkpointAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    kind: v.string(),
    message: v.string(),
    metadataJson: v.string(),
    owner: v.string(),
    sequence: v.number(),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (
      !attempt ||
      attempt.leaseOwner !== args.owner ||
      attempt.status === "terminal" ||
      (attempt.leaseExpiresAt !== undefined && attempt.leaseExpiresAt < now())
    ) {
      return false;
    }
    const existing = await ctx.db
      .query("workAttemptEvents")
      .withIndex("by_attempt_and_sequence", (q: any) =>
        q.eq("attemptId", attempt._id).eq("sequence", args.sequence)
      )
      .unique();
    if (!existing) {
      await ctx.db.insert("workAttemptEvents", {
        attemptId: attempt._id,
        kind: args.kind,
        message: args.message,
        metadataJson: args.metadataJson,
        occurredAt: now(),
        sequence: args.sequence,
      });
    }
    await ctx.db.patch(attempt._id, {
      leaseExpiresAt: now() + 60_000,
      status: "running",
    });
    return true;
  },
});

// Single resolver mutation: record the attempt outcome, then either schedule
// the next attempt within the kit retry policy or settle the Run/Work at the
// terminal classification. Replaces the old completeAttempt that always reset
// Work to "ready" and ignored the retry policy.
export const finishAttempt = internalMutation({
  args: {
    attemptId: v.id("workAttempts"),
    classification: attemptClassification,
    owner: v.string(),
    retryable: v.boolean(),
    summary: v.string(),
  },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (
      !attempt ||
      attempt.leaseOwner !== args.owner ||
      attempt.status === "terminal" ||
      (attempt.leaseExpiresAt !== undefined && attempt.leaseExpiresAt < now())
    ) {
      return null;
    }
    const endedAt = now();
    await ctx.db.patch(attempt._id, {
      classification: args.classification,
      endedAt,
      leaseExpiresAt: undefined,
      status: "terminal",
      summary: args.summary,
    });
    const outcome = {
      classification: args.classification,
      retryable: args.retryable,
      summary: args.summary,
    };
    const resolution = resolveOutcome(
      outcome,
      attempt.number,
      defaultCodingKitV0.retryPolicy
    );
    const run = await ctx.db.get(attempt.runId);
    if (!run) {
      return null;
    }
    await appendWorkEvent(
      ctx,
      attempt.workId,
      "attempt.completed",
      `attempt-completed:${attempt._id}`,
      String(attempt._id),
      JSON.stringify({
        classification: args.classification,
        retried: resolution.kind === "retry",
        summary: args.summary,
      })
    );
    await appendWorkEvent(
      ctx,
      attempt.workId,
      "resolver.decided",
      `resolver-decided:${attempt._id}`,
      String(attempt._id),
      JSON.stringify({
        classification: args.classification,
        decision: resolution.kind,
      })
    );
    if (resolution.kind === "retry") {
      await ctx.db.insert("resolverDecisions", {
        attemptId: attempt._id,
        attemptNumber: attempt.number,
        classification: args.classification,
        createdAt: endedAt,
        decision: "retry",
        runId: run._id,
        summary: args.summary,
        workId: attempt.workId,
      });
      // Keep Run running and Work executing; spin the next attempt.
      const nextAttemptId = await ctx.db.insert("workAttempts", {
        number: attempt.number + 1,
        runId: run._id,
        status: "queued",
        workId: attempt.workId,
      });
      await ctx.scheduler.runAfter(0, executeRef, {
        attemptId: nextAttemptId,
        scenario: run.scenario,
      });
      return { nextAttemptId, retried: true };
    }
    await ctx.db.patch(run._id, {
      endedAt,
      status: "terminal",
      terminalClassification: args.classification,
      terminalSummary: args.summary,
    });
    const work = await ctx.db.get(run.workId);
    const slice = await getRunSlice(ctx, run);
    let resultingWorkStatus = resolution.workStatus;
    if (args.classification === "Succeeded" && slice) {
      await ctx.db.patch(slice._id, { status: "completed" });
      await appendWorkEvent(
        ctx,
        attempt.workId,
        "slice.completed",
        `slice-completed:${run._id}:${slice.sliceId}`,
        String(slice._id)
      );
      const slices = await ctx.db
        .query("workSlices")
        .withIndex("by_workId_and_designVersion", (q) =>
          q
            .eq("workId", attempt.workId)
            .eq("designVersion", slice.designVersion)
        )
        .collect();
      const nextSlice = slices
        .sort((left, right) => left.ordinal - right.ordinal)
        .find((candidate) => candidate.status === "planned");
      if (nextSlice) {
        await ctx.db.patch(nextSlice._id, { status: "ready" });
        resultingWorkStatus = "ready";
        await appendWorkEvent(
          ctx,
          attempt.workId,
          "slice.ready",
          `slice-ready:${run._id}:${nextSlice.sliceId}`,
          String(nextSlice._id)
        );
      }
    } else if (slice) {
      const blocked = ["NeedsInput", "Blocked", "VerificationFailed"].includes(
        args.classification
      );
      await ctx.db.patch(slice._id, { status: blocked ? "blocked" : "ready" });
    }
    if (work) {
      await ctx.db.patch(work._id, {
        status: resultingWorkStatus,
        updatedAt: endedAt,
      });
    }
    await ctx.db.insert("resolverDecisions", {
      attemptId: attempt._id,
      attemptNumber: attempt.number,
      classification: args.classification,
      createdAt: endedAt,
      decision: "terminal",
      resultingWorkStatus,
      runId: run._id,
      summary: args.summary,
      workId: attempt.workId,
    });
    await appendWorkEvent(
      ctx,
      attempt.workId,
      "run.completed",
      `run-completed:${run._id}`,
      String(run._id),
      JSON.stringify({
        classification: args.classification,
        workStatus: resultingWorkStatus,
      })
    );
    return { retried: false };
  },
});

export const executeFakeAttempt = internalAction({
  args: {
    attemptId: v.id("workAttempts"),
    scenario: v.union(
      v.literal("success"),
      v.literal("transient-failure-then-success"),
      v.literal("needs-input"),
      v.literal("permanent-failure"),
      v.literal("cancelled")
    ),
  },
  handler: async (ctx, args) => {
    const owner = `fake-worker:${args.attemptId}`;
    const claimed = await ctx.runMutation(claimRef, {
      attemptId: args.attemptId,
      leaseMs: 60_000,
      owner,
    });
    if (!claimed) {
      return null;
    }
    const [events, outcome] = await Effect.runPromise(
      FakeHarnessLive().run({
        attemptNumber: claimed.number,
        scenario: args.scenario,
      })
    );
    for (const item of events) {
      await ctx.runMutation(checkpointRef, {
        attemptId: args.attemptId,
        kind: item.kind,
        message: item.message,
        metadataJson: JSON.stringify(item.metadata),
        owner,
        sequence: item.sequence,
      });
    }
    await ctx.runMutation(finishRef, {
      attemptId: args.attemptId,
      classification: outcome.classification,
      owner,
      retryable: outcome.retryable,
      summary: outcome.summary,
    });
    return null;
  },
});

// Bounded recovery: a claimed/running attempt whose lease expired means the
// worker died mid-flight. Terminate that attempt, then either resume within
// the kit budget or settle the Run/Work so nothing stays "running" forever.
export const reconcileExpiredAttempts = internalMutation({
  args: {},
  handler: async (ctx) => {
    const timestamp = now();
    const [claimed, running] = await Promise.all([
      ctx.db
        .query("workAttempts")
        .withIndex("by_status_and_leaseExpiresAt", (q) =>
          q.eq("status", "claimed").lt("leaseExpiresAt", timestamp)
        )
        .collect(),
      ctx.db
        .query("workAttempts")
        .withIndex("by_status_and_leaseExpiresAt", (q) =>
          q.eq("status", "running").lt("leaseExpiresAt", timestamp)
        )
        .collect(),
    ]);
    const active = [...claimed, ...running];
    let reconciled = 0;
    for (const attempt of active) {
      if (
        (attempt.status === "claimed" || attempt.status === "running") &&
        attempt.leaseExpiresAt !== undefined &&
        attempt.leaseExpiresAt < now()
      ) {
        await ctx.db.patch(attempt._id, {
          classification: "Blocked",
          endedAt: now(),
          leaseExpiresAt: undefined,
          leaseOwner: undefined,
          status: "terminal",
          summary: "Attempt lease expired and was reconciled",
        });
        const run = await ctx.db.get(attempt.runId);
        await appendWorkEvent(
          ctx,
          attempt.workId,
          "attempt.reconciled",
          `attempt-reconciled:${attempt._id}`,
          String(attempt._id),
          JSON.stringify({ attemptNumber: attempt.number })
        );
        reconciled += 1;
        if (!run) {
          continue;
        }
        const attempts = await ctx.db
          .query("workAttempts")
          .withIndex("by_runId_and_number", (q) => q.eq("runId", run._id))
          .collect();
        const withinBudget =
          attempts.length < defaultCodingKitV0.retryPolicy.maxAttempts;
        const retry = run.status === "running" && withinBudget;
        await ctx.db.insert("resolverDecisions", {
          attemptId: attempt._id,
          attemptNumber: attempt.number,
          classification: "Blocked",
          createdAt: now(),
          decision: retry ? "retry" : "terminal",
          resultingWorkStatus: retry ? undefined : "blocked",
          runId: run._id,
          summary: "Attempt lease expired and was reconciled",
          workId: attempt.workId,
        });
        await appendWorkEvent(
          ctx,
          attempt.workId,
          "resolver.decided",
          `resolver-decided:${attempt._id}`,
          String(attempt._id),
          JSON.stringify({
            classification: "Blocked",
            decision: retry ? "retry" : "terminal",
          })
        );
        if (retry) {
          const nextAttemptId = await ctx.db.insert("workAttempts", {
            number: attempts.length + 1,
            runId: run._id,
            status: "queued",
            workId: attempt.workId,
          });
          await ctx.scheduler.runAfter(0, executeRef, {
            attemptId: nextAttemptId,
            scenario: run.scenario,
          });
        } else {
          await ctx.db.patch(run._id, {
            endedAt: now(),
            status: "terminal",
            terminalClassification: "Blocked",
            terminalSummary: "Run reconciled after expired lease",
          });
          const work = await ctx.db.get(run.workId);
          if (work) {
            await ctx.db.patch(work._id, {
              status: "blocked",
              updatedAt: now(),
            });
          }
          const slice = await getRunSlice(ctx, run);
          if (slice) {
            await ctx.db.patch(slice._id, { status: "blocked" });
          }
        }
      }
    }
    return { reconciled };
  },
});

export const listRunEvents = query({
  args: { attemptId: v.id("workAttempts") },
  handler: async (ctx, args) => {
    const attempt = await ctx.db.get(args.attemptId);
    if (!attempt) {
      return [];
    }
    const work = await ctx.db.get(attempt.workId);
    if (!work) {
      return [];
    }
    await requireProjectMember(ctx, work.projectId);
    return await ctx.db
      .query("workAttemptEvents")
      .withIndex("by_attempt_and_sequence", (q) =>
        q.eq("attemptId", args.attemptId)
      )
      .order("asc")
      .collect();
  },
});