import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { existsSync, readFileSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";
import { fileURLToPath } from "node:url";

/**
 * pi-oc-style-agent-switcher
 *
 * Keyboard-driven switching between primary Pi agent profiles.
 *
 * Source of truth:
 *   settings.json → agent.<name>
 *
 * Entries under `agent` participate in the cycle unless they explicitly set
 * `mode: "subagent"`. The extension reads and applies `agent.<name>.model`
 * on cycle switches, and `agent.<name>.thinking` on startup (by matching the
 * current model) and whenever the model changes.
 */

interface SwitcherConfig {
  cyclePrev?: string;
  cycleNext?: string;
}

interface AgentProfile {
  mode?: "primary" | "subagent";
  model?: string;
  thinking?: string;
  variant?: string;
  tools?: string[] | Record<string, boolean>;
  /** ANSI-compatible hex color (e.g. "#6FCF97") for the status line. Auto-assigned if omitted. */
  color?: string;
  /** Icon shown before the agent name in the status line. Auto-assigned if omitted. */
  icon?: string;
  /** Enable the orchestrator's stronger ULW reminder by default for this session. */
  ulw?: boolean;
}

export { type AgentProfile };

/** Session custom-entry type consumed by pi-sidebar and other extensions. */
export const ULW_STATE_ENTRY_TYPE = "pi-oc-style-agent-switcher:ulw";

const TEMPLATES_DIR = fileURLToPath(
  new URL("./templates/agents", import.meta.url),
);
const GLOBAL_AGENTS_DIR = join(homedir(), ".pi", "agent", "agents");

const DEFAULT_AGENT_TEMPLATES = ["orchestrator", "build"];
const MAIN_SESSION_PROMPT_PROFILES = new Set(["orchestrator"]);
const FILE_OPERATION_TOOLS = new Set(["read", "Read", "write", "Write", "edit", "Edit"]);
const ORCHESTRATOR_WORKFLOW_REMINDER =
  "Workflow reminder (not an automatic action): choose the lightest workflow; direct execution is okay for isolated, low-risk work. Otherwise, plan lanes and dependencies, dispatch appropriate background specialists, track their results, reconcile findings, and verify the outcome.";
const ORCHESTRATOR_ULW_REMINDER =
  "ULW mode is active (prompt-mediated only; this extension hook only injects reminders and never directly spawns subagents): the orchestrator is a scheduler/coordinator first, not the default implementation worker. Before substantive work, perform a delegation check and write a short work graph covering independent lanes, dependency order, file ownership, and validation ownership. Direct work is only for one isolated, clear, low-risk action where delegation costs more; route multi-file or multi-stage work, broad discovery, unfamiliar research, implementation, tests, UI/design, architecture/review, and complex debugging to appropriate specialists. Prefer installed `@esso0428/pi-subagents` tools when available; otherwise inspect the current session's tools and choose a compatible alternative. Do not hard-bind OpenCode `task` or an unavailable package/API. Prefer background execution for independent lanes; do not immediately wait after spawning, do not use `wait_for_user` to wait for children, continue only non-overlapping coordination work, then end the turn and rely on available completion notifications/events. Wait or status-check only when a dependent decision truly needs the result, using whatever runtime provides. Track child IDs, state, ownership, dependencies, and terminal results; avoid overlapping writes and duplicate objectives. Reconcile every output against the original goal—specialist output is input, not final truth—and verify the final state. Delegation Recovery: an aborted, errored, stopped, partial, or incomplete child is not a reason to take over. Inspect the failure and, if recoverable, adjust only the failing constraint (raise max turns, narrow prompt/scope, fix tool or extension access, change specialist/model, resume/revive, or respawn); never reissue unchanged work. Take over only after justified recovery fails, no subagent capability exists, or safety/user choice requires it. Choose child budgets proportionately; do not give broad uncertain work a tiny max_turns, and recover when a child reaches its budget rather than silently working locally. If no subagent capability exists, report that limitation and proceed proportionately.";
const FILE_OPERATION_NUDGE =
  "Delegation reminder (not an automatic action): reconsider whether the next non-trivial implementation or research step should go to a specialist.";
const FILE_OPERATION_ULW_NUDGE =
  "ULW file-operation reminder (prompt-mediated only; this extension hook only injects reminders and never directly spawns subagents): after this read/write/edit, re-run the scheduler-first delegation check. For substantive work, use a short ownership/dependency/validation graph; route non-trivial work to an appropriate specialist, prefer installed `@esso0428/pi-subagents` or a compatible current-session alternative, and use background lanes for independent work without immediately waiting or using `wait_for_user`. Track and reconcile child state/results, avoid overlapping writes and duplicate objectives, and verify the final state. An aborted/error/stopped/partial child triggers Delegation Recovery—inspect it and adjust only the failing constraint (budget, scope, access, specialist/model, resume/revive, or respawn); do not reissue unchanged work or take over unless justified recovery fails, no capability exists, or safety/user choice requires it. Choose proportionate max_turns and recover budget exhaustion. This hook never spawns subagents directly.";

let currentIndex = 0;
let profiles: Record<string, AgentProfile> = {};
let cycle: string[] = [];

function readJson(path: string): Record<string, any> {
  try {
    return existsSync(path) ? JSON.parse(readFileSync(path, "utf-8")) : {};
  } catch {
    return {};
  }
}

/** Project values override global values field-by-field. */
function readMergedSettings(): Record<string, any> {
  const global = readJson(join(homedir(), ".pi", "agent", "settings.json"));
  const project = readJson(join(process.cwd(), ".pi", "settings.json"));
  const names = new Set([
    ...Object.keys(global.agent ?? {}),
    ...Object.keys(project.agent ?? {}),
  ]);
  const agent: Record<string, AgentProfile> = {};
  for (const name of names) {
    agent[name] = {
      ...(global.agent?.[name] ?? {}),
      ...(project.agent?.[name] ?? {}),
    };
  }
  return { agent };
}

function loadConfig(): SwitcherConfig {
  const config = readJson(join(homedir(), ".pi", "agent", "agent-switcher.json"));
  return config && typeof config === "object" ? config : {};
}

/** Parse the `model` field from a markdown file's YAML frontmatter. */
function parseModelFromFrontmatter(filePath: string): string | undefined {
  try {
    const content = readFileSync(filePath, "utf-8");
    const match = content.match(/^---\n([\s\S]*?)\n---/);
    if (!match) return undefined;
    const yaml = match[1];
    for (const line of yaml.split("\n")) {
      const kv = line.match(/^model:\s*"?([^\s"]+)"?$/);
      if (kv) return kv[1];
    }
  } catch { /* ignore */ }
  return undefined;
}

