import { randomUUID } from "node:crypto";
import type {
  ExtensionAPI,
  ExtensionCommandContext,
  SessionManager,
} from "@earendil-works/pi-coding-agent";
import { Type } from "typebox";

const STATE_ENTRY = "pi-fresh-loop";
const STATUS_KEY = "pi-fresh-loop";
const STOP_TOOL = "loop_stop";

interface RunningMarker {
  version: 1;
  kind: "running";
  runId: string;
  prompt: string;
  iteration: number;
  modelProvider?: string;
  modelId?: string;
}

interface StopMarker {
  version: 1;
  kind: "stop";
  runId: string;
  iteration: number;
  source: "command" | "agent" | "system";
  reason?: string;
}

type LoopMarker = RunningMarker | StopMarker;
type BranchReader = Pick<SessionManager, "getBranch">;
type NewSessionOptions = NonNullable<Parameters<ExtensionCommandContext["newSession"]>[0]>;
type ReplacedSessionContext = Parameters<NonNullable<NewSessionOptions["withSession"]>>[0];

function isLoopMarker(value: unknown): value is LoopMarker {
  if (!value || typeof value !== "object") return false;
  const marker = value as Partial<LoopMarker>;
  return (
    marker.version === 1 &&
    typeof marker.runId === "string" &&
    typeof marker.iteration === "number" &&
    (marker.kind === "running" || marker.kind === "stop")
  );
}

export function getLatestLoopMarker(sessionManager: BranchReader): LoopMarker | undefined {
  const branch = sessionManager.getBranch();
  for (let index = branch.length - 1; index >= 0; index--) {
    const entry = branch[index];
    if (entry.type !== "custom" || entry.customType !== STATE_ENTRY) continue;
    if (isLoopMarker(entry.data)) return entry.data;
  }
  return undefined;
}

function getRunningMarker(sessionManager: BranchReader): RunningMarker | undefined {
  const marker = getLatestLoopMarker(sessionManager);
  return marker?.kind === "running" ? marker : undefined;
}

function getLastAssistantFailure(
  sessionManager: BranchReader,
): { stopReason: "aborted" | "error"; errorMessage?: string } | undefined {
  const branch = sessionManager.getBranch();
  for (let index = branch.length - 1; index >= 0; index--) {
    const entry = branch[index];
    if (entry.type !== "message" || entry.message.role !== "assistant") continue;
    if (entry.message.stopReason === "aborted" || entry.message.stopReason === "error") {
      return {
        stopReason: entry.message.stopReason,
        errorMessage: entry.message.errorMessage,
      };
    }
    return undefined;
  }
  return undefined;
}

function appendSystemStop(sessionManager: BranchReader, marker: RunningMarker, reason: string): void {
  // ReplacedSessionContext intentionally exposes a read-only SessionManager type.
  // Custom state does not affect model context, so direct append is safe here and
  // lets orchestration close state after its originating extension instance is stale.
  (sessionManager as SessionManager).appendCustomEntry(STATE_ENTRY, {
    version: 1,
    kind: "stop",
    runId: marker.runId,
    iteration: marker.iteration,
    source: "system",
    reason,
  } satisfies StopMarker);
}

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

function notifySafely(
  ctx: ExtensionCommandContext | ReplacedSessionContext | undefined,
  message: string,
  level: "info" | "warning" | "error",
): void {
  try {
    ctx?.ui.notify(message, level);
  } catch {
    // Session replacement can invalidate a context while an error unwinds.
  }
}

