import { z } from "zod";
import { CapabilityService } from "./capability";
import { buildTransitionEvent } from "./events";
import { ClusterScheduler } from "./scheduler";
import type { HostState } from "./status";
import {
  AuditWriteError,
  StaleRevisionError,
  type ControllerStore,
  type TransitionIntent,
  type TransitionVerb,
} from "./store";

const ADMISSION_RECONCILE_TIMEOUT_MS = 8_000;

function unhealthyChecks(record: {
  healthStorage: boolean;
  healthRunner: boolean;
  healthOffload: boolean;
}): string[] {
  const failing: string[] = [];
  if (!record.healthStorage) failing.push("storage");
  if (!record.healthRunner) failing.push("runner");
  if (!record.healthOffload) failing.push("offload");
  return failing;
}

export const TRANSITION_VERBS = [
  "box-drain",
  "box-restore",
  "host-quarantine",
  "host-unquarantine",
  "admission-reconcile",
  "job-retry",
  "ci-reconcile",
  "recall-spill",
] as const;

const VerbSchema = z.enum(TRANSITION_VERBS);

const BoxHostArgsSchema = z.object({
  host: z.string().min(1),
});

const QuarantineArgsSchema = z.object({
  host: z.string().min(1),
  command: z.string().min(1),
});

const AdmissionReconcileArgsSchema = z.object({
  reason: z.string().optional(),
}).strict();

const JobRetryArgsSchema = z.object({
  jobId: z.string().min(1),
});

const OptionalHostArgsSchema = z.object({
  host: z.string().min(1).optional(),
});

export type TransitionSuccess = {
  status: 200;
  body: { revision: number; result: unknown };
};

export type TransitionFailure = {
  status: 409 | 422 | 403;
  body: Record<string, unknown>;
};

export type TransitionResponse = TransitionSuccess | TransitionFailure;

type TransitionRequest = {
  expectedRevision: unknown;
  idempotencyKey: string;
  args?: Record<string, unknown>;
};

type CommitParams = {
  verb: TransitionVerb;
  idempotencyKey: string;
  args: Record<string, unknown>;
  expectedRevision: number;
  journalId: number;
  skipJournal?: boolean;
  skipIdempotency?: boolean;
  skipRevisionCheck?: boolean;
};

type TransitionGuard =
  | { ok: true; result: unknown; apply: () => void; acceptedRevision?: number }
  | { ok: false; detail: string };

export class TransitionEngine {
  constructor(
    private readonly store: ControllerStore,
    private readonly now: () => number = () => Date.now(),
    private readonly capability: CapabilityService = new CapabilityService(store),
    private readonly scheduler: ClusterScheduler = new ClusterScheduler(store, { now }),
  ) {}

  async resumePending(): Promise<TransitionResponse[]> {
    const responses: TransitionResponse[] = [];
    for (const pending of this.store.listPendingJournals()) {
      if (pending.verb === "delivery-feature-reconcile") continue;
      const response = await this.executeCommitted({
        verb: pending.verb,
        idempotencyKey: pending.idempotencyKey,
        args: pending.args,
        expectedRevision: pending.expectedRevision,
        journalId: pending.id,
        skipJournal: true,
        skipIdempotency: true,
        skipRevisionCheck: true,
      });
      responses.push(response);
    }
    return responses;
  }