/**
 * For each packaged default agent template (orchestrator, build), check
 * whether a project/global markdown file exists and add a fallback profile
 * entry if none is already configured in settings.json → agent.
 *
 * Packaged templates remain immutable fallbacks. They are intentionally not
 * copied into the user's global agents directory, so package upgrades apply
 * to users who have not created a project/global override.
 */
function discoverAgentFiles(profiles: Record<string, AgentProfile>): Record<string, AgentProfile> {
  const result = { ...profiles };
  for (const name of DEFAULT_AGENT_TEMPLATES) {
    if (name in result) continue;
    const filePath = resolveAgentPromptPath(name);
    if (!filePath) continue;
    const model = parseModelFromFrontmatter(filePath);
    result[name] = model ? { model } : {};
  }
  return result;
}

function reloadProfiles(): void {
  const settings = readMergedSettings();
  profiles = discoverAgentFiles(settings.agent ?? {});
  assignAgentColors(profiles);
  cycle = Object.entries(profiles)
    .filter(([, profile]) => profile?.mode !== "subagent")
    .map(([name]) => name);
  if (currentIndex >= cycle.length) currentIndex = 0;
}

function selectedName(): string | undefined {
  return cycle[currentIndex];
}

const ANSI_RESET = "\x1b[0m";

function hexToAnsiFg(hex: string): string {
  const r = parseInt(hex.slice(1, 3), 16);
  const g = parseInt(hex.slice(3, 5), 16);
  const b = parseInt(hex.slice(5, 7), 16);
  return `\x1b[38;2;${r};${g};${b}m`;
}

