import React, { useCallback } from "react";
import { create } from "zustand";
import { useShallow } from "zustand/react/shallow";
import { useAppStore, useModel } from "./state.ts";
import { useSession } from "./session-context.ts";
import { listPreviousSessions, type Session } from "./session-history/index.ts";
import { Auth, mergeEnvVar, useConfig, useSetConfig, Config } from "./config.ts";
import { ModelSetup } from "./components/auto-detect-models.tsx";
import { AutofixModelMenu } from "./components/autofix-model-menu.tsx";
import { ConfirmDialog } from "./components/confirm-dialog.tsx";
import { readAuthForModel } from "./config.ts";
import { keyFromName, SYNTHETIC_PROVIDER, providerForBaseUrl } from "./providers.ts";
import { KbShortcutPanel } from "./components/kb-select/kb-shortcut-panel.tsx";
import { Item, ShortcutArray, Keymap } from "./components/kb-select/kb-shortcut-select.tsx";
import { MenuQuotaIndicator } from "./components/menu-quota-indicator.tsx";
import { CustomAuthFlow } from "./components/add-model-flow.tsx";
import { Span, useApp } from "paintcannon-react";
import { useKeyboard } from "./hooks/use-keyboard.ts";
import { LoadSessionMenu } from "./session-history/load-session-menu.tsx";
type MenuMode =
  | "main-menu"
  | "load-session"
  | "settings-menu"
  | "model-select"
  | "add-model"
  | "diff-apply-toggle"
  | "fix-json-toggle"
  | "set-default-model"
  | "quit-confirm"
  | "remove-model"
  | "clear-confirm"
  | "notifications-menu";