export default function freshLoopExtension(pi: ExtensionAPI) {
  function requestStop(
    marker: RunningMarker,
    source: StopMarker["source"],
    reason?: string,
  ): void {
    pi.appendEntry(STATE_ENTRY, {
      version: 1,
      kind: "stop",
      runId: marker.runId,
      iteration: marker.iteration,
      source,
      reason,
    } satisfies StopMarker);
  }

  pi.on("session_start", async (event, ctx) => {
    const marker = getRunningMarker(ctx.sessionManager);
    if (!marker) return;

    if (event.reason === "new") {
      // newSession.setup writes the live marker before the replacement runtime
      // emits session_start. Re-select the model captured by the originating
      // command through this replacement extension instance; the old pi is stale.
      if (!marker.modelProvider || !marker.modelId) return;
      if (ctx.model?.provider === marker.modelProvider && ctx.model.id === marker.modelId) return;
      const model = ctx.modelRegistry.find(marker.modelProvider, marker.modelId);
      try {
        if (!model || !(await pi.setModel(model))) {
          requestStop(marker, "system", "model-unavailable");
        }
      } catch {
        requestStop(marker, "system", "model-unavailable");
      }
      return;
    }

    // Other start reasons mean a stopped process, reload, fork, or manual resume
    // reopened an orphaned iteration.
    requestStop(marker, "system", "session-restarted");
    ctx.ui.setStatus(STATUS_KEY, undefined);
    ctx.ui.notify("Previous loop stopped because session was restarted or resumed.", "warning");
  });

  pi.registerTool({
    name: STOP_TOOL,
    label: "Stop Loop",
    description:
      "Stop active /loop after current iteration. Call when prompt's stopping condition is satisfied.",
    promptSnippet: "Stop active /loop when its prompt-defined stopping condition is satisfied",
    promptGuidelines: [
      "When an active /loop prompt says to stop under a condition and that condition is satisfied, call loop_stop. Do not call loop_stop while more loop work remains.",
    ],
    parameters: Type.Object({}),
    async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
      const marker = getRunningMarker(ctx.sessionManager);
      if (!marker) {
        return {
          content: [{ type: "text" as const, text: "No loop is running." }],
          details: { stopped: false },
        };
      }

      requestStop(marker, "agent");
      ctx.ui.setStatus(STATUS_KEY, `loop: stopping after #${marker.iteration}`);
      return {
        content: [
          {
            type: "text" as const,
            text: "Loop stop requested. Finish this iteration; no next iteration will start.",
          },
        ],
        details: { stopped: true, runId: marker.runId, iteration: marker.iteration },
      };
    },
  });

  pi.registerCommand("loop", {
    description: "Repeat a prompt in fresh sessions. Usage: /loop <prompt> or /loop stop",
    getArgumentCompletions: (prefix) =>
      "stop".startsWith(prefix.trim().toLowerCase())
        ? [{ value: "stop", label: "stop", description: "Stop after current iteration" }]
        : null,
    handler: async (args, ctx) => {
      if (args.trim().toLowerCase() === "stop") {
        const marker = getRunningMarker(ctx.sessionManager);
        if (!marker) {
          ctx.ui.notify("No loop is running.", "info");
          return;
        }

        requestStop(marker, "command");
        ctx.ui.setStatus(STATUS_KEY, `loop: stopping after #${marker.iteration}`);
        ctx.ui.notify(`Loop will stop after iteration ${marker.iteration}.`, "info");
        return;
      }

      if (!args.trim()) {
        ctx.ui.notify("Usage: /loop <prompt> or /loop stop", "warning");
        return;
      }

      if (!ctx.isIdle()) {
        ctx.ui.notify("Agent busy. Start loop after current run finishes.", "warning");
        return;
      }

      const existing = getRunningMarker(ctx.sessionManager);
      if (existing) {
        ctx.ui.notify(
          `Loop already running at iteration ${existing.iteration}. Use /loop stop first.`,
          "warning",
        );
        return;
      }

      if (!ctx.model) {
        ctx.ui.notify("No model selected.", "error");
        return;
      }

      const prompt = args;
      const runId = randomUUID();
      const modelProvider = ctx.model.provider;
      const modelId = ctx.model.id;
      let iteration = 1;
      let currentCtx: ExtensionCommandContext = ctx;

      while (true) {
        let iterationCtx: ReplacedSessionContext | undefined;
        const runningMarker: RunningMarker = {
          version: 1,
          kind: "running",
          runId,
          prompt,
          iteration,
          modelProvider,
          modelId,
        };

        try {
          const result = await currentCtx.newSession({
            setup: async (sessionManager) => {
              sessionManager.appendCustomEntry(STATE_ENTRY, runningMarker);
            },
            withSession: async (replacementCtx) => {
              iterationCtx = replacementCtx;
              replacementCtx.ui.setStatus(STATUS_KEY, `loop: #${iteration}`);

              const latest = getLatestLoopMarker(replacementCtx.sessionManager);
              if (latest?.kind === "stop" && latest.runId === runId) return;
              if (
                replacementCtx.model?.provider !== modelProvider ||
                replacementCtx.model?.id !== modelId
              ) {
                appendSystemStop(replacementCtx.sessionManager, runningMarker, "model-unavailable");
                return;
              }

              await replacementCtx.sendUserMessage(prompt);
            },
          });

          if (result.cancelled) {
            const active = getRunningMarker(currentCtx.sessionManager);
            if (active?.runId === runId) appendSystemStop(currentCtx.sessionManager, active, "cancelled");
            currentCtx.ui.setStatus(STATUS_KEY, undefined);
            currentCtx.ui.notify("Loop stopped: new session was cancelled.", "warning");
            return;
          }
        } catch (error) {
          const activeCtx = iterationCtx ?? currentCtx;
          try {
            const active = getRunningMarker(activeCtx.sessionManager);
            if (active?.runId === runId) appendSystemStop(activeCtx.sessionManager, active, "failure");
            activeCtx.ui.setStatus(STATUS_KEY, undefined);
          } catch {
            // Context may already be stale after external session replacement.
          }
          notifySafely(activeCtx, `Loop stopped: ${errorMessage(error)}`, "error");
          return;
        }

        if (!iterationCtx) {
          notifySafely(currentCtx, "Loop stopped: replacement session unavailable.", "error");
          return;
        }

        try {
          const latest = getLatestLoopMarker(iterationCtx.sessionManager);
          if (latest?.kind === "stop" && latest.runId === runId) {
            iterationCtx.ui.setStatus(STATUS_KEY, undefined);
            if (latest.source === "system" && latest.reason === "model-unavailable") {
              iterationCtx.ui.notify(
                `Loop stopped: model ${modelProvider}/${modelId} is unavailable.`,
                "error",
              );
            } else {
              iterationCtx.ui.notify(`Loop stopped after ${iteration} iteration(s).`, "info");
            }
            return;
          }

          const failure = getLastAssistantFailure(iterationCtx.sessionManager);
          if (failure) {
            appendSystemStop(iterationCtx.sessionManager, runningMarker, failure.stopReason);
            iterationCtx.ui.setStatus(STATUS_KEY, undefined);
            const detail = failure.errorMessage ? `: ${failure.errorMessage}` : "";
            iterationCtx.ui.notify(
              failure.stopReason === "aborted"
                ? "Loop stopped: iteration aborted."
                : `Loop stopped: iteration failed${detail}`,
              failure.stopReason === "aborted" ? "warning" : "error",
            );
            return;
          }

          currentCtx = iterationCtx;
          iteration++;
        } catch (error) {
          notifySafely(iterationCtx, `Loop stopped: ${errorMessage(error)}`, "error");
          return;
        }
      }
    },
  });
}