/**
 * Light pastel palette (HSL ~75% lightness, ~70% saturation).
 * Indices 0-2 are reserved for named agents;
 * remaining slots are distributed evenly around the hue wheel.
 */
const NAMED_AGENT_COLORS: Record<string, string> = {
  orchestrator: "#6FCF97", // light green
  build: "#6BB5FF",         // light blue
  plan: "#FFD93D",          // light yellow
};

const NAMED_AGENT_ICONS: Record<string, string> = {
  orchestrator: "\uD83C\uDFAF", // 🎯
  build: "\uD83D\uDD27",         // 🔧
  plan: "\uD83D\uDCCB",          // 📋
};

const DEFAULT_AGENT_ICON = "\uD83E\uDDE0"; // 🧠

/** Auto-palette for unnamed agents (evenly spaced hues). */
const AUTO_PALETTE = [
  "#FF8A8A", // light red
  "#7EC8E3", // light cyan
  "#C084FC", // light purple
  "#FBBF24", // amber
  "#34D399", // emerald
  "#F472B6", // pink
  "#A78BFA", // violet
  "#60A5FA", // blue
  "#FB923C", // orange
  "#4ADE80", // green
  "#E879F9", // magenta
  "#2DD4BF", // teal
];

function assignAgentColors(profiles: Record<string, AgentProfile>): void {
  const named = Object.keys(NAMED_AGENT_COLORS);
  const entries = Object.entries(profiles);
  let autoIndex = 0;
  for (const [name, profile] of entries) {
    // Assign icon
    if (!profile.icon) {
      profile.icon = NAMED_AGENT_ICONS[name] ?? DEFAULT_AGENT_ICON;
    }
    // Assign color
    if (profile.color) continue; // explicit color wins
    if (named.includes(name)) {
      profile.color = NAMED_AGENT_COLORS[name];
    } else {
      profile.color = AUTO_PALETTE[autoIndex % AUTO_PALETTE.length];
      autoIndex++;
    }
  }
}

function formatKey(key: string): string {
  return key.replace(/left/i, "←").replace(/right/i, "→");
}

export function resolveEffectiveModelId(
  profile: AgentProfile | undefined,
  currentModelId?: string,
): string | undefined {
  return profile?.model || currentModelId;
}

export function formatStatusText(
  name: string | undefined,
  profile: AgentProfile | undefined,
  currentModelId?: string,
  ulw?: boolean,
): string {
  if (!name) return `${DEFAULT_AGENT_ICON} no primary agent`;
  const icon = profile?.icon ?? DEFAULT_AGENT_ICON;
  const colored = profile?.color
    ? `${hexToAnsiFg(profile.color)}${name}${ANSI_RESET}`
    : name;
  const ulwEnabled = ulw ?? (profile?.ulw !== false);
  const ulwText = name === "orchestrator"
    ? ` · ULW ${ulwEnabled ? "ON" : "OFF"}`
    : "";
  if (profile?.model) {
    const badge = profile.model === currentModelId ? "✓" : "✗";
    return `${icon} ${colored} (${badge}) ${profile.model}${ulwText}`;
  }
  return `${icon} ${colored}${currentModelId ? ` · ${currentModelId}` : ""}${ulwText}`;
}

export function shouldApplyConfiguredModel(
  profile: AgentProfile | undefined,
  currentModelId?: string,
): boolean {
  return !!profile?.model && profile.model !== currentModelId;
}