  handle(verbInput: "admission-reconcile", body: TransitionRequest): Promise<TransitionResponse>;
  handle(
    verbInput: Exclude<TransitionVerb, "admission-reconcile">,
    body: TransitionRequest,
  ): TransitionResponse;
  handle(
    verbInput: string,
    body: TransitionRequest,
  ): TransitionResponse | Promise<TransitionResponse>;
  handle(
    verbInput: string,
    body: TransitionRequest,
  ): TransitionResponse | Promise<TransitionResponse> {
    const verbResult = VerbSchema.safeParse(verbInput);
    if (!verbResult.success) {
      return {
        status: 422,
        body: { error: "invalid-args", detail: `unknown verb: ${verbInput}` },
      };
    }
    const verb = verbResult.data;

    const expectedRevision = coerceExpectedRevision(body.expectedRevision);
    if (expectedRevision === null) {
      return {
        status: 422,
        body: { error: "invalid-args", detail: "expectedRevision must be a number" },
      };
    }

    if (!body.idempotencyKey || typeof body.idempotencyKey !== "string") {
      return {
        status: 422,
        body: { error: "invalid-args", detail: "idempotencyKey is required" },
      };
    }

    const args = body.args ?? {};
    const argsValidation = validateVerbArgs(verb, args);
    if (!argsValidation.ok) {
      return {
        status: 422,
        body: { error: "invalid-args", detail: argsValidation.detail },
      };
    }

    const replay = this.store.getIdempotencyResult(body.idempotencyKey);
    if (replay) {
      return {
        status: 200,
        body: { revision: replay.revision, result: replay.result },
      };
    }

    let journalId: number;
    try {
      journalId = this.store.journalIntent({
        verb,
        idempotencyKey: body.idempotencyKey,
        args,
        expectedRevision,
      });
    } catch (err) {
      if (err instanceof AuditWriteError) {
        return { status: 403, body: { error: "audit-write-failed" } };
      }
      throw err;
    }

    const currentRevision = this.store.getRevision();
    if (expectedRevision !== currentRevision) {
      this.store.cancelJournal(journalId);
      return {
        status: 409,
        body: { error: "stale-revision", currentRevision },
      };
    }

    return this.executeCommitted({
      verb,
      idempotencyKey: body.idempotencyKey,
      args,
      expectedRevision,
      journalId,
    });
  }

  private executeCommitted(params: CommitParams): TransitionResponse | Promise<TransitionResponse> {
    if (!params.skipRevisionCheck) {
      const currentRevision = this.store.getRevision();
      if (params.expectedRevision !== currentRevision) {
        this.store.cancelJournal(params.journalId);
        return {
          status: 409,
          body: { error: "stale-revision", currentRevision },
        };
      }
    }

    if (params.verb === "admission-reconcile") {
      return this.executeAdmissionCommitted(params);
    }

    const guard = this.guardExecution(params.verb, params.args);
    return this.commitGuard(params, guard);
  }

  private async executeAdmissionCommitted(params: CommitParams): Promise<TransitionResponse> {
    const guard = await this.admissionReconcile(params.args);
    return this.commitGuard(params, guard);
  }

  private commitGuard(params: CommitParams, guard: TransitionGuard): TransitionResponse {
    if (!guard.ok) {
      this.store.cancelJournal(params.journalId);
      return {
        status: 422,
        body: { error: "invalid-args", detail: guard.detail },
      };
    }

    const result = guard.result;
    const event = buildTransitionEvent(
      this.store,
      params.verb,
      params.args,
      result,
      this.now,
    );
    try {
      const revision = this.store.commitTransition({
        journalId: params.journalId,
        idempotencyKey: params.idempotencyKey,
        result,
        apply: guard.apply,
        event,
        acceptedRevision: guard.acceptedRevision,
      });
      return { status: 200, body: { revision, result } };
    } catch (error) {
      if (error instanceof StaleRevisionError) {
        this.store.cancelJournal(params.journalId);
        return {
          status: 409,
          body: { error: "stale-revision", currentRevision: error.currentRevision },
        };
      }
      throw error;
    }
  }

  private guardExecution(
    verb: TransitionVerb,
    args: Record<string, unknown>,
  ): TransitionGuard {
    switch (verb) {
      case "box-drain":
        return this.boxDrain(args);
      case "box-restore":
        return this.boxRestore(args);
      case "host-quarantine":
        return this.hostQuarantine(args);
      case "host-unquarantine":
        return this.hostUnquarantine(args);
      case "admission-reconcile":
        return { ok: false, detail: "admission-reconcile requires async execution" };
      case "job-retry":
        return this.jobRetry(args);
      case "ci-reconcile":
        return this.ciReconcile(args);
      case "recall-spill":
        return this.recallSpill(args);
      default:
        return { ok: false, detail: `unsupported verb: ${verb}` };
    }
  }

  private boxDrain(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = BoxHostArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const { host } = parsed.data;
    const record = this.store.getHost(host);
    if (!record) return { ok: false, detail: `unknown host: ${host}` };
    if (record.state !== "available") {
      return { ok: false, detail: `host ${host} is ${record.state}, expected available` };
    }
    return {
      ok: true,
      result: { host, from: record.state, to: "draining" as HostState },
      apply: () => {
        this.store.setHostState(host, "draining");
        this.syncClusterObserved();
      },
    };
  }

