import {
  Action,
  ActionPanel,
  Alert,
  Form,
  Icon,
  List,
  Toast,
  confirmAlert,
  openExtensionPreferences,
  showToast,
} from "@vicinae/api";
import { useCallback, useEffect, useState } from "react";
import { usePreferences } from "./lib/hooks/usePreferences";
import {
  deleteModel,
  describePullProgress,
  isSelfHosted,
  listModels,
  listRunningModels,
  pullModel,
  unloadModel,
  type RunningModel,
} from "./lib/ollama";
import { clearModelSelection, selectModel, type ResolvedPreferences } from "./lib/preferences";
import type { ModelInfo } from "./lib/types";
import { errorMessage, formatBytes } from "./lib/utils";

/**
 * Browse, pick, and manage the models the configured Ollama host can reach.
 *
 * Four jobs in one command, because they are four answers to the same question
 * ("what models do I have?"): pick one, unload one, delete one, pull a new one.
 * Unloading is the odd one out and the reason a row needs two states rather than
 * one: a model can be installed without being loaded, and only the loaded ones
 * are holding memory. Delete and pull only appear on a host the user runs, since
 * there is nothing to pull onto Ollama Cloud and nothing to delete from it.
 *
 * ## Where a pick is stored
 *
 * Not in the `Model` preference, because an extension cannot write its own
 * preferences. Vicinae reads them once at startup and holds them in memory, so
 * even writing its config file directly does nothing until it restarts, which was
 * measured rather than assumed (`providers["@adrianbonpin/lam-agent"].preferences`
 * in `~/.config/vicinae/settings.json`, edited while a command was then launched:
 * the requests still went to the old host).
 *
 * So a pick goes to LocalStorage and `loadPreferences` layers it over the
 * preference. The cost of that arrangement is one thing this screen has to say
 * out loud, because it is invisible otherwise: while a pick is active, editing
 * the `Model` preference has no effect until the pick is reset. Which is why the
 * setup screen shows both values and offers the reset.
 *
 * ## Why the pull keeps the window open
 *
 * The download runs inside this command's worker, so closing the launcher cancels
 * it. That is a real constraint and the command description says so.
 */

/** Which surface the command is showing. The form is a view rather than a dialog, which Vicinae has no equivalent of. */
type View = "list" | "pull";

export default function ModelsCommand() {
  const { preferences, reload } = usePreferences();

  // Nothing renders until the model in use is known, because the whole list is
  // marked up relative to it: which row is current, and whether picking one
  // would change anything.
  if (!preferences) {
    return <List isLoading navigationTitle="Ollama Models" searchBarPlaceholder="Loading models…" />;
  }

  return <ModelsScreen preferences={preferences} reloadPreferences={reload} />;
}

interface ModelsScreenProps {
  preferences: ResolvedPreferences;
  /** Re-reads preferences after this screen changes the model selection. */
  reloadPreferences: () => void;
}