export function pickStatusModelId(
  profile: AgentProfile | undefined,
  currentModelId: string | undefined,
  applyOk: boolean,
): string | undefined {
  if (!applyOk) return currentModelId;
  return resolveEffectiveModelId(profile, currentModelId);
}

export function extractAgentPromptBody(markdown: string): string {
  return markdown.replace(/^---\n[\s\S]*?\n---\n*/, "").trim();
}

export function appendAgentPrompt(basePrompt: string, agentPrompt: string): string {
  if (!agentPrompt.trim()) return basePrompt;
  if (!basePrompt.trim()) return agentPrompt.trim();
  return `${basePrompt}\n\n${agentPrompt.trim()}`;
}

export function shouldInjectMainSessionPrompt(name: string | undefined): boolean {
  return !!name && MAIN_SESSION_PROMPT_PROFILES.has(name);
}

export function choosePreferredPromptPath(
  localPath?: string,
  globalPath?: string,
  packagePath?: string,
): string | undefined {
  return localPath || globalPath || packagePath;
}

function resolveAgentPromptPath(name: string): string | undefined {
  const projectCandidate = join(process.cwd(), ".pi", "agents", `${name}.md`);
  const globalCandidate = join(GLOBAL_AGENTS_DIR, `${name}.md`);
  const packageCandidate = join(TEMPLATES_DIR, `${name}.md`);
  return choosePreferredPromptPath(
    existsSync(projectCandidate) ? projectCandidate : undefined,
    existsSync(globalCandidate) ? globalCandidate : undefined,
    existsSync(packageCandidate) ? packageCandidate : undefined,
  );
}

function loadAgentPromptBody(name: string): string | undefined {
  const filePath = resolveAgentPromptPath(name);
  if (!filePath) return undefined;
  try {
    return extractAgentPromptBody(readFileSync(filePath, "utf-8"));
  } catch {
    return undefined;
  }
}

