import { getAgentDir, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
import { Buffer } from "node:buffer";
import { join } from "node:path";
import { discoverDefaultAgents, type DiscoverAgentsResult } from "./agents.js";
import { loadConfig, type LoadConfigResult } from "./config.js";
import { registerSubagentsUi } from "./dashboard.js";
import { sanitizeTerminalText } from "./job-status.js";
import { truncateUtf8 } from "./output.js";
import { PiRpcSessionRunner, type SessionRunner } from "./session-runner.js";
import { SubagentManager, type SubagentManagerEvent } from "./subagent-manager.js";
import { registerSubagentTools, type ToolServices } from "./tools.js";
import { REPORT_MAX_BYTES, type AgentProfile, type JobRequest, type SubagentReport } from "./types.js";

export {
  REPORT_DEDUPE_MAX_BYTES,
  REPORT_DEDUPE_MAX_ITEMS,
  REPORT_ID_MAX_BYTES,
  REPORT_INBOX_MAX_BYTES,
  REPORT_INBOX_MAX_ITEMS,
  REPORT_MAX_BYTES,
  REPORT_RECORD_FIXED_BYTES,
  SESSION_TOMBSTONE_MAX_ITEMS,
  TASK_MAX_BYTES,
} from "./types.js";

const HELP_MESSAGE_TYPE = "simple-subagents-help";
const SESSION_FAILURE_MESSAGE_TYPE = "simple-subagents-session-failed";
const HELP_SESSION_ID_MAX_BYTES = 512;
const SESSION_FAILURE_MESSAGE_MAX_BYTES = 1024;
const HELP_DELIVERY_FAILURE = "Unable to deliver subagent help context.";
const SESSION_FAILURE_DELIVERY_FAILURE = "Unable to deliver subagent failure context.";

const safeSessionId = (sessionId: string): string => {
  const sanitized = sanitizeTerminalText(sessionId).replace(/\s+/gu, " ").trim();
  return truncateUtf8(sanitized, HELP_SESSION_ID_MAX_BYTES).text || "unknown";
};

const safeGeneration = (generation: number): string =>
  Number.isSafeInteger(generation) && generation >= 0 ? String(generation) : "unknown";

const formatHelpMessage = (report: SubagentReport): { sessionId: string; content: string } => {
  const sessionId = safeSessionId(report.sessionId);
  const generation = safeGeneration(report.generation);
  const prefix = `Subagent help request\nSession: ${sessionId}\nGeneration: ${generation}\n\n`;
  const helpBudget = Math.max(0, REPORT_MAX_BYTES - Buffer.byteLength(prefix, "utf8"));
  const help = truncateUtf8(sanitizeTerminalText(report.message), helpBudget).text;
  return { sessionId, content: `${prefix}${help}` };
};

const formatSessionFailureMessage = (
  event: Extract<SubagentManagerEvent, { type: "session_failed" }>,
): { sessionId: string; content: string } => {
  const sessionId = safeSessionId(event.sessionId);
  const content = `Subagent session exited unexpectedly.\nSession: ${sessionId}\nGeneration: ${safeGeneration(event.generation)}\nPartial result ready: ${event.partialResultReady === true ? "yes" : "no"}.`;
  return { sessionId, content: truncateUtf8(content, SESSION_FAILURE_MESSAGE_MAX_BYTES).text };
};

export interface ExtensionDependencies {
  createManager?: (runner: SessionRunner) => SubagentManager;
  loadConfig?: (path: string) => Promise<LoadConfigResult>;
  discoverProfiles?: () => Promise<DiscoverAgentsResult>;
  getAgentDir?: () => string;
}

export function createSimpleSubagentsExtension(dependencies: ExtensionDependencies = {}): (pi: ExtensionAPI) => void {
  return (pi) => {
    const runner = new PiRpcSessionRunner();
    const manager = dependencies.createManager?.(runner) ?? new SubagentManager({ runner });
    const readConfig = dependencies.loadConfig ?? loadConfig;
    const discoverProfiles = dependencies.discoverProfiles ?? discoverDefaultAgents;
    const resolveAgentDir = dependencies.getAgentDir ?? getAgentDir;
    let config = { confirmWrites: false, allowThinkingOverrides: false };
    let profiles = new Map<string, AgentProfile>();
    let toolsRegistered = false;
    let shutdown: Promise<void> | undefined;
    let replacementGuard: Promise<{ cancel: true } | undefined> | undefined;
    let activeHelpListener: object | undefined;
    let cleanupHelpListener: (() => void) | undefined;
    let cleanupUi: (() => void) | undefined = registerSubagentsUi(pi, manager);

    const clearUi = (): void => {
      const cleanup = cleanupUi;
      cleanupUi = undefined;
      cleanup?.();
    };

    const clearHelpListener = (): void => {
      activeHelpListener = undefined;
      const cleanup = cleanupHelpListener;
      cleanupHelpListener = undefined;
      cleanup?.();
    };

    const subscribeToHelp = (ctx: ExtensionContext): void => {
      clearHelpListener();
      if (typeof manager.subscribeEvents !== "function") return;
      const listenerToken = {};
      activeHelpListener = listenerToken;
      const unsubscribe = manager.subscribeEvents((event) => {
        if (activeHelpListener !== listenerToken) return;
        const notify = (message: string): void => {
          if (!ctx.hasUI) return;
          try {
            ctx.ui.notify(message, "warning");
          } catch {}
        };

        if (event.type === "help_waiting") {
          const help = formatHelpMessage(event.report);
          try {
            pi.sendMessage({
              customType: HELP_MESSAGE_TYPE,
              content: help.content,
              display: true,
            }, { deliverAs: "steer", triggerTurn: true });
          } catch {
            notify(HELP_DELIVERY_FAILURE);
            return;
          }
          notify(`Subagent ${help.sessionId} requested help.`);
          return;
        }

        if (event.type !== "session_failed") return;
        const failure = formatSessionFailureMessage(event);
        try {
          pi.sendMessage({
            customType: SESSION_FAILURE_MESSAGE_TYPE,
            content: failure.content,
            display: true,
          }, { deliverAs: "nextTurn" });
        } catch {
          notify(SESSION_FAILURE_DELIVERY_FAILURE);
          return;
        }
        notify(`Subagent ${failure.sessionId} exited unexpectedly.`);
      });
      cleanupHelpListener = () => {
        if (activeHelpListener === listenerToken) activeHelpListener = undefined;
        unsubscribe();
      };
    };

    const confirmParentReplacement = (ctx: ExtensionContext): Promise<{ cancel: true } | undefined> => {
      if (!manager.hasOpenChildren()) return Promise.resolve(undefined);
      if (replacementGuard) return replacementGuard;

      const operation = (async (): Promise<{ cancel: true } | undefined> => {
        if (!ctx.hasUI || typeof ctx.ui?.confirm !== "function") return { cancel: true };
        let confirmed: unknown;
        try {
          confirmed = await ctx.ui.confirm(
            "Close child sessions?",
            "All child sessions will close before replacing this parent session.",
          );
        } catch {
          return { cancel: true };
        }
        if (confirmed !== true) return { cancel: true };
        try {
          await manager.closeAll();
        } catch {
          return { cancel: true };
        }
        if (manager.hasOpenChildren()) return { cancel: true };
        return undefined;
      })();
      replacementGuard = operation;
      void operation.finally(() => {
        if (replacementGuard === operation) replacementGuard = undefined;
      });
      return operation;
    };

    const services: ToolServices = {
      manager,
      getProfiles: async () => profiles,
      confirmWritable: async (requests: readonly JobRequest[], ctx: ExtensionContext) => {
        if (!config.confirmWrites || requests.length === 0) return "approved";
        if (!ctx.hasUI) return "unavailable";
        const pluralSuffix = requests.length === 1 ? "" : "s";
        const confirmed = await ctx.ui.confirm(
          "Allow writable subagents?",
          `Allow ${requests.length} writable background job${pluralSuffix}? Their work should not overlap.`,
        );
        if (confirmed) return "approved";
        return "declined";
      },
      defaults: (ctx) => {
        let parentModel: string | undefined;
        if (ctx.model) {
          parentModel = `${ctx.model.provider}/${ctx.model.id}`;
        } else {
          parentModel = undefined;
        }
        return {
          cwd: ctx.cwd,
          parentModel,
          thinkingLevel: ctx.thinkingLevel,
        };
      },
    };

    pi.on("session_start", async (_event, ctx) => {
      const priorShutdown = shutdown;
      if (priorShutdown) {
        await priorShutdown;
        if (shutdown === priorShutdown) {
          shutdown = undefined;
          cleanupUi = registerSubagentsUi(pi, manager);
        }
      }
      subscribeToHelp(ctx);
      const [loadedConfig, discovered] = await Promise.all([
        readConfig(join(resolveAgentDir(), "simple-subagents.json")),
        discoverProfiles(),
      ]);
      config = loadedConfig.config;
      profiles = new Map(discovered.agents.map((profile) => [profile.name, profile]));
      if (!toolsRegistered) {
        registerSubagentTools(pi, services, config.allowThinkingOverrides);
        toolsRegistered = true;
      }
      if (loadedConfig.warning) ctx.ui.notify(loadedConfig.warning, "warning");
      for (const diagnostic of discovered.diagnostics) ctx.ui.notify(diagnostic, "warning");
    });

    pi.on("session_before_switch", async (_event, ctx) => confirmParentReplacement(ctx));
    pi.on("session_before_fork", async (_event, ctx) => confirmParentReplacement(ctx));

    pi.on("session_shutdown", async () => {
      clearHelpListener();
      clearUi();
      if (shutdown === undefined) {
        shutdown = Promise.resolve()
          .then(() => manager.shutdown())
          .catch(() => {});
      }
      await shutdown;
    });
  };
}

export default createSimpleSubagentsExtension();