'use client';

import { Badge } from '@tale/ui/badge';
import { useLocale } from '@tale/ui/i18n/locale-provider';
import { ChevronDown, Cpu } from 'lucide-react';
import {
  type ReactNode,
  useCallback,
  useEffect,
  useMemo,
  useState,
} from 'react';

import {
  SearchableSelect,
  type SearchableSelectOption,
} from '@/app/components/ui/forms/searchable-select';
import { useAccessibleModels } from '@/app/features/settings/governance/hooks/queries';
import { useListProviders } from '@/app/features/settings/providers/hooks/queries';
import { useT } from '@/lib/i18n/client';
import {
  expandModelVariants,
  getVariantBadgeLabel,
} from '@/lib/shared/utils/expand-model-variants';
import {
  parseModelRef,
  stripModelRefQualifier,
} from '@/lib/shared/utils/model-ref';
import { resolveModelLocale } from '@/lib/shared/utils/resolve-provider-locale';

import { useChatAgents } from '../../hooks/queries';
import { useEffectiveAgent } from '../../hooks/use-effective-agent';
import { ModelTagIcons } from '../model-tag-icons';
import { useArenaMode } from './arena-mode-context';

interface ArenaModelSelectorProps {
  organizationId: string;
}

function getModelShortName(modelId: string): string {
  const slash = modelId.lastIndexOf('/');
  return slash >= 0 ? modelId.slice(slash + 1) : modelId;
}

export function ArenaModelSelector({
  organizationId,
}: ArenaModelSelectorProps) {
  const { t } = useT('chat');
  const { agent: effectiveAgent } = useEffectiveAgent(organizationId);
  const { agents } = useChatAgents(organizationId);
  const { providers } = useListProviders(organizationId);
  const { locale } = useLocale();
  const { modelA, modelB, setModelA, setModelB } = useArenaMode();
  const [openA, setOpenA] = useState(false);
  const [openB, setOpenB] = useState(false);

  const supportedModels = useMemo(() => {
    const agent = agents?.find((a) => a.name === effectiveAgent?.name);
    return agent?.supportedModels ?? [];
  }, [agents, effectiveAgent?.name]);

  const modelInfoMap = useMemo(() => {
    const map = new Map<
      string,
      {
        displayName: string;
        description?: string;
        tags: string[];
        quantizations?: string[];
      }
    >();
    for (const provider of providers) {
      if (
        !provider ||
        !('models' in provider) ||
        !Array.isArray(provider.models)
      )
        continue;
      for (const model of provider.models) {
        const resolved = resolveModelLocale(model, provider.i18n, locale);
        map.set(model.id, {
          displayName: resolved.displayName || model.displayName,
          description: resolved.description || undefined,
          tags: model.tags ?? [],
          quantizations: Array.isArray(model.quantizations)
            ? model.quantizations
            : undefined,
        });
      }
    }
    return map;
  }, [providers, locale]);

  const renderTagIcons = useCallback(
    (option: SearchableSelectOption): ReactNode => {
      const info = modelInfoMap.get(stripModelRefQualifier(option.value));
      if (!info?.tags.length) return null;
      return <ModelTagIcons tags={info.tags} t={t} />;
    },
    [modelInfoMap, t],
  );

  const getDisplayName = useCallback(
    (ref: string) => {
      const { modelId, quantization } = parseModelRef(ref);
      const base =
        modelInfoMap.get(modelId)?.displayName ?? getModelShortName(modelId);
      // Append the variant in the closed trigger so fp8 vs fp4 selections
      // are distinguishable without opening the menu.
      return quantization
        ? `${base} (${getVariantBadgeLabel(quantization)})`
        : base;
    },
    [modelInfoMap],
  );

  // Governance policies match on plain model ids; strip qualifiers before asking.
  const plainModelIds = useMemo(
    () => supportedModels.map(stripModelRefQualifier),
    [supportedModels],
  );
  const { data: accessibleModelIds } = useAccessibleModels(
    organizationId,
    plainModelIds,
  );

  const filteredModels = useMemo(() => {
    const accessible = accessibleModelIds ? new Set(accessibleModelIds) : null;
    const accessFiltered = accessible
      ? supportedModels.filter((ref) =>
          accessible.has(stripModelRefQualifier(ref)),
        )
      : supportedModels;
    return expandModelVariants(
      accessFiltered,
      (bareId) => modelInfoMap.get(bareId)?.quantizations,
    );
  }, [supportedModels, accessibleModelIds, modelInfoMap]);

  const options = useMemo(
    () =>
      filteredModels.map((ref) => {
        const { quantization } = parseModelRef(ref);
        return {
          value: ref,
          label: getDisplayName(ref),
          labelBadge: quantization ? (
            <Badge variant="outline" className="text-[10px] font-normal">
              {getVariantBadgeLabel(quantization)}
            </Badge>
          ) : undefined,
          description: modelInfoMap.get(stripModelRefQualifier(ref))
            ?.description,
        };
      }),
    [filteredModels, getDisplayName, modelInfoMap],
  );

  // A previously persisted selection may no longer appear in `filteredModels`
  // — e.g. governance tightened, or expansion now offers `@fp8`/`@fp4`
  // variants in place of a bare `glm-5.1` ref. Treat the stale value as
  // empty when computing the trigger so SearchableSelect's `value` always
  // matches an option.
  const isInList = useCallback(
    (ref: string | null): boolean =>
      ref != null && filteredModels.includes(ref),
    [filteredModels],
  );
  const currentModelA = isInList(modelA) ? modelA : (filteredModels[0] ?? null);
  const currentModelB = isInList(modelB)
    ? modelB
    : (filteredModels[1] ?? filteredModels[0] ?? null);

  // Sync default selections back to context so sendMessage can read them.
  // Also clear stale values that no longer appear in `filteredModels` so the
  // backend doesn't dispatch to a model the user can't see in the picker.
  useEffect(() => {
    if (filteredModels.length < 2) return;
    if (modelA && !filteredModels.includes(modelA)) {
      setModelA(filteredModels[0]);
    } else if (!modelA && filteredModels[0]) {
      setModelA(filteredModels[0]);
    }
    if (modelB && !filteredModels.includes(modelB)) {
      setModelB(filteredModels[1] ?? filteredModels[0]);
    } else if (!modelB && filteredModels[1]) {
      setModelB(filteredModels[1]);
    }
  }, [filteredModels, modelA, modelB, setModelA, setModelB]);

  if (filteredModels.length < 2) return null;

  return (
    <div className="flex items-center gap-3 px-2">
      <div className="flex items-center gap-1.5">
        <span className="text-muted-foreground text-xs font-medium">A</span>
        <SearchableSelect
          value={currentModelA}
          onValueChange={setModelA}
          options={options}
          open={openA}
          onOpenChange={setOpenA}
          align="start"
          side="top"
          sideOffset={8}
          contentClassName="w-[22rem]"
          searchPlaceholder={t('modelSelector.searchPlaceholder')}
          emptyText={t('modelSelector.noResults')}
          aria-label={t('arena.modelA')}
          optionAction={renderTagIcons}
          trigger={
            <button
              type="button"
              className="text-muted-foreground hover:text-foreground flex items-center gap-1.5 text-xs transition-colors"
              aria-label={t('arena.modelA')}
            >
              <Cpu className="size-3.5" aria-hidden="true" />
              <span>
                {currentModelA
                  ? getDisplayName(currentModelA)
                  : t('arena.selectModel')}
              </span>
              <ChevronDown className="size-3" aria-hidden="true" />
            </button>
          }
        />
      </div>
      <span className="text-muted-foreground text-xs">{t('arena.vs')}</span>
      <div className="flex items-center gap-1.5">
        <span className="text-muted-foreground text-xs font-medium">B</span>
        <SearchableSelect
          value={currentModelB}
          onValueChange={setModelB}
          options={options}
          open={openB}
          onOpenChange={setOpenB}
          align="start"
          side="top"
          sideOffset={8}
          contentClassName="w-[22rem]"
          searchPlaceholder={t('modelSelector.searchPlaceholder')}
          emptyText={t('modelSelector.noResults')}
          aria-label={t('arena.modelB')}
          optionAction={renderTagIcons}
          trigger={
            <button
              type="button"
              className="text-muted-foreground hover:text-foreground flex items-center gap-1.5 text-xs transition-colors"
              aria-label={t('arena.modelB')}
            >
              <Cpu className="size-3.5" aria-hidden="true" />
              <span>
                {currentModelB
                  ? getDisplayName(currentModelB)
                  : t('arena.selectModel')}
              </span>
              <ChevronDown className="size-3" aria-hidden="true" />
            </button>
          }
        />
      </div>
    </div>
  );
}
