import type {
  ExtensionAPI,
  ExtensionContext,
  Skill,
  SlashCommandInfo,
  ToolInfo,
} from "@earendil-works/pi-coding-agent";
import {
  formatSkillsForPrompt,
  getSettingsListTheme,
} from "@earendil-works/pi-coding-agent";
import {
  Container,
  type SettingItem,
  SettingsList,
  Text,
} from "@earendil-works/pi-tui";

export type ContextState = {
  disabledSkillNames: string[];
  disabledToolNames: string[];
};

type StateEntry = {
  type: string;
  customType?: string;
  data?: unknown;
};

const STATE_TYPE = "context-control-state";
const STATUS_KEY = "context-control";
const DESCRIPTION_MAX_LENGTH = 400;

export default function contextControl(pi: ExtensionAPI) {
  let disabledSkillNames = new Set<string>();
  let disabledToolNames = new Set<string>();

  function persist() {
    pi.appendEntry<ContextState>(STATE_TYPE, {
      disabledSkillNames: [...disabledSkillNames].sort(),
      disabledToolNames: [...disabledToolNames].sort(),
    });
  }

  function restore(ctx: ExtensionContext) {
    const saved = readStateFromBranch(ctx.sessionManager.getBranch());
    disabledSkillNames = new Set(saved.disabledSkillNames);
    disabledToolNames = new Set(saved.disabledToolNames);
    applyTools();
    renderStatus(ctx);
  }

  function applyTools() {
    const active = pi
      .getActiveTools()
      .filter((name) => !disabledToolNames.has(name));
    pi.setActiveTools(active);
  }

  function renderStatus(ctx: ExtensionContext) {
    const count = disabledSkillNames.size + disabledToolNames.size;
    ctx.ui.setStatus(
      STATUS_KEY,
      count ? ctx.ui.theme.fg("warning", `context -${count}`) : undefined,
    );
  }

  function getExtensionTools() {
    return pi
      .getAllTools()
      .filter(
        (tool) =>
          tool.sourceInfo.source !== "builtin" &&
          tool.sourceInfo.source !== "sdk",
      );
  }

  function getCommandOnlyExtensions(toolNames: Set<string>) {
    return pi
      .getCommands()
      .filter((command) => command.source === "extension")
      .filter((command) => !toolNames.has(command.name));
  }

  pi.registerCommand("context", {
    description: "Toggle skills and extension tools for this session",
    handler: async (_args, ctx) => {
      if (ctx.mode !== "tui") {
        ctx.ui.notify("/context requires TUI mode", "error");
        return;
      }

      const skills = ctx.getSystemPromptOptions().skills ?? [];
      const extensionTools = getExtensionTools();
      const toolNames = new Set(extensionTools.map((tool) => tool.name));
      const commandOnlyExtensions = getCommandOnlyExtensions(toolNames);
      const items = buildSettingItems(
        skills,
        extensionTools,
        commandOnlyExtensions,
        pi.getActiveTools(),
        {
          disabledSkillNames: [...disabledSkillNames],
          disabledToolNames: [...disabledToolNames],
        },
      );

      await ctx.ui.custom((tui, theme, _kb, done) => {
        const container = new Container();
        container.addChild(
          new Text(theme.fg("accent", theme.bold("Context Control")), 1, 0),
        );

        const settingsList = new SettingsList(
          items,
          Math.min(items.length + 2, 18),
          getSettingsListTheme(),
          (id, value) => {
            const [kind, name] = id.split(":", 2);
            if (!name) return;

            const disabled = value === "disabled";
            if (kind === "skill") setFlag(disabledSkillNames, name, disabled);
            if (kind === "tool") {
              setFlag(disabledToolNames, name, disabled);
              pi.setActiveTools(
                applyToolToggle(pi.getActiveTools(), name, disabled),
              );
            }

            persist();
            renderStatus(ctx);
          },
          () => done(undefined),
          { enableSearch: true },
        );

        container.addChild(settingsList);
        return {
          render: (width: number) => container.render(width),
          invalidate: () => container.invalidate(),
          handleInput: (data: string) => {
            settingsList.handleInput?.(data);
            tui.requestRender();
          },
        };
      });
    },
  });

  pi.registerCommand("context-status", {
    description: "Show disabled skills and extension tools",
    handler: async (_args, ctx) => {
      const skills = [...disabledSkillNames].sort();
      const tools = [...disabledToolNames].sort();
      ctx.ui.notify(
        `Disabled skills: ${skills.join(", ") || "none"}\nDisabled extension tools: ${tools.join(", ") || "none"}`,
        "info",
      );
    },
  });

  pi.on("session_start", async (_event, ctx) => restore(ctx));
  pi.on("session_tree", async (_event, ctx) => restore(ctx));
  pi.on("session_shutdown", async (_event, ctx) =>
    ctx.ui.setStatus(STATUS_KEY, undefined),
  );

  pi.on("before_agent_start", async (event) => {
    if (disabledSkillNames.size === 0) return;

    const skills = (event.systemPromptOptions.skills ?? []).filter(
      (skill: Skill) => !disabledSkillNames.has(skill.name),
    );
    const systemPrompt = replaceSkillsSection(event.systemPrompt, skills);
    return { systemPrompt };
  });
}