  private boxRestore(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = BoxHostArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const { host } = parsed.data;
    const record = this.store.getHost(host);
    if (!record) return { ok: false, detail: `unknown host: ${host}` };

    const restorable: HostState[] = ["draining", "maintenance", "degraded"];
    if (!restorable.includes(record.state) && record.state !== "restoring") {
      return {
        ok: false,
        detail: `host ${host} is ${record.state}, cannot restore`,
      };
    }

    if (record.state === "restoring") {
      if (!this.store.isHostHealthGreen(host)) {
        return {
          ok: false,
          detail: `host ${host} restoring but health checks are not green`,
        };
      }
      return {
        ok: true,
        result: { host, from: "restoring", to: "available" as HostState },
        apply: () => {
          const current = this.store.getHost(host);
          if (!current) throw new Error(`unknown host: ${host}`);
          this.store.upsertHost({
            ...current,
            state: "available",
            enrolling: false,
          });
          this.syncClusterObserved();
        },
      };
    }

    return {
      ok: true,
      result: { host, from: record.state, to: "restoring" as HostState },
      apply: () => {
        this.store.setHostState(host, "restoring");
        if (this.store.isHostHealthGreen(host)) {
          const current = this.store.getHost(host);
          if (!current) throw new Error(`unknown host: ${host}`);
          this.store.upsertHost({
            ...current,
            state: "available",
            enrolling: false,
          });
        }
        this.syncClusterObserved();
      },
    };
  }

  private hostQuarantine(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = QuarantineArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const { host, command } = parsed.data;
    if (!this.store.getHost(host)) return { ok: false, detail: `unknown host: ${host}` };
    return {
      ok: true,
      result: { host, command, circuitOpen: true },
      apply: () => {
        this.capability.quarantine(host, command);
      },
    };
  }

  private hostUnquarantine(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = QuarantineArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const { host, command } = parsed.data;
    if (!this.store.getHost(host)) return { ok: false, detail: `unknown host: ${host}` };
    return {
      ok: true,
      result: { host, command, circuitOpen: false, probe: "half-open" },
      apply: () => {
        this.capability.halfOpen(host, command);
      },
    };
  }

  private async admissionReconcile(
    args: Record<string, unknown>,
  ): Promise<
    | { ok: true; result: unknown; apply: () => void; acceptedRevision: number }
    | { ok: false; detail: string }
  > {
    const parsed = AdmissionReconcileArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }

    const acceptedRevision = this.store.getRevision();
    const advanced: Array<{ host: string; from: HostState; to: HostState }> = [];
    for (const host of this.store.listHosts()) {
      if (host.state === "draining" && this.store.isHostIdle(host.hostname)) {
        advanced.push({
          host: host.hostname,
          from: "draining",
          to: "maintenance",
        });
      }
    }

    const enrollingHosts = this.store
      .listHosts()
      .filter((host) => host.enrolling && host.state === "maintenance");

    let timeout: ReturnType<typeof setTimeout> | undefined;
    const probeResults = await Promise.race([
      Promise.all(
        enrollingHosts.map(async (host) => ({
          host: host.hostname,
          result: await this.capability.probeAdmission(host.hostname),
        })),
      ),
      new Promise<null>((resolve) => {
        timeout = setTimeout(() => resolve(null), ADMISSION_RECONCILE_TIMEOUT_MS);
      }),
    ]).finally(() => {
      if (timeout) clearTimeout(timeout);
    });

    const probes = probeResults === null
      ? enrollingHosts.map((host) => ({
          host: host.hostname,
          result: {
            admitted: false,
            missing: ["probe-timeout"] as string[],
            reason:
              `admission probe did not finish within ${ADMISSION_RECONCILE_TIMEOUT_MS}ms`,
          },
        }))
      : probeResults;
    const checkedAt = new Date(this.now()).toISOString();

