/**
 * Which model answers, and what that costs.
 *
 * This is a small control carrying a large argument. Every entry in the
 * catalogue — NVIDIA's Nemotron, DeepSeek, Anthropic's Opus, OpenAI's GPT — is
 * served by Prime Intellect's own inference on a single API key. Nowhere else
 * in PIG is that visible; a transcript footer naming the model is a fact, but a
 * menu of five vendors under one key is the pitch. So the menu says so, once,
 * quietly, at the bottom.
 *
 * The catalogue is the server's, never this file's. The relay refuses any id it
 * did not send, so a hard-coded option that has been retired upstream is a menu
 * entry whose only effect is a 400 — and a *stored* id that has been retired is
 * the same 400 on every turn until someone clears their browser storage. Both
 * are handled below rather than left to the user to discover.
 *
 * On presenting cost: "$0.05 / $0.20 per Mtok" is the unit providers publish
 * and it is meaningless to the sales lead this product is for. The headline
 * figure is therefore an estimate of what a hundred questions cost, derived
 * from a real measured read-only turn, with the raw per-Mtok rates kept on a
 * secondary line for whoever wants to check the arithmetic.
 */
import { useCallback, useEffect, useMemo, useRef, useState, type JSX } from 'react';
import { useQuery } from '@tanstack/react-query';
// `Cpu`, not `Sparkles`: Sparkles used to mean both "Piggy" and "model", and a
// glyph that means two things means neither. Piggy is `PiggyMark` everywhere
// now, so this control needs a neutral mark of its own — and inference running
// on Prime Intellect's own silicon is the thing this menu is about.
import { ChevronsUpDown, Cpu } from 'lucide-react';
import type { PiggyModelOption } from '@pig/core';
import { fetchPiggyModels } from '@/lib/piggy-chat';
import { useIdentityQuery } from '@/lib/identity';
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuLabel,
  DropdownMenuRadioGroup,
  DropdownMenuRadioItem,
  DropdownMenuSeparator,
  DropdownMenuTrigger,
} from '@/components/ui/dropdown-menu';
import { Badge, Button, Label, Skeleton, cn } from '@/components/ui';

// ------------------------------------------------------------------ catalogue

export const PIGGY_MODELS_QUERY_KEY = ['piggy', 'models'] as const;

/**
 * A stable empty array for the pending and failed cases.
 *
 * `?? []` would hand every render a new reference, which is enough to re-run
 * any effect below that depends on the list — including the one that repairs an
 * invalid selection, which would then loop.
 */
const NO_MODELS: PiggyModelOption[] = [];

export interface PiggyModelCatalogue {
  models: PiggyModelOption[];
  defaultModelId: string | null;
  isLoading: boolean;
  error: Error | null;
}

/**
 * The models this deployment offers.
 *
 * Cached for the life of the tab: the catalogue is deployment configuration,
 * not data. It cannot change while the page is open, and a refetch on window
 * focus would put a network round trip behind a control the user is in the act
 * of opening.
 */
export function usePiggyModels(): PiggyModelCatalogue {
  const query = useQuery({
    queryKey: PIGGY_MODELS_QUERY_KEY,
    queryFn: fetchPiggyModels,
    staleTime: Infinity,
    gcTime: Infinity,
    retry: 1,
  });

  return {
    models: query.data?.models ?? NO_MODELS,
    defaultModelId: query.data?.defaultModelId ?? null,
    isLoading: query.isLoading,
    error: toError(query.error),
  };
}

function toError(value: unknown): Error | null {
  if (!value) return null;
  return value instanceof Error ? value : new Error(String(value));
}

// ----------------------------------------------------------------- what it costs

/**
 * A real read-only turn, measured against nemotron on the live stack: roughly
 * 5,300 tokens in (the system prompt, the tool schemas and the CRM context
 * dominate) and 150 out. Every price in this menu is that same turn priced on a
 * different model, which is the only way five figures spanning a hundredfold
 * are comparable at a glance.
 */
const TYPICAL_INPUT_TOKENS = 5_300;
const TYPICAL_OUTPUT_TOKENS = 150;