function setFlag(set: Set<string>, value: string, enabled: boolean) {
  if (enabled) set.add(value);
  else set.delete(value);
}

export function readStateFromBranch(
  entries: Iterable<StateEntry>,
): ContextState {
  let saved: ContextState | undefined;
  for (const entry of entries) {
    if (entry.type !== "custom" || entry.customType !== STATE_TYPE) continue;
    if (!isContextState(entry.data)) continue;
    saved = entry.data;
  }

  return saved ?? { disabledSkillNames: [], disabledToolNames: [] };
}

const truncateText = (description: string, maxLength: number) => {
  if (description.length <= maxLength) return description;
  return description.slice(0, maxLength - 4) + " ...";
};

export function buildSettingItems(
  skills: Skill[],
  extensionTools: ToolInfo[],
  commandOnlyExtensions: SlashCommandInfo[],
  activeToolNames: Iterable<string>,
  state: ContextState,
): SettingItem[] {
  const activeTools = new Set(activeToolNames);
  const disabledSkills = new Set(state.disabledSkillNames);
  const disabledTools = new Set(state.disabledToolNames);

  return [
    ...skills.map((skill) => ({
      id: `skill:${skill.name}`,
      label: `skill ${skill.name}`,
      description: `${truncateText(skill.description, DESCRIPTION_MAX_LENGTH)}\n${skill.filePath}`,
      currentValue: disabledSkills.has(skill.name) ? "disabled" : "enabled",
      values: ["enabled", "disabled"],
    })),
    ...extensionTools.map((tool) => ({
      id: `tool:${tool.name}`,
      label: `tool ${tool.name}`,
      description: `${truncateText(tool.description, DESCRIPTION_MAX_LENGTH)}\n${tool.sourceInfo.path}`,
      currentValue:
        disabledTools.has(tool.name) || !activeTools.has(tool.name)
          ? "disabled"
          : "enabled",
      values: ["enabled", "disabled"],
    })),
    ...commandOnlyExtensions.map((command) => ({
      id: `command:${command.name}`,
      label: `cmd /${command.name}`,
      description: `Commands cannot be disabled session-only. Use pi config for full resource filtering.\n${command.sourceInfo.path}`,
      currentValue: "loaded",
    })),
  ];
}

export function applyToolToggle(
  activeToolNames: Iterable<string>,
  toolName: string,
  disabled: boolean,
): string[] {
  const activeTools = new Set(activeToolNames);
  if (disabled) activeTools.delete(toolName);
  else activeTools.add(toolName);
  return [...activeTools];
}

function isContextState(value: unknown): value is ContextState {
  if (!value || typeof value !== "object") return false;
  const state = value as Partial<ContextState>;
  return (
    Array.isArray(state.disabledSkillNames) &&
    Array.isArray(state.disabledToolNames)
  );
}

export function replaceSkillsSection(systemPrompt: string, skills: Skill[]) {
  const skillBlock =
    /\n\nThe following skills provide specialized instructions for specific tasks\.[\s\S]*?<\/available_skills>/;
  if (!skillBlock.test(systemPrompt)) return systemPrompt;
  return systemPrompt.replace(skillBlock, formatSkillsForPrompt(skills));
}