function ModelsScreen({ preferences, reloadPreferences }: ModelsScreenProps) {
  const [models, setModels] = useState<ModelInfo[] | undefined>(undefined);
  /** By name, so a row can ask whether it is the one holding memory. */
  const [running, setRunning] = useState<Map<string, RunningModel>>(new Map());
  const [error, setError] = useState<string | undefined>(undefined);
  const [isLoading, setIsLoading] = useState(true);
  const [view, setView] = useState<View>("list");

  /** Whether this host's models can be managed at all. */
  const selfHosted = isSelfHosted(preferences.host);

  const load = useCallback(async () => {
    setIsLoading(true);

    try {
      // Both at once: what is installed and what is loaded are different
      // questions, and the second must not hold up the first.
      const [installed, loaded] = await Promise.all([
        listModels({ host: preferences.host, apiKey: preferences.apiKey }),
        listRunningModels({ host: preferences.host, apiKey: preferences.apiKey }),
      ]);

      setModels(installed);
      setRunning(new Map(loaded.map((model) => [model.name, model])));
      setError(undefined);
    } catch (thrown) {
      setModels([]);
      setRunning(new Map());
      setError(errorMessage(thrown));
    } finally {
      setIsLoading(false);
    }
  }, [preferences.host, preferences.apiKey]);

  useEffect(() => {
    void load();
  }, [load]);

  /**
   * Releases a model from memory.
   *
   * No confirmation, unlike deleting: nothing is removed from disk, and the
   * worst case is that the next question pays to load it again. The toast reports
   * what was freed, because that is the whole reason to do this.
   */
  const unload = useCallback(
    async (model: ModelInfo, loaded: RunningModel) => {
      try {
        await unloadModel({ host: preferences.host, apiKey: preferences.apiKey, model: model.name });
      } catch (thrown) {
        await showToast({
          style: Toast.Style.Failure,
          title: `Could not unload ${model.name}`,
          message: errorMessage(thrown),
        });
        return;
      }

      const freed = formatBytes(loaded.sizeVram);
      await showToast({
        style: Toast.Style.Success,
        title: `Unloaded ${model.name}`,
        message: freed ? `${freed} of memory freed` : undefined,
      });

      void load();
    },
    [load, preferences.host, preferences.apiKey],
  );

  const choose = useCallback(
    async (name: string) => {
      if (name === preferences.model) {
        await showToast({ style: Toast.Style.Success, title: `${name} is already in use` });
        return;
      }

      await selectModel(name);
      reloadPreferences();

      await showToast({
        style: Toast.Style.Success,
        title: `Lam will use ${name}`,
        // Said here because it is the one surprising part of this screen: the
        // model is chosen in this command, not in Vicinae's settings.
        message: "Chosen here. The Model preference is left alone.",
      });
    },
    [preferences.model, reloadPreferences],
  );

  const remove = useCallback(
    async (model: ModelInfo) => {
      const confirmed = await confirmAlert({
        title: `Delete ${model.name}?`,
        message: model.remoteHost
          ? `This removes the local reference. The model itself lives at ${model.remoteHost} and is not deleted there.`
          : "The weights are removed from this machine, and pulling it again means downloading all of it.",
        primaryAction: { title: "Delete", style: Alert.ActionStyle.Destructive },
      });
      if (!confirmed) return;

      try {
        await deleteModel({ host: preferences.host, apiKey: preferences.apiKey, model: model.name });
      } catch (thrown) {
        await showToast({
          style: Toast.Style.Failure,
          title: `Could not delete ${model.name}`,
          message: errorMessage(thrown),
        });
        return;
      }

      // The model that just went may be the one in use, in which case the next
      // question asked would fail with a 404. A selection can be cleared here;
      // a preference cannot, so that case is reported instead of silently
      // leaving a dangling name behind.
      const wasInUse = model.name === preferences.model;
      const selected = preferences.modelSource === "selection";

      if (wasInUse && selected) {
        await clearModelSelection();
        reloadPreferences();
      }

      await showToast({
        style: Toast.Style.Success,
        title: `Deleted ${model.name}`,
        message: wasInUse
          ? selected
            ? "It was the model in use, so that choice has been cleared."
            : "It is still what the Model preference points at. Change it in preferences."
          : undefined,
      });

      void load();
    },
    [load, preferences.host, preferences.apiKey, preferences.model, preferences.modelSource, reloadPreferences],
  );

  if (view === "pull") {
    return (
      <PullModelForm
        host={preferences.host}
        apiKey={preferences.apiKey}
        onDone={() => {
          setView("list");
          void load();
        }}
        onCancel={() => setView("list")}
      />
    );
  }

  const pullAction = selfHosted ? (
    <Action title="Pull a Model…" icon={Icon.Download} onAction={() => setView("pull")} />
  ) : null;

  const listActions = (
    <ActionPanel>
      <Action title="Refresh" icon={Icon.ArrowClockwise} onAction={() => void load()} />
      {pullAction ? <ActionPanel.Section>{pullAction}</ActionPanel.Section> : null}
      <ActionPanel.Section>
        <Action
          title="Open Extension Preferences"
          icon={Icon.BulletPoints}
          onAction={() => void openExtensionPreferences()}
        />
      </ActionPanel.Section>
    </ActionPanel>
  );

  const loaded = models ?? [];

  return (
    <List
      isLoading={isLoading}
      navigationTitle="Ollama Models"
      searchBarPlaceholder="Search models…"
      actions={listActions}
    >
      {error ? (
        <List.EmptyView
          icon={Icon.QuestionMarkCircle}
          title="Could not reach Ollama"
          description={error}
          actions={
            <ActionPanel>
              <Action title="Try Again" icon={Icon.ArrowClockwise} onAction={() => void load()} />
              <Action
                title="Open Extension Preferences"
                icon={Icon.BulletPoints}
                onAction={() => void openExtensionPreferences()}
              />
            </ActionPanel>
          }
        />
      ) : models !== undefined && loaded.length === 0 ? (
        <List.EmptyView
          icon={Icon.Box}
          title={selfHosted ? "No models on this server" : "No models in this account"}
          description={
            selfHosted
              ? "Pull one from the action panel, or with `ollama pull` in a terminal."
              : "This host has nothing installed, and models are managed at ollama.com."
          }
          actions={listActions}
        />
      ) : null}

      {loaded.map((model) => {
        const isCurrent = model.name === preferences.model;
        const size = model.remoteHost ? undefined : formatBytes(model.size);
        const inMemory = running.get(model.name);

        return (
          <List.Item
            key={model.name}
            id={model.name}
            icon={isCurrent ? Icon.CheckCircle : Icon.Box}
            title={model.name}
            subtitle={[model.parameterSize, size].filter(Boolean).join(" · ") || undefined}
            accessories={[
              ...(inMemory ? [{ tag: { value: "loaded", color: "green" as const } }] : []),
              ...(isCurrent ? [{ tag: { value: "current", color: "blue" as const } }] : []),
              ...(model.capabilities?.includes("vision") ? [{ tag: "vision" }] : []),
              { tag: model.remoteHost ? "cloud" : selfHosted ? "local" : "cloud" },
            ]}
            actions={
              <ActionPanel>
                <Action
                  title={isCurrent ? "Current Model" : "Use This Model"}
                  icon={isCurrent ? Icon.CheckCircle : Icon.Stars}
                  onAction={() => void choose(model.name)}
                />
                {inMemory ? (
                  <Action
                    title="Unload Model"
                    icon={Icon.Eject}
                    onAction={() => void unload(model, inMemory)}
                  />
                ) : null}
                <Action.CopyToClipboard title="Copy Model Name" content={model.name} />
                {selfHosted ? (
                  <Action
                    title="Delete Model"
                    icon={Icon.Trash}
                    style="destructive"
                    onAction={() => void remove(model)}
                  />
                ) : null}
                <ActionPanel.Section>
                  <Action title="Refresh" icon={Icon.ArrowClockwise} onAction={() => void load()} />
                  {pullAction}
                  <Action
                    title="Open Extension Preferences"
                    icon={Icon.BulletPoints}
                    onAction={() => void openExtensionPreferences()}
                  />
                </ActionPanel.Section>
              </ActionPanel>
            }
          />
        );
      })}
    </List>
  );
}

