import type {
  ExtensionAPI,
  ExtensionContext,
  Theme,
} from "@earendil-works/pi-coding-agent";
import { truncateToWidth } from "@earendil-works/pi-tui";
import type { ContextTransitionRuntime } from "./context-transition.ts";
import type { WorkflowRun } from "./domain.ts";
import { deriveRunLabels } from "./run-label.ts";

const WIDGET_KEY = "workflow-active-run";
const CORE_STAGES = new Set([
  "discover",
  "design",
  "prototype",
  "implement",
  "verify",
  "review",
]);

export interface WorkflowProgress {
  runId: string;
  label: string;
  stage: string;
  completed: number;
  total: number;
  gate?: string;
}

function activityTime(run: WorkflowRun) {
  return Math.max(
    run.lease?.heartbeatAt ?? 0,
    ...run.gates.map((gate) => gate.openedAt),
    ...run.stages.flatMap((stage) => [stage.startedAt ?? 0, stage.finishedAt ?? 0]),
  );
}

function readableGate(kind: string) {
  return kind
    .split("-")
    .map((part) => `${part.slice(0, 1).toUpperCase()}${part.slice(1)}`)
    .join(" ");
}

export function newestAttachedActiveRun(
  runs: WorkflowRun[],
  sessionId: string | undefined,
) {
  if (!sessionId) return undefined;
  return runs
    .filter(
      (run) =>
        (run.status === "running" || run.status === "waiting") &&
        run.parentSessionIds.includes(sessionId),
    )
    .sort((left, right) => activityTime(right) - activityTime(left))[0];
}

export function projectWorkflowProgress(
  runs: WorkflowRun[],
  sessionId: string | undefined,
): WorkflowProgress | undefined {
  const run = newestAttachedActiveRun(runs, sessionId);
  if (!run) return undefined;
  const active = [...run.stages]
    .reverse()
    .find((stage) => stage.status === "running" || stage.status === "waiting");
  const gate = [...run.gates].reverse().find((entry) => entry.status === "open");
  const label = deriveRunLabels(runs)[runs.indexOf(run)] ?? run.runId.slice(0, 8);
  const stages = run.stages.filter((stage) => CORE_STAGES.has(stage.kind));
  return {
    runId: run.runId,
    label,
    stage: String(active?.kind ?? run.status).toUpperCase(),
    completed: stages.filter((stage) => stage.status === "completed").length,
    total: stages.length,
    ...(gate ? { gate: readableGate(gate.kind) } : {}),
  };
}

export function renderWorkflowProgress(
  progress: WorkflowProgress,
  width: number,
  theme?: Pick<Theme, "fg">,
) {
  const paint = (token: "accent" | "dim" | "muted" | "text", value: string) =>
    theme ? theme.fg(token, value) : value;
  const base = `${paint("accent", "workflow")}  ${paint("text", progress.label)}  ${paint("muted", progress.stage)}  ${progress.completed}/${progress.total}`;
  const full = progress.gate ? `${base}  ${paint("dim", `Gate: ${progress.gate}`)}` : base;
  return truncateToWidth(full, Math.max(1, width), "");
}

export async function refreshWorkflowProgress(
  runtime: ContextTransitionRuntime,
  ctx: ExtensionContext,
) {
  if (ctx.mode !== "tui") return;
  const progress = projectWorkflowProgress(
    await runtime.store.listRuns(),
    ctx.sessionManager.getSessionId(),
  );
  if (!progress) {
    ctx.ui.setWidget(WIDGET_KEY, undefined);
    return;
  }
  ctx.ui.setWidget(WIDGET_KEY, (_tui, theme) => ({
    render(width: number) {
      return [renderWorkflowProgress(progress, width, theme)];
    },
    invalidate() {},
  }));
}

export function registerWorkflowProgress(
  pi: ExtensionAPI,
  runtime: ContextTransitionRuntime,
) {
  pi.on("session_start", async (_event, ctx) => {
    await refreshWorkflowProgress(runtime, ctx);
  });
  pi.on("session_shutdown", async (_event, ctx) => {
    if (ctx.mode === "tui") ctx.ui.setWidget(WIDGET_KEY, undefined);
  });
}