import type {
  ExtensionAPI,
  ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Text, type AutocompleteItem } from "@earendil-works/pi-tui";
import {
  discoverRoles,
  findRole,
  normalizeRoleKey,
  type RoleDefinition,
  type RoleDiscoveryResult,
  type RoleSource,
} from "./roles.js";

const ROLE_FLAG = "role";
const ROLE_ENTRY_TYPE = "pi-role:state";
const ROLE_REMOVED_MESSAGE_TYPE = "pi-role:removed";
const STATUS_KEY = "pi-role";
const SUBCOMMANDS = ["list", "current", "reload", "none"] as const;

export const ROLE_REMOVED_AGENT_MESSAGE =
  "Your previously assigned role has been removed. Stop following that role’s instructions and follow the current system prompt and user instructions.";

export interface RoleStateEntry {
  version: 1;
  active: boolean;
  action: "activate" | "reload" | "disable";
  message: string;
  key?: string;
  name?: string;
  source?: RoleSource;
  path?: string;
}

interface RuntimeState {
  activeRole: RoleDefinition | null;
  discovery: RoleDiscoveryResult;
  ctx: ExtensionContext | null;
}

function emptyDiscovery(): RoleDiscoveryResult {
  return { roles: [], warnings: [], userRolesDir: "" };
}

export function roleStatusText(name: string): string {
  return `Role: ${name} 👤 `;
}

export function composeRoleSystemPrompt(
  upstream: string,
  role: RoleDefinition | null,
): string | undefined {
  if (!role) return undefined;
  const body = role.body.trim();
  if (!body) return upstream || undefined;
  return [upstream, body].filter((part) => part.length > 0).join("\n\n");
}

export function findLatestRoleState(
  entries: readonly unknown[],
): RoleStateEntry | undefined {
  for (let index = entries.length - 1; index >= 0; index--) {
    const entry = entries[index];
    if (!entry || typeof entry !== "object") continue;
    const candidate = entry as {
      type?: unknown;
      customType?: unknown;
      data?: unknown;
    };
    if (
      candidate.type !== "custom" ||
      candidate.customType !== ROLE_ENTRY_TYPE
    ) {
      continue;
    }
    if (!candidate.data || typeof candidate.data !== "object") continue;
    const data = candidate.data as Partial<RoleStateEntry>;
    if (
      data.version !== 1 ||
      typeof data.active !== "boolean" ||
      (data.action !== "activate" &&
        data.action !== "reload" &&
        data.action !== "disable") ||
      typeof data.message !== "string"
    ) {
      continue;
    }
    return data as RoleStateEntry;
  }
  return undefined;
}

export function roleCompletions(
  prefix: string,
  roles: readonly RoleDefinition[],
): AutocompleteItem[] | null {
  const needle = normalizeRoleKey(prefix.trimStart());
  const items: AutocompleteItem[] = [];

  for (const command of SUBCOMMANDS) {
    if (command.startsWith(needle)) {
      items.push({
        value: command,
        label: command,
        description:
          command === "none"
            ? "Disable the active role"
            : command === "list"
              ? "List available roles"
              : command === "current"
                ? "Show the active role"
                : "Reload the active role file",
      });
    }
  }

  for (const role of roles) {
    if (normalizeRoleKey(role.key).startsWith(needle)) {
      items.push({
        value: role.key,
        label: role.key,
        description: `${role.name} [${role.source}] — ${role.description}`,
      });
    }
  }

  return items.length > 0 ? items : null;
}

function report(
  ctx: ExtensionContext,
  message: string,
  type: "info" | "warning" | "error",
): void {
  if (ctx.hasUI) {
    ctx.ui.notify(message, type);
  } else if (type !== "info") {
    console.error(message);
  }
}

function roleStateFor(
  role: RoleDefinition,
  action: "activate" | "reload",
): RoleStateEntry {
  return {
    version: 1,
    active: true,
    action,
    message:
      action === "reload"
        ? `Role type reloaded: ${role.name}.`
        : `Role type switched to ${role.name}.`,
    key: role.key,
    name: role.name,
    source: role.source,
    path: role.path,
  };
}