type MenuState = {
  menuMode: MenuMode;
  setMenuMode: (mode: MenuMode) => void;
};
const useMenuState = create<MenuState>((set, _) => ({
  menuMode: "main-menu",
  setMenuMode: menuMode => {
    set({
      menuMode,
    });
  },
}));
export function Menu({ onSessionChange }: { onSessionChange: (session: Session) => void }) {
  const { menuMode, setMenuMode } = useMenuState(
    useShallow(state => ({
      menuMode: state.menuMode,
      setMenuMode: state.setMenuMode,
    })),
  );
  if (menuMode === "main-menu") return <MainMenu />;
  if (menuMode === "load-session") {
    return (
      <LoadSessionMenu onBack={() => setMenuMode("main-menu")} onSessionChange={onSessionChange} />
    );
  }
  if (menuMode === "settings-menu") return <SettingsMenu />;
  if (menuMode === "model-select") return <SwitchModelMenu />;
  if (menuMode === "set-default-model") return <SetDefaultModelMenu />;
  if (menuMode === "quit-confirm") return <QuitConfirm />;
  if (menuMode === "clear-confirm")
    return <ClearConversationConfirm onSessionChange={onSessionChange} />;
  if (menuMode === "remove-model") return <RemoveModelMenu />;
  if (menuMode === "diff-apply-toggle") return <DiffApplyToggle />;
  if (menuMode === "fix-json-toggle") return <FixJsonToggle />;
  if (menuMode === "notifications-menu") return <NotificationsMenu />;
  const _: "add-model" = menuMode;
  return <AddModelMenuFlow />;
}
function AutofixToggle({
  configKey,
  modelNickname,
  disableNotification,
  enableNotification,
  defaultModel,
  children,
}: {
  disableNotification: string;
  enableNotification: string;
  defaultModel: string;
  modelNickname: string;
  configKey: "diffApply" | "fixJson";
  children: React.ReactNode;
}) {
  const config = useConfig();
  const setConfig = useSetConfig();
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const { toggleMenu, notify } = useAppStore(
    useShallow(state => ({
      toggleMenu: state.toggleMenu,
      notify: state.notify,
    })),
  );
  const session = useSession();
  useKeyboard(event => {
    if (event.key === "Escape") setMenuMode("main-menu");
  });
  if (config[configKey]) {
    return (
      <ConfirmDialog
        rejectLabel={`Disable ${modelNickname}`}
        confirmLabel={`Keep ${modelNickname} on (recommended)`}
        onReject={async () => {
          const newconf = {
            ...config,
          };
          delete newconf[configKey];
          await setConfig(newconf);
          setMenuMode("main-menu");
          toggleMenu();
          notify(disableNotification, session);
        }}
        onConfirm={() => {
          setMenuMode("main-menu");
        }}
      />
    );
  }
  return (
    <AutofixModelMenu
      defaultModel={defaultModel}
      modelNickname={modelNickname}
      config={config}
      onOverrideDefaultApiKey={async apiEnvVar => {
        await setConfig({
          ...config,
          defaultApiKeyOverrides: {
            ...(config.defaultApiKeyOverrides || {}),
            [keyFromName(SYNTHETIC_PROVIDER.name)]: apiEnvVar,
          },
        });
      }}
      onComplete={async setting => {
        await setConfig({
          ...config,
          [configKey]: setting,
        });
        setMenuMode("main-menu");
        toggleMenu();
        notify(enableNotification, session);
      }}
      onCancel={() => {
        setMenuMode("main-menu");
      }}
    >
      {children}
    </AutofixModelMenu>
  );
}
function DiffApplyToggle() {
  return (
    <AutofixToggle
      defaultModel="hf:syntheticlab/diff-apply"
      configKey="diffApply"
      modelNickname="diff-apply"
      enableNotification="Fast diff apply enabled"
      disableNotification="Fast diff apply disabled"
    >
      <Span>
        Even good coding models sometimes make minor mistakes generating code diffs, which can cause
        slow retries and can confuse them, since models often aren't trained as well to handle edit
        failures as they are successes. Diff-apply is a fast, small model that fixes minor code diff
        edit inaccuracies. It speeds up iteration and can significantly improve model performance.
      </Span>
    </AutofixToggle>
  );
}
function FixJsonToggle() {
  return (
    <AutofixToggle
      defaultModel="hf:syntheticlab/fix-json"
      configKey="fixJson"
      modelNickname="fix-json"
      enableNotification="JSON auto-fix enabled"
      disableNotification="JSON auto-fix disabled"
    >
      <Span>
        Octo uses tools to work with your underlying codebase. Some model providers don't support
        strict constraints on how tool calls are generated, and models can make mistakes generating
        JSON, the format used for all of Octo's tool calls.
      </Span>
      <Span>
        The fix-json model can automatically fix broken JSON for Octo, helping models avoid failures
        more quickly and cheaply than retrying the main model. It also may help reduce the main
        model's confusion.
      </Span>
    </AutofixToggle>
  );
}
function SwitchModelMenu() {
  const { setModelOverride, toggleMenu } = useAppStore(
    useShallow(state => ({
      setModelOverride: state.setModelOverride,
      toggleMenu: state.toggleMenu,
    })),
  );
  const session = useSession();
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const config = useConfig();
  const setConfig = useSetConfig();
  const [pendingModel, setPendingModel] = React.useState<null | Config["models"][number]>(null);
  useKeyboard(event => {
    if (event.key === "Escape" && pendingModel == null) setMenuMode("main-menu");
  });
  const onSelect = useCallback(
    async (item: Item<`model-${string}` | "back">) => {
      if (item.value === "back") {
        setMenuMode("main-menu");
        return;
      }
      const target = item.value.replace("model-", "");
      const model = config.models.find(m => m.nickname === target)!;
      const auth = await readAuthForModel(model, config);
      if (!auth.ok) {
        setPendingModel(model);
        return;
      }
      setModelOverride(target, session);
      setMenuMode("main-menu");
      toggleMenu();
    },
    [config, setMenuMode, setModelOverride, toggleMenu, session],
  );
  if (pendingModel) {
    return (
      <CustomAuthFlow
        config={config}
        authData={
          pendingModel.type === "codex"
            ? {
                modelType: "codex",
              }
            : {
                modelType: pendingModel.type,
                baseUrl: pendingModel.baseUrl,
              }
        }
        onComplete={async (auth?: Auth) => {
          const index = config.models.indexOf(pendingModel);
          if (auth && index >= 0) {
            if (pendingModel.type === "codex") {
              if (auth.type === "codex") {
                const models = [...config.models];
                models[index] = {
                  ...pendingModel,
                  auth,
                };
                await setConfig({
                  ...config,
                  models,
                });
              }
            } else if (auth.type === "env") {
              await setConfig(mergeEnvVar(config, pendingModel, auth.name));
            } else if (auth.type === "command") {
              const models = [...config.models];
              models[index] = {
                ...pendingModel,
                auth,
              };
              await setConfig({
                ...config,
                models,
              });
            }
          }
          setModelOverride(pendingModel.nickname, session);
          setPendingModel(null);
          setMenuMode("main-menu");
          toggleMenu();
        }}
        onCancel={() => {
          setPendingModel(null);
        }}
      />
    );
  }
  const numericItems = config.models.map(model => {
    return {
      label: model.nickname,
      value: `model-${model.nickname}` as const,
    };
  });
  const shortcutItems: ShortcutArray<`model-${string}` | "back"> = [
    {
      type: "auto-list" as const,
      order: numericItems,
    },
    {
      type: "key" as const,
      mapping: {
        b: {
          label: "Back to main menu",
          value: "back",
        },
      },
    },
  ];
  return (
    <KbShortcutPanel
      title="Which model should Octo use now?"
      shortcutItems={shortcutItems}
      onSelect={onSelect}
    />
  );
}
type SettingsValues =
  | "set-default-model"
  | "remove-model"
  | "disable-diff-apply"
  | "disable-fix-json";