interface PullModelFormProps {
  host: string;
  apiKey: string;
  onDone: () => void;
  onCancel: () => void;
}

/**
 * Pull a model by name.
 *
 * A `Form` rather than a `List`, because the model being pulled is usually not
 * one of the models already installed, which is the whole reason for pulling it.
 * A picker over the library would be nicer and needs an endpoint this API does
 * not expose, so the name is typed and the description carries a few examples.
 */
function PullModelForm({ host, apiKey, onDone, onCancel }: PullModelFormProps) {
  const [isPulling, setIsPulling] = useState(false);

  const pull = useCallback(
    async (values: Form.Values) => {
      const name = String(values.model ?? "").trim();

      if (!name) {
        await showToast({ style: Toast.Style.Failure, title: "Give the model a name" });
        return;
      }

      setIsPulling(true);
      const toast = await showToast({ style: Toast.Style.Animated, title: `Pulling ${name}` });

      try {
        let last = "";

        for await (const progress of pullModel({ host, apiKey, model: name })) {
          // Ollama sends a tick per chunk, which for a large model is far more
          // often than the toast needs updating. Only write when the line
          // actually changed, so the percentage is what limits the traffic.
          const line = describePullProgress(progress);
          if (line !== last) {
            last = line;
            toast.message = line;
          }
        }

        toast.style = Toast.Style.Success;
        toast.title = `Pulled ${name}`;
        toast.message = "It is in the list now";
        onDone();
      } catch (thrown) {
        toast.style = Toast.Style.Failure;
        toast.title = `Could not pull ${name}`;
        toast.message = errorMessage(thrown);
      } finally {
        setIsPulling(false);
      }
    },
    [apiKey, host, onDone],
  );

  return (
    <Form
      isLoading={isPulling}
      navigationTitle="Pull a Model"
      actions={
        <ActionPanel>
          <Action.SubmitForm title="Pull Model" icon={Icon.Download} onSubmit={pull} />
          <Action title="Cancel" icon={Icon.Xmark} onAction={onCancel} />
        </ActionPanel>
      }
    >
      <Form.TextField
        id="model"
        title="Model"
        placeholder="qwen3:8b"
        autoFocus
        info="The name as it appears in the Ollama library, including the tag. Leave the tag off to get the default one."
      />
      <Form.Description
        title="Examples"
        text={["qwen3:8b", "llama3.2:3b", "gemma3:12b", "deepseek-r1:14b"].join("\n")}
      />
      <Form.Description
        title="While it pulls"
        text="Leave the launcher open. The download runs in this window, so closing it cancels the pull."
      />
    </Form>
  );
}