export default function piRoleExtension(pi: ExtensionAPI): void {
  const state: RuntimeState = {
    activeRole: null,
    discovery: emptyDiscovery(),
    ctx: null,
  };

  const clearRole = (ctx: ExtensionContext): void => {
    state.activeRole = null;
    if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
  };

  const refresh = (ctx: ExtensionContext, showWarnings: boolean): void => {
    state.discovery = discoverRoles(ctx.cwd, {
      projectTrusted: ctx.isProjectTrusted(),
    });
    if (showWarnings) {
      for (const warning of state.discovery.warnings) {
        report(ctx, `pi-role: ${warning}`, "warning");
      }
    }
  };

  const setRoleWithoutEntry = (
    role: RoleDefinition,
    ctx: ExtensionContext,
  ): void => {
    state.activeRole = role;
    if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, roleStatusText(role.name));
  };

  const applyExplicitRole = (
    role: RoleDefinition,
    action: "activate" | "reload",
    ctx: ExtensionContext,
  ): void => {
    setRoleWithoutEntry(role, ctx);
    pi.appendEntry<RoleStateEntry>(ROLE_ENTRY_TYPE, roleStateFor(role, action));
  };

  const restoreRole = (ctx: ExtensionContext): void => {
    const persisted = findLatestRoleState(ctx.sessionManager.getBranch());
    if (!persisted || !persisted.active || !persisted.key) {
      clearRole(ctx);
      return;
    }

    const role = findRole(state.discovery.roles, persisted.key);
    if (!role) {
      clearRole(ctx);
      report(
        ctx,
        `pi-role: cannot restore role "${persisted.key}" because its role file is unavailable or invalid.`,
        "warning",
      );
      return;
    }

    setRoleWithoutEntry(role, ctx);
  };

  const listRoles = (ctx: ExtensionContext): void => {
    if (state.discovery.roles.length === 0) {
      report(ctx, "pi-role: no roles found.", "info");
      return;
    }

    const lines = state.discovery.roles.map((role) => {
      const active =
        state.activeRole &&
        normalizeRoleKey(state.activeRole.key) === normalizeRoleKey(role.key)
          ? "* "
          : "  ";
      return `${active}${role.key} [${role.source}] — ${role.name}: ${role.description}`;
    });
    report(ctx, lines.join("\n"), "info");
  };

  pi.registerFlag(ROLE_FLAG, {
    type: "string",
    description:
      "Start with a configured role (for example, --role architect).",
  });

  pi.registerEntryRenderer<RoleStateEntry>(
    ROLE_ENTRY_TYPE,
    (entry, _options, theme) => {
      const message = entry.data?.message ?? "Role state changed.";
      return new Text(theme.fg("dim", message), 1, 0);
    },
  );

  pi.on("session_start", (event, ctx) => {
    state.ctx = ctx;
    refresh(ctx, true);

    if (event.reason === "new") {
      clearRole(ctx);
      return;
    }

    if (event.reason === "startup") {
      const flag = pi.getFlag(ROLE_FLAG);
      if (typeof flag === "string" && flag.trim().length > 0) {
        if (normalizeRoleKey(flag) === "none") {
          clearRole(ctx);
          return;
        }
        const role = findRole(state.discovery.roles, flag);
        if (!role) {
          clearRole(ctx);
          report(ctx, `pi-role: unknown role "${flag}".`, "error");
          return;
        }
        applyExplicitRole(role, "activate", ctx);
        return;
      }
    }

    restoreRole(ctx);
  });

  pi.on("before_agent_start", (event) => {
    const systemPrompt = composeRoleSystemPrompt(
      event.systemPrompt,
      state.activeRole,
    );
    return systemPrompt === undefined ? undefined : { systemPrompt };
  });

  pi.registerCommand("role", {
    description: "List, activate, reload, inspect, or disable session roles",
    getArgumentCompletions: (prefix) => {
      if (state.ctx) refresh(state.ctx, false);
      return roleCompletions(prefix, state.discovery.roles);
    },
    handler: async (args, ctx) => {
      state.ctx = ctx;
      refresh(ctx, true);
      const requested = args.trim();
      const command = normalizeRoleKey(requested);

      if (!requested || command === "list") {
        listRoles(ctx);
        return;
      }

      if (command === "current") {
        if (!state.activeRole) {
          report(ctx, "pi-role: no role is active.", "info");
          return;
        }
        const role = state.activeRole;
        report(
          ctx,
          `Role: ${role.name}\nKey: ${role.key}\nSource: ${role.source}\n${role.description}\n${role.path}`,
          "info",
        );
        return;
      }

      if (command === "none") {
        if (!state.activeRole) {
          report(ctx, "pi-role: no role is active.", "info");
          return;
        }
        clearRole(ctx);
        pi.appendEntry<RoleStateEntry>(ROLE_ENTRY_TYPE, {
          version: 1,
          active: false,
          action: "disable",
          message: "Role type disabled.",
        });
        pi.sendMessage({
          customType: ROLE_REMOVED_MESSAGE_TYPE,
          content: ROLE_REMOVED_AGENT_MESSAGE,
          display: false,
        });
        return;
      }

      if (command === "reload") {
        if (!state.activeRole) {
          report(ctx, "pi-role: no active role to reload.", "error");
          return;
        }
        const role = findRole(state.discovery.roles, state.activeRole.key);
        if (!role) {
          report(
            ctx,
            `pi-role: cannot reload role "${state.activeRole.key}" because its role file is unavailable or invalid.`,
            "error",
          );
          return;
        }
        applyExplicitRole(role, "reload", ctx);
        return;
      }

      const role = findRole(state.discovery.roles, requested);
      if (!role) {
        report(ctx, `pi-role: unknown role "${requested}".`, "error");
        return;
      }

      const action =
        state.activeRole &&
        normalizeRoleKey(state.activeRole.key) === normalizeRoleKey(role.key)
          ? "reload"
          : "activate";
      applyExplicitRole(role, action, ctx);
    },
  });
}