const SETTINGS_ITEMS = {
  c: {
    label: "Change the default model",
    value: "set-default-model",
  },
  r: {
    label: "Remove a model",
    value: "remove-model",
  },
  d: {
    label: "Disable fast diff application",
    value: "disable-diff-apply",
  },
  t: {
    label: "Disable auto-fixing JSON tool calls",
    value: "disable-fix-json",
  },
} satisfies Keymap<SettingsValues>;
function filterSettings(config: Config) {
  let items: Keymap<SettingsValues> = {};
  if (config.models.length > 1) {
    items = {
      ...items,
      c: SETTINGS_ITEMS.c,
      r: SETTINGS_ITEMS.r,
    };
  }
  if (config.diffApply) {
    items = {
      ...items,
      d: SETTINGS_ITEMS.d,
    };
  }
  if (config.fixJson) {
    items = {
      ...items,
      t: SETTINGS_ITEMS.t,
    };
  }
  return items;
}
function MainMenu() {
  const { toggleMenu, notify, resetPreMenuVimMode } = useAppStore(
    useShallow(state => ({
      toggleMenu: state.toggleMenu,
      notify: state.notify,
      resetPreMenuVimMode: state.resetPreMenuVimMode,
    })),
  );
  const session = useSession();
  const hasPreviousSessions =
    listPreviousSessions(session.metadata.cwd, session.metadata.sessionId).length > 0;
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const config = useConfig();
  const setConfig = useSetConfig();
  const model = useModel();
  const provider = model.type === "codex" ? null : providerForBaseUrl(model.baseUrl);
  const isSynthetic = provider === SYNTHETIC_PROVIDER;
  useKeyboard(event => {
    if (event.key === "Escape") toggleMenu();
  });
  type Value =
    | "model-select"
    | "add-model"
    | "load-session"
    | "vim-toggle"
    | "return"
    | "quit"
    | "fix-json-toggle"
    | "diff-apply-toggle"
    | "settings-menu"
    | "clear-confirm"
    | "notifications-menu";
  let items: Keymap<Value> = {
    m: {
      label: "⤭ Switch model",
      value: "model-select" as const,
    },
    a: {
      label: "+ Add a new model",
      value: "add-model" as const,
    },
    c: {
      label: "⦿ New conversation",
      value: "clear-confirm" as const,
    },
  };
  if (hasPreviousSessions) {
    items = {
      ...items,
      s: {
        label: "⥻ Load previous session",
        value: "load-session" as const,
      },
    };
  }
  if (config.vimEmulation?.enabled) {
    items = {
      ...items,
      e: {
        label: "♺ Switch to Emacs mode",
        value: "vim-toggle" as const,
      },
    };
  } else {
    items = {
      ...items,
      v: {
        label: "♺ Switch to Vim mode",
        value: "vim-toggle" as const,
      },
    };
  }
  if (config.fixJson == null) {
    items = {
      ...items,
      f: {
        label: "🪄 Enable auto-fixing JSON tool calls",
        value: "fix-json-toggle" as const,
      },
    };
  }
  if (config.diffApply == null) {
    items = {
      ...items,
      d: {
        label: "💫 Enable fast diff application",
        value: "diff-apply-toggle" as const,
      },
    };
  }
  if (config.notifications?.notifyCommand) {
    items = {
      ...items,
      n: {
        label: "🕭 Notifications",
        value: "notifications-menu" as const,
      },
    };
  }
  const settings = filterSettings(config);
  if (Object.values(settings).length !== 0) {
    items = {
      ...items,
      t: {
        label: "⛭ Settings",
        value: "settings-menu" as const,
      },
    };
  }
  items = {
    ...items,
    b: {
      label: "⟵ Back to Octo",
      value: "return" as const,
    },
    q: {
      label: "× Quit",
      value: "quit" as const,
    },
  };
  const onSelect = useCallback(
    async (item: Item<Value>) => {
      if (item.value === "return") toggleMenu();
      else if (item.value === "quit") setMenuMode("quit-confirm");
      else if (item.value === "vim-toggle") {
        const wasEnabled = config.vimEmulation?.["enabled"] ?? false;

        // Write ONLY to config - single source of truth
        await setConfig({
          ...config,
          vimEmulation: {
            enabled: !wasEnabled,
          },
        });

        // When switching from Emacs to Vim, default to INSERT mode
        if (!wasEnabled) {
          resetPreMenuVimMode();
        }

        // Notify user
        notify(`Switched to ${wasEnabled ? "Emacs" : "Vim"} mode`, session);
        return;
      } else if (item.value === "clear-confirm") setMenuMode("clear-confirm");
      else setMenuMode(item.value);
    },
    [config, setConfig, notify, session],
  );
  return (
    <KbShortcutPanel
      title="Main Menu"
      shortcutItems={[
        {
          type: "key" as const,
          mapping: items,
        },
      ]}
      onSelect={onSelect}
    >
      {isSynthetic ? <MenuQuotaIndicator /> : null}
    </KbShortcutPanel>
  );
}
function SettingsMenu() {
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const config = useConfig();
  useKeyboard(event => {
    if (event.key === "Escape") setMenuMode("main-menu");
  });
  const settingsItems = filterSettings(config);
  let items: Keymap<SettingsValues | "back"> = {
    ...settingsItems,
    b: {
      label: "Back",
      value: "back" as const,
    },
  };
  const onSelect = useCallback((item: Item<SettingsValues | "back">) => {
    if (item.value === "disable-diff-apply") setMenuMode("diff-apply-toggle");
    else if (item.value === "disable-fix-json") setMenuMode("fix-json-toggle");
    else if (item.value === "back") setMenuMode("main-menu");
    else setMenuMode(item.value);
  }, []);
  return (
    <KbShortcutPanel
      title="Settings Menu"
      shortcutItems={[
        {
          type: "key" as const,
          mapping: items,
        },
      ]}
      onSelect={onSelect}
    />
  );
}
type NotificationValue = "always-notify" | "session-notify" | "notify-once" | "back";
function NotificationsMenu() {
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const config = useConfig();
  const setConfig = useSetConfig();
  const { sessionAutoNotify, notifyOnce, toggleMenu, setNotifyOnce, setNotifySession } =
    useAppStore(
      useShallow(state => ({
        sessionAutoNotify: state.sessionAutoNotify,
        notifyOnce: state.notifyOnce,
        toggleMenu: state.toggleMenu,
        setNotifyOnce: state.setNotifyOnce,
        setNotifySession: state.setNotifySession,
      })),
    );
  useKeyboard(event => {
    if (event.key === "Escape") setMenuMode("main-menu");
  });
  const alwaysNotify = config.notifications?.alwaysNotify;
  const items: Keymap<NotificationValue> = {
    o: {
      label: notifyOnce
        ? "Do not notify the next time Octo needs input"
        : "Notify the next time Octo needs input",
      value: "notify-once" as const,
    },
    s: {
      label: sessionAutoNotify
        ? "Stop auto-notifying this session"
        : "Auto-notify for the rest of this session",
      value: "session-notify" as const,
    },
    a: {
      label: alwaysNotify ? "Stop always auto-notifying" : "Always auto-notify",
      value: "always-notify" as const,
    },
    b: {
      label: "Back",
      value: "back" as const,
    },
  };
  const onSelect = useCallback(
    async (item: Item<NotificationValue>) => {
      if (item.value === "always-notify") {
        await setConfig({
          ...config,
          notifications: {
            ...config.notifications,
            alwaysNotify: !alwaysNotify,
          } as Config["notifications"],
        });
      } else if (item.value === "session-notify") {
        setNotifySession(!sessionAutoNotify);
      } else if (item.value === "notify-once") {
        setNotifyOnce(!notifyOnce);
        setMenuMode("main-menu");
        toggleMenu();
      } else if (item.value === "back") {
        setMenuMode("main-menu");
      }
    },
    [config, setConfig, alwaysNotify, sessionAutoNotify, notifyOnce, toggleMenu],
  );
  return (
    <KbShortcutPanel
      title="Notifications"
      shortcutItems={[
        {
          type: "key" as const,
          mapping: items,
        },
      ]}
      onSelect={onSelect}
    />
  );
}
function QuitConfirm() {
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const app = useApp();
  const session = useSession();
  return (
    <ConfirmDialog
      confirmLabel="Yes, quit"
      rejectLabel="Never mind, take me back"
      onConfirm={() => {
        /*
         * Restore the stashed pre-menu state (which may be an in-flight tool batch), then
         * record skip markers for any un-run tool calls so the session history stays
         * well-formed after exit.
         */
        const state = useAppStore.getState();
        state.closeMenu();
        state.abortResponse(session, { exiting: true });
        app.exit();
      }}
      onReject={() => setMenuMode("main-menu")}
      rejectFirst={true}
    />
  );
}
function ClearConversationConfirm({
  onSessionChange,
}: {
  onSessionChange: (session: Session) => void;
}) {
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const { startNewSession, notify } = useAppStore(
    useShallow(state => ({
      startNewSession: state.startNewSession,
      notify: state.notify,
    })),
  );
  const session = useSession();
  return (
    <ConfirmDialog
      confirmLabel="Yes, start new conversation"
      rejectLabel="Never mind, take me back"
      onConfirm={async () => {
        const { cwd, cliArgs } = session.metadata;
        const oldSessionId = session.metadata.sessionId;
        const newSession = startNewSession(cwd, cliArgs);
        onSessionChange(newSession);
        setMenuMode("main-menu");
        const resumeHint =
          oldSessionId != null
            ? ` To resume the previous session, run \`octo --resume ${oldSessionId}\``
            : "";
        notify(`New conversation started.${resumeHint}`, newSession);
      }}
      onReject={() => setMenuMode("main-menu")}
    />
  );
}
function SetDefaultModelMenu() {
  const { setModelOverride, toggleMenu } = useAppStore(
    useShallow(state => ({
      setModelOverride: state.setModelOverride,
      toggleMenu: state.toggleMenu,
    })),
  );
  const session = useSession();
  const config = useConfig();
  const setConfig = useSetConfig();
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  useKeyboard(event => {
    if (event.key === "Escape") setMenuMode("main-menu");
  });
  const numericItems = config.models.map(model => {
    return {
      label: model.nickname,
      value: `model-${model.nickname}` as const,
    };
  });
  const shortcutItems: ShortcutArray<`model-${string}` | "back"> = [
    {
      type: "auto-list" as const,
      order: numericItems,
    },
    {
      type: "key" as const,
      mapping: {
        b: {
          label: "Back to main menu",
          value: "back",
        },
      },
    },
  ];
  const onSelect = useCallback(
    async (item: Item<`model-${string}` | "back">) => {
      if (item.value === "back") {
        setMenuMode("main-menu");
        return;
      }
      const target = item.value.replace("model-", "");
      const model = config.models.find(m => m.nickname === target)!;
      const rest = config.models.filter(m => m.nickname !== target);
      await setConfig({
        ...config,
        models: [model, ...rest],
      });
      setModelOverride(target, session);
      setMenuMode("main-menu");
      toggleMenu();
    },
    [config, setModelOverride, toggleMenu, session],
  );
  return (
    <KbShortcutPanel
      title="Which model should be the default?"
      shortcutItems={shortcutItems}
      onSelect={onSelect}
    />
  );
}
function RemoveModelMenu() {
  const { setModelOverride, toggleMenu } = useAppStore(
    useShallow(state => ({
      setModelOverride: state.setModelOverride,
      toggleMenu: state.toggleMenu,
    })),
  );
  const session = useSession();
  const config = useConfig();
  const setConfig = useSetConfig();
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  useKeyboard(event => {
    if (event.key === "Escape") setMenuMode("main-menu");
  });
  const numericItems = config.models.map(model => {
    return {
      label: model.nickname,
      value: `model-${model.nickname}` as const,
    };
  });
  const shortcutItems: ShortcutArray<`model-${string}` | "back"> = [
    {
      type: "auto-list" as const,
      order: numericItems,
    },
    {
      type: "key" as const,
      mapping: {
        b: {
          label: "Back to main menu",
          value: "back",
        },
      },
    },
  ];
  const onSelect = useCallback(
    async (item: Item<`model-${string}` | "back">) => {
      if (item.value === "back") {
        setMenuMode("main-menu");
        return;
      }
      const target = item.value.replace("model-", "");
      const rest = config.models.filter(m => m.nickname !== target);
      await setConfig({
        ...config,
        models: [...rest],
      });
      const current = rest[0];
      setModelOverride(current.nickname, session);
      setMenuMode("main-menu");
      toggleMenu();
    },
    [config, setModelOverride, toggleMenu, session],
  );
  return (
    <KbShortcutPanel
      title="Which model do you want to remove?"
      shortcutItems={shortcutItems}
      onSelect={onSelect}
    />
  );
}
function AddModelMenuFlow() {
  const { setMenuMode } = useMenuState(
    useShallow(state => ({
      setMenuMode: state.setMenuMode,
    })),
  );
  const setConfig = useSetConfig();
  const config = useConfig();
  const onComplete = useCallback(
    async (models: Config["models"]) => {
      await setConfig({
        ...config,
        models: [...config.models, ...models],
      });
      setMenuMode("model-select");
    },
    [config, setConfig],
  );
  const onCancel = useCallback(() => {
    setMenuMode("main-menu");
  }, [setMenuMode]);
  const onOverrideDefaultApiKey = useCallback(
    async (overrides: Record<string, string>) => {
      await setConfig({
        ...config,
        defaultApiKeyOverrides: {
          ...(config.defaultApiKeyOverrides || {}),
          ...overrides,
        },
      });
    },
    [config, setConfig],
  );
  return (
    <ModelSetup
      config={config}
      onComplete={onComplete}
      onCancel={onCancel}
      onOverrideDefaultApiKey={onOverrideDefaultApiKey}
    />
  );
}