    return {
      ok: true,
      acceptedRevision,
      result: {
        reconciled: true,
        reason: parsed.data.reason ?? null,
        advanced,
        probes: probes.map(({ host, result }) => ({
          host,
          admitted: result.admitted,
          reason: result.reason,
        })),
      },
      apply: () => {
        for (const item of advanced) {
          this.store.setHostState(item.host, item.to);
        }
        for (const { host, result } of probes) {
          const record = this.store.getHost(host);
          if (!record) continue;
          const capabilityOk = result.admitted;
          const next = {
            ...record,
            capabilityOk,
            capabilityReason: result.reason,
            capabilityCheckedAt: checkedAt,
          };
          if (
            capabilityOk &&
            record.healthStorage &&
            record.healthRunner &&
            record.healthOffload
          ) {
            next.state = "available";
            next.enrolling = false;
          } else if (capabilityOk) {
            next.capabilityReason = `admitted, but health checks failing: ${
              unhealthyChecks(record).join(", ")
            }`;
          }
          this.store.upsertHost(next);
        }
        this.syncClusterObserved();
        this.scheduler.reconcile();
      },
    };
  }

  private jobRetry(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = JobRetryArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const job = this.store.getJob(parsed.data.jobId);
    if (!job) return { ok: false, detail: `unknown job: ${parsed.data.jobId}` };
    if (!job.infraFailure) {
      return { ok: false, detail: `job ${job.id} is not a typed infra failure` };
    }
    const nextAttempt = job.attempt + 1;
    return {
      ok: true,
      result: { jobId: job.id, attempt: nextAttempt },
      apply: () => {
        this.store.upsertJob({
          ...job,
          attempt: nextAttempt,
          stage: "queued",
          infraFailure: false,
        });
      },
    };
  }

  private ciReconcile(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = OptionalHostArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const targets = parsed.data.host
      ? [parsed.data.host]
      : this.store.listHosts().map((h) => h.hostname);
    const reconciled: string[] = [];
    return {
      ok: true,
      result: { hosts: targets, serviceState: "aligned" },
      apply: () => {
        for (const hostname of targets) {
          const host = this.store.getHost(hostname);
          if (!host) continue;
          this.store.upsertHost({ ...host, ciJobsRunning: 0 });
          reconciled.push(hostname);
        }
      },
    };
  }

  private recallSpill(
    args: Record<string, unknown>,
  ):
    | { ok: true; result: unknown; apply: () => void }
    | { ok: false; detail: string } {
    const parsed = OptionalHostArgsSchema.safeParse(args);
    if (!parsed.success) {
      return { ok: false, detail: parsed.error.message };
    }
    const lease = this.store.getLease();
    if (parsed.data.host && lease.host && lease.host !== parsed.data.host) {
      return {
        ok: false,
        detail: `spill lease is on ${lease.host}, not ${parsed.data.host}`,
      };
    }
    return {
      ok: true,
      result: { recalled: true, host: lease.host },
      apply: () => {
        this.scheduler.recallSpill(parsed.data.host);
      },
    };
  }

  private syncClusterObserved(): void {
    const hosts = this.store.listHosts();
    const meta = this.store.getMetaSnapshot();
    if (hosts.some((h) => h.state === "degraded")) {
      this.store.setClusterState(meta.desired, "degraded");
      return;
    }
    if (hosts.some((h) => h.state === "maintenance" || h.state === "draining")) {
      this.store.setClusterState(meta.desired, "maintenance");
      return;
    }
    if (hosts.some((h) => h.state === "restoring")) {
      this.store.setClusterState(meta.desired, "restoring");
      return;
    }
    this.store.setClusterState(meta.desired, "available");
  }
}

export function coerceExpectedRevision(value: unknown): number | null {
  if (typeof value === "number" && Number.isInteger(value)) {
    return value;
  }
  if (typeof value === "string" && value.trim() !== "") {
    const parsed = Number(value);
    if (Number.isInteger(parsed)) return parsed;
  }
  return null;
}

function validateVerbArgs(
  verb: TransitionVerb,
  args: Record<string, unknown>,
): { ok: true } | { ok: false; detail: string } {
  const schema = verbArgsSchema(verb);
  const result = schema.safeParse(args);
  if (!result.success) {
    return { ok: false, detail: result.error.message };
  }
  return { ok: true };
}

function verbArgsSchema(verb: TransitionVerb): z.ZodType<unknown> {
  switch (verb) {
    case "box-drain":
    case "box-restore":
      return BoxHostArgsSchema;
    case "host-quarantine":
    case "host-unquarantine":
      return QuarantineArgsSchema;
    case "admission-reconcile":
      return AdmissionReconcileArgsSchema;
    case "job-retry":
      return JobRetryArgsSchema;
    case "ci-reconcile":
    case "recall-spill":
      return OptionalHostArgsSchema;
    default:
      return z.object({});
  }
}

export function grantFallbackLease(
  store: ControllerStore,
  params: { host: string; reason: string; ttlMs?: number },
  now: () => number = () => Date.now(),
): void {
  const expiresAt = new Date(now() + (params.ttlMs ?? 5 * 60 * 1000)).toISOString();
  store.setLease({
    active: true,
    expiresAt,
    host: params.host,
    reason: params.reason,
  });
}

export type { TransitionIntent };