export default function (pi: ExtensionAPI) {
  let pendingFileOperationNudge = false;
  let sessionUlw = true;
  reloadProfiles();

  function readPersistedSessionUlw(ctx: any): boolean | undefined {
    try {
      const manager = ctx?.sessionManager;
      if (!manager) return undefined;
      const getEntries = (manager as {
        getBranch?: () => unknown;
        getEntries?: () => unknown;
      }).getBranch ?? (manager as {
        getBranch?: () => unknown;
        getEntries?: () => unknown;
      }).getEntries;
      if (typeof getEntries !== "function") return undefined;
      const entries = getEntries.call(manager);
      if (!Array.isArray(entries)) return undefined;
      for (let i = entries.length - 1; i >= 0; i--) {
        const entry = entries[i] as Record<string, unknown> | undefined;
        if (entry?.type !== "custom" || entry.customType !== ULW_STATE_ENTRY_TYPE) {
          continue;
        }
        const data = entry.data;
        if (!data || typeof data !== "object" || Array.isArray(data)) continue;
        const state = data as Record<string, unknown>;
        if (typeof state.sessionUlw === "boolean") return state.sessionUlw;
        // Older entries did not persist sessionUlw. Only orchestrator ON/OFF
        // entries carry an unambiguous value; N/A/build entries are ignored.
        if (state.profile === "orchestrator" && (state.status === "ON" || state.status === "OFF")) {
          return state.status === "ON";
        }
      }
    } catch { /* session history is best-effort */ }
    return undefined;
  }

  function initializeSessionUlw(ctx: any): void {
    sessionUlw = profiles.orchestrator?.ulw !== false;
    const persisted = readPersistedSessionUlw(ctx);
    if (persisted !== undefined) sessionUlw = persisted;
  }

  function publishUlwState(): void {
    try {
      const profile = selectedName();
      const enabled = profile === "orchestrator";
      pi.appendEntry(ULW_STATE_ENTRY_TYPE, {
        status: enabled ? (sessionUlw ? "ON" : "OFF") : "N/A",
        profile: profile ?? "",
        enabled,
        sessionUlw,
      });
    } catch { /* custom entries are best-effort */ }
  }

  const config = loadConfig();
  const cyclePrevKey = config.cyclePrev || "alt+shift+left";
  const cycleNextKey = config.cycleNext || "alt+shift+right";

  function updateStatus(
    ui: { setStatus: (id: string, text: string) => void },
    currentModelId?: string,
  ): void {
    try {
      const name = selectedName();
      ui.setStatus(
        "agent",
        formatStatusText(
          name,
          name ? profiles[name] : undefined,
          currentModelId,
          name === "orchestrator" ? sessionUlw : undefined,
        ),
      );
    } catch { /* never crash Pi */ }
  }

  /** Apply thinking from a profile whose model matches the current model. */
  function applyThinkingForCurrentModel(currentModelId?: string): void {
    if (!currentModelId) return;
    for (const entry of Object.values(profiles)) {
      if (entry?.thinking && entry.model === currentModelId) {
        try {
          pi.setThinkingLevel(entry.thinking as any);
        } catch { /* clamped or unsupported, ignore */ }
        return;
      }
    }
  }

  /**
   * Apply the selected cycle agent's model (and its thinking) to Pi.
   * Returns true on success/no-op.
   */
  async function applySelectedModel(ctx: any): Promise<boolean> {
    const name = selectedName();
    if (!name) {
      return false;
    }
    const profile = profiles[name];

    // Always apply thinking from the selected profile, if set
    if (profile?.thinking) {
      try {
        pi.setThinkingLevel(profile.thinking as any);
      } catch { /* clamped or unsupported by model, ignore */ }
    }

    // Apply model (only if it differs from the current one)
    const explicitModelId = profile?.model;
    if (!shouldApplyConfiguredModel(profile, ctx.model?.id)) {
      return true;
    }
    const modelId = explicitModelId!;
    const separator = modelId.indexOf("/");
    if (separator <= 0) {
      return false;
    }
    try {
      const provider = modelId.slice(0, separator);
      const modelName = modelId.slice(separator + 1);
      const model = ctx.modelRegistry.find(provider, modelName);
      if (!model) {
        return false;
      }
      return await pi.setModel(model);
    } catch {
      return false;
    }
  }

  try {
    pi.on("session_start", async (_event, ctx) => {
      pendingFileOperationNudge = false;
      reloadProfiles();
      initializeSessionUlw(ctx);
      // Apply thinking from whichever profile matches the current model
      applyThinkingForCurrentModel(ctx.model?.id);
      updateStatus(ctx.ui, ctx.model?.id);
      publishUlwState();
    });
  } catch { /* never crash Pi */ }

  try {
    pi.on("tool_result", async (event) => {
      const name = selectedName();
      if (
        shouldInjectMainSessionPrompt(name) &&
        FILE_OPERATION_TOOLS.has(event.toolName)
      ) {
        pendingFileOperationNudge = true;
      }
    });
  } catch { /* never crash Pi */ }

  try {
    pi.on("before_agent_start", async (event) => {
      const name = selectedName();
      if (!shouldInjectMainSessionPrompt(name)) {
        pendingFileOperationNudge = false;
        return undefined;
      }

      let systemPrompt = event.systemPrompt;
      const agentPrompt = loadAgentPromptBody(name!);
      if (agentPrompt) {
        systemPrompt = appendAgentPrompt(systemPrompt, agentPrompt);
      }
      systemPrompt = appendAgentPrompt(
        systemPrompt,
        sessionUlw ? ORCHESTRATOR_ULW_REMINDER : ORCHESTRATOR_WORKFLOW_REMINDER,
      );
      if (pendingFileOperationNudge) {
        systemPrompt = appendAgentPrompt(
          systemPrompt,
          sessionUlw ? FILE_OPERATION_ULW_NUDGE : FILE_OPERATION_NUDGE,
        );
        pendingFileOperationNudge = false;
      }
      return { systemPrompt };
    });
  } catch { /* never crash Pi */ }

  try {
    pi.on("model_select", async (event, ctx) => {
      // Apply thinking for the newly selected model
      applyThinkingForCurrentModel(`${event.model.provider}/${event.model.id}`);
      updateStatus(ctx.ui, `${event.model.provider}/${event.model.id}`);
    });
  } catch { /* never crash Pi */ }

  try {
    pi.registerShortcut(cyclePrevKey, {
      description: "Select and apply previous primary agent model",
      handler: async (ctx) => {
        try {
          if (cycle.length > 0) {
            currentIndex = (currentIndex - 1 + cycle.length) % cycle.length;
          }
          const ok = await applySelectedModel(ctx);
          const name = selectedName();
          if (!shouldInjectMainSessionPrompt(name)) pendingFileOperationNudge = false;
          updateStatus(
            ctx.ui,
            pickStatusModelId(name ? profiles[name] : undefined, ctx.model?.id, ok),
          );
          publishUlwState();
        } catch (error) {
          ctx.ui.notify(`Agent switch failed: ${error}`, "error");
        }
      },
    });
  } catch { /* duplicate/invalid key */ }

  try {
    pi.registerShortcut(cycleNextKey, {
      description: "Select and apply next primary agent model",
      handler: async (ctx) => {
        try {
          if (cycle.length > 0) {
            currentIndex = (currentIndex + 1) % cycle.length;
          }
          const ok = await applySelectedModel(ctx);
          const name = selectedName();
          if (!shouldInjectMainSessionPrompt(name)) pendingFileOperationNudge = false;
          updateStatus(
            ctx.ui,
            pickStatusModelId(name ? profiles[name] : undefined, ctx.model?.id, ok),
          );
          publishUlwState();
        } catch (error) {
          ctx.ui.notify(`Agent switch failed: ${error}`, "error");
        }
      },
    });
  } catch { /* duplicate/invalid key */ }

  try {
    pi.registerCommand("ulw", {
      description: "Toggle orchestrator ULW mode for this session",
      handler: async (_args, ctx) => {
        try {
          if (!shouldInjectMainSessionPrompt(selectedName())) {
            ctx.ui.notify("ULW only applies to orchestrator.", "info");
            return;
          }
          sessionUlw = !sessionUlw;
          updateStatus(ctx.ui, ctx.model?.id);
          publishUlwState();
          ctx.ui.notify(`ULW ${sessionUlw ? "ON" : "OFF"}`, "info");
        } catch (error) {
          ctx.ui.notify(`/ulw failed: ${error}`, "error");
        }
      },
    });
  } catch { /* command conflict */ }

  try {
    pi.registerCommand("agents-list", {
      description: "List switchable primary agent models",
      handler: async (_args, ctx) => {
        try {
          reloadProfiles();
          const lines = [
            "Primary agent model switcher",
            "",
            ...cycle.map((name, index) => {
              const marker = index === currentIndex ? "►" : " ";
              const model = resolveEffectiveModelId(profiles[name], ctx.model?.id) ?? "(no model)";
              const p = profiles[name];
              const icon = p?.icon ?? DEFAULT_AGENT_ICON;
              const colored = p?.color
                ? `${hexToAnsiFg(p.color)}${name}${ANSI_RESET}`
                : name;
              return `${marker} ${icon} ${colored} · ${model}`;
            }),
            "",
            `${formatKey(cyclePrevKey)} / ${formatKey(cycleNextKey)}  select and apply`,
            "",
            "All settings.agent entries are listed unless mode is set to subagent.",
          ];
          ctx.ui.notify(lines.join("\n"), "info");
        } catch (error) {
          ctx.ui.notify(`/agents failed: ${error}`, "error");
        }
      },
    });
  } catch { /* command conflict */ }
}