/**
 * A single question on the cheapest model costs three hundredths of a cent, and
 * "$0.0003" is a number nobody can rank against another number. Quoting a
 * hundred questions puts the whole catalogue in the range people actually price
 * things in — three cents to three dollars.
 */
const QUOTED_QUESTIONS = 100;

const QUOTE_FORMAT = new Intl.NumberFormat('en-US', {
  style: 'currency',
  currency: 'USD',
  minimumFractionDigits: 2,
  maximumFractionDigits: 2,
});

/**
 * Rates are published as round dollars ($5, $25) as often as fractions ($0.05).
 * Two formatters rather than one: a single `minimumFractionDigits: 0` renders
 * $0.20 as "$0.2", which reads as a typo next to "$0.05" in the same column.
 */
const WHOLE_RATE_FORMAT = new Intl.NumberFormat('en-US', {
  style: 'currency',
  currency: 'USD',
  minimumFractionDigits: 0,
  maximumFractionDigits: 0,
});
const PART_RATE_FORMAT = new Intl.NumberFormat('en-US', {
  style: 'currency',
  currency: 'USD',
  minimumFractionDigits: 2,
  maximumFractionDigits: 2,
});

function formatRate(dollars: number): string {
  return Number.isInteger(dollars)
    ? WHOLE_RATE_FORMAT.format(dollars)
    : PART_RATE_FORMAT.format(dollars);
}

/**
 * US dollars for one typical question.
 *
 * `costPerMTok*` is dollars per million tokens — not cents, and deliberately
 * not, per the note in the protocol. Nothing here may be run through the
 * `Cents` formatters in lib/api.
 */
function dollarsPerQuestion(model: PiggyModelOption): number {
  return (
    (TYPICAL_INPUT_TOKENS * model.costPerMTokIn + TYPICAL_OUTPUT_TOKENS * model.costPerMTokOut) /
    1_000_000
  );
}

function formatQuote(model: PiggyModelOption): string {
  const total = dollarsPerQuestion(model) * QUOTED_QUESTIONS;
  // A model cheap enough to round to zero would otherwise be quoted "$0.00",
  // which reads as free rather than as very cheap.
  return total > 0 && total < 0.01 ? 'under $0.01' : QUOTE_FORMAT.format(total);
}

function formatRates(model: PiggyModelOption): string {
  return `${formatRate(model.costPerMTokIn)} in / ${formatRate(model.costPerMTokOut)} out per Mtok`;
}

function formatContext(tokens: number): string {
  return tokens >= 1_000 ? `${Math.round(tokens / 1_000)}K` : String(tokens);
}

/**
 * How close to the dearest model a model has to be to count as top tier.
 *
 * The two frontier entries are priced within a few per cent of each other, and
 * naming only the very dearest "most capable" would be PIG picking a winner
 * between two vendors on a rounding difference. A band names the tier instead,
 * which is the true statement.
 */
const TOP_TIER_RATIO = 0.85;

/**
 * Which entry is cheapest, and which are the ones to reach for when it matters.
 *
 * Price is the proxy for capability, because the catalogue carries no capability
 * score and price is the only ordering the server actually sends. The
 * alternative — a list of model ids ranked in this file — is a second source of
 * truth that goes stale the first time the deployment adds a model.
 */
function priceBands(models: PiggyModelOption[]): {
  cheapestId: string | null;
  topTierIds: ReadonlySet<string>;
} {
  const empty: ReadonlySet<string> = new Set<string>();
  if (models.length < 2) return { cheapestId: null, topTierIds: empty };

  let cheapest: PiggyModelOption | null = null;
  let dearest = 0;
  for (const model of models) {
    const cost = dollarsPerQuestion(model);
    if (!cheapest || cost < dollarsPerQuestion(cheapest)) cheapest = model;
    if (cost > dearest) dearest = cost;
  }
  if (!cheapest || dearest <= 0) return { cheapestId: null, topTierIds: empty };

  const cheapestId = cheapest.id;
  const topTierIds = new Set(
    models
      // The cheapest model is never also the top tier, however flat the
      // catalogue's pricing happens to be.
      .filter(
        (model) =>
          model.id !== cheapestId && dollarsPerQuestion(model) >= dearest * TOP_TIER_RATIO,
      )
      .map((model) => model.id),
  );
  return { cheapestId, topTierIds };
}

// ------------------------------------------------------------------- persistence

const STORAGE_PREFIX = 'pig:piggy:model';

function storageKey(userId: string | null): string {
  return userId ? `${STORAGE_PREFIX}:${userId}` : STORAGE_PREFIX;
}

export function readStoredPiggyModelId(userId: string | null): string | null {
  try {
    return localStorage.getItem(storageKey(userId));
  } catch {
    /* Private browsing throws on localStorage. The deployment default is fine. */
    return null;
  }
}

export function writeStoredPiggyModelId(userId: string | null, modelId: string | null): void {
  try {
    const key = storageKey(userId);
    if (modelId) localStorage.setItem(key, modelId);
    else localStorage.removeItem(key);
  } catch {
    /* As above: an unpersisted preference is a smaller problem than a throw. */
  }
}

export interface PiggyModelChoice extends PiggyModelCatalogue {
  /**
   * The id to send with a turn: the stored choice when the catalogue still
   * lists it, otherwise the deployment default. Null only while the catalogue
   * is loading or unavailable, in which case send no `modelId` at all.
   */
  modelId: string | null;
  /** True when the user picked this, false when it is the deployment default. */
  isExplicit: boolean;
  setModelId: (modelId: string) => void;
}

/**
 * The choice, persisted per user, with the server as the authority on validity.
 *
 * Pair this with `PiggyModelPicker` — `value={modelId}` and
 * `onChange={setModelId}` — rather than a plain `useState`, or the preference
 * is remembered for the session only.
 */
export function usePiggyModelChoice(): PiggyModelChoice {
  const catalogue = usePiggyModels();
  const userId = useIdentityQuery().data?.id ?? null;
  const [stored, setStored] = useState<string | null>(() => readStoredPiggyModelId(userId));

  // On a cold cache the signed-in user arrives a tick after first render, so
  // the key this first read from was the anonymous one. Re-read once it settles
  // rather than showing whatever the previous person on this browser chose.
  useEffect(() => {
    setStored(readStoredPiggyModelId(userId));
  }, [userId]);

  const known = catalogue.models.some((model) => model.id === stored);

  // A stored id the relay no longer lists is a guaranteed 400 on every
  // subsequent turn, and the user has no way to connect that error to a choice
  // they made weeks ago. Drop it as soon as the catalogue contradicts it.
  useEffect(() => {
    if (!stored || known) return;
    if (catalogue.isLoading || catalogue.models.length === 0) return;
    writeStoredPiggyModelId(userId, null);
    setStored(null);
  }, [stored, known, catalogue.isLoading, catalogue.models, userId]);

  const setModelId = useCallback(
    (modelId: string) => {
      writeStoredPiggyModelId(userId, modelId);
      setStored(modelId);
    },
    [userId],
  );

  const isExplicit = Boolean(stored) && known;
  return {
    ...catalogue,
    modelId: isExplicit ? stored : catalogue.defaultModelId,
    isExplicit,
    setModelId,
  };
}

// ------------------------------------------------------------------------ picker

/**
 * The label the tight trigger shows.
 *
 * The protocol carries no short form, and the docked panel header is 22rem
 * wide — enough for about a dozen characters once the icon and chevron are
 * paid for. Dropping the family prefix from a long name keeps the part that
 * distinguishes it ("Nano 30B", "Super 120B") rather than the part every entry
 * in a family shares; shorter names are already short enough to leave alone.
 */
const LONG_LABEL_WORDS = 4;

function shortLabel(label: string): string {
  const words = label.split(/\s+/).filter(Boolean);
  return words.length >= LONG_LABEL_WORDS ? words.slice(-2).join(' ') : label;
}

/**
 * The human name for a model id, for surfaces that only have the id.
 *
 * The transcript footer used to print `nvidia/nemotron-3-super-120b` in
 * monospace under every answer — the raw wire value, in the one typeface this
 * product reserves for machine text, on the line meant to tell a sales lead
 * which model answered them. The catalogue already carries the name the picker
 * shows two inches above it, so the footer says the same thing the menu says.
 *
 * The fallback drops the provider prefix rather than inventing capitalisation:
 * a deployment can list a model this browser's cached catalogue has never
 * seen, and a guessed name is worse than an honest identifier.
 */
export function usePiggyModelLabel(modelId: string | undefined): string | undefined {
  const { models } = usePiggyModels();
  if (!modelId) return undefined;
  const known = models.find((model) => model.id === modelId);
  if (known) return known.label;
  const slash = modelId.lastIndexOf('/');
  return slash === -1 ? modelId : modelId.slice(slash + 1);
}

export interface PiggyModelPickerProps {
  /** The chosen model id, or null to follow the deployment default. */
  value: string | null;
  onChange: (modelId: string) => void;
  /** The tight rendering, for the 22rem docked panel header. */
  compact?: boolean;
  disabled?: boolean;
}

export function PiggyModelPicker({
  value,
  onChange,
  compact = false,
  disabled = false,
}: PiggyModelPickerProps): JSX.Element {
  const { models, defaultModelId, isLoading, error } = usePiggyModels();
  const userId = useIdentityQuery().data?.id ?? null;

  const selected = models.find((model) => model.id === value) ?? null;
  const fallback =
    models.find((model) => model.id === defaultModelId) ??
    models.find((model) => model.isDefault) ??
    null;
  const inForce = selected ?? fallback;
  /**
   * Whether the model in force is the deployment's own default.
   *
   * Deliberately a fact about the model, not about how it was arrived at. The
   * caller may resolve a null preference to the default id before passing it
   * in, so "the user made no choice" is not reliably visible here — and it is
   * not the interesting question anyway. What the user needs to know is which
   * model is answering and whether that is the shipped one.
   */
  const isDeploymentDefault = Boolean(inForce && inForce.id === defaultModelId);

  const { cheapestId, topTierIds } = useMemo(() => priceBands(models), [models]);

  /**
   * Repair a selection the catalogue does not list.
   *
   * The owner of `value` may be persisting it itself, or restoring it from
   * somewhere this component cannot see, so showing the default while the
   * caller still holds a retired id would render one model and send another.
   * Correcting the caller is the only fix that reaches the request. Guarded by
   * the id already repaired, so a caller that ignores `onChange` gets one
   * attempt rather than an infinite loop.
   */
  const repaired = useRef<string | null>(null);
  useEffect(() => {
    if (!value || isLoading || models.length === 0 || !defaultModelId) return;
    if (models.some((model) => model.id === value)) return;
    if (repaired.current === value) return;
    repaired.current = value;
    writeStoredPiggyModelId(userId, null);
    onChange(defaultModelId);
  }, [value, isLoading, models, defaultModelId, onChange, userId]);

  const handleSelect = useCallback(
    (modelId: string) => {
      // Written here as well as in `usePiggyModelChoice` so the preference
      // survives a reload however the caller holds it. Writing the same value
      // twice costs nothing; losing it because the caller used `useState`
      // costs the user their choice on every visit.
      writeStoredPiggyModelId(userId, modelId);
      onChange(modelId);
    },
    [onChange, userId],
  );

  if (isLoading) {
    return <Skeleton className={cn('h-11 rounded-lg', compact ? 'w-28' : 'w-44')} />;
  }

  if (!inForce) {
    return (
      <Button
        variant="outline"
        size="sm"
        disabled
        className={cn('gap-1.5 font-normal', compact ? 'px-2' : 'px-2.5 text-sm')}
        aria-label="The model list is unavailable"
        title={error?.message ?? 'Piggy did not return a model list.'}
      >
        <Cpu className="size-4 shrink-0 text-muted" aria-hidden />
        {compact ? null : <span className="text-muted">Model unavailable</span>}
      </Button>
    );
  }

  return (
    <DropdownMenu>
      <DropdownMenuTrigger asChild>
        <Button
          variant={compact ? 'ghost' : 'outline'}
          size="sm"
          disabled={disabled}
          aria-label={`Model: ${inForce.label}${isDeploymentDefault ? ', the deployment default' : ''}. Change the model Piggy answers with.`}
          className={cn(
            'gap-1.5 font-normal data-[state=open]:bg-surface-2',
            compact ? 'max-w-[11rem] px-2' : 'max-w-[18rem] px-2.5 text-sm',
          )}
        >
          <Cpu className="size-4 shrink-0 text-muted" aria-hidden />
          <span className="truncate text-fg">
            {compact ? shortLabel(inForce.label) : inForce.label}
          </span>
          {!compact && isDeploymentDefault ? (
            <span className="shrink-0 text-xs text-muted">Default</span>
          ) : null}
          <ChevronsUpDown className="size-3.5 shrink-0 text-muted" aria-hidden />
        </Button>
      </DropdownMenuTrigger>

      {/*
        Above the sheet and drawer primitives, which sit at z-50 themselves: the
        same trigger appears inside the mobile drawer, and a menu that opens
        behind the surface that spawned it is a control that simply does not
        work on a phone.
      */}
      <DropdownMenuContent
        align="end"
        sideOffset={6}
        className="z-[60] w-[min(26rem,calc(100vw-1.5rem))] p-1.5"
      >
        <DropdownMenuLabel className="px-2 pb-1.5 pt-1">
          <Label className="flex items-baseline justify-between gap-2">
            Model
            <span className="normal-case tracking-normal">{models.length} available</span>
          </Label>
        </DropdownMenuLabel>

        <DropdownMenuRadioGroup value={inForce.id} onValueChange={handleSelect}>
          {models.map((model) => (
            <DropdownMenuRadioItem
              key={model.id}
              value={model.id}
              // The indicator is absolutely positioned with no `top`, so it
              // would ride the top edge of a three-line row; nudged down to sit
              // against the label rather than the padding above it.
              className="items-start gap-2 rounded-md py-2.5 pl-8 pr-2 [&>span]:top-3"
            >
              <div className="flex min-w-0 flex-1 flex-col gap-1">
                <div className="flex flex-wrap items-center gap-1.5">
                  <span className="text-sm font-medium text-fg">{model.label}</span>
                  {model.id === defaultModelId ? (
                    <Badge tone="neutral" className="px-1.5 py-0">
                      Default
                    </Badge>
                  ) : null}
                  {model.id === cheapestId ? (
                    <Badge tone="positive" className="px-1.5 py-0">
                      Cheapest
                    </Badge>
                  ) : null}
                  {topTierIds.has(model.id) ? (
                    <Badge tone="accent" className="px-1.5 py-0">
                      Most capable
                    </Badge>
                  ) : null}
                </div>
                {model.hint ? (
                  <p className="whitespace-normal text-xs leading-snug text-muted">{model.hint}</p>
                ) : null}
                <p className="nums whitespace-normal text-xs leading-snug text-muted">
                  {formatRates(model)} · {formatContext(model.contextWindow)} context
                </p>
              </div>
              <div className="flex shrink-0 flex-col items-end pl-1 text-right">
                <span className="nums text-sm font-medium text-fg">{formatQuote(model)}</span>
                <span className="text-xs leading-tight text-muted">per 100 questions</span>
              </div>
            </DropdownMenuRadioItem>
          ))}
        </DropdownMenuRadioGroup>

        <DropdownMenuSeparator />

        <p className="whitespace-normal px-2 pb-1 pt-1.5 text-xs leading-snug text-muted">
          Every model here is served by Prime Intellect inference on one API key. Prices are
          estimated from a measured question — about {TYPICAL_INPUT_TOKENS.toLocaleString('en-US')}{' '}
          tokens in and {TYPICAL_OUTPUT_TOKENS} out.
        </p>
      </DropdownMenuContent>
    </DropdownMenu>
  );
}
