platform/src/apps/officer-web/Screens/Dashboard/Settings/ServerSettings/AIHarnessesSection.tsx

import { useState } from 'react';
import { Plus, Save, Trash2, RefreshCw, Loader2, X } from 'lucide-react';
import { useQuery, useQueryClient } from '@tanstack/react-query';
import { toast } from 'sonner';
import { Button } from '@/components/ui/button';
import { Input } from '@/components/ui/input';
import {
  CommandDialog,
  CommandEmpty,
  CommandGroup,
  CommandInput,
  CommandItem,
  CommandList,
} from '@/components/ui/command';
import { useClient } from 'hooks/useClient';

type VersionInfo = { version: string | null; path: string | null; globalPath: string | null };
type StoredApiKeys = { keys: { provider: string; value: string }[] };
type LocalProviderEntry = {
  id: string;
  name: string;
  url: string;
  apiType: 'ollama' | 'openai-compatible' | 'lmstudio';
  auth?: { type: 'api-key' | 'basic' };
};
type ProbeResult = {
  success: boolean;
  apiType?: LocalProviderEntry['apiType'];
  name?: string;
  needsAuth?: boolean;
  authType?: 'api-key' | 'basic' | 'unknown';
  models?: string[];
  error?: string;
};

const PI_PROVIDERS: { key: string; piId: string }[] = [
  { key: 'Anthropic', piId: 'anthropic' },
  { key: 'OpenAI', piId: 'openai' },
  { key: 'Google', piId: 'google' },
  { key: 'Groq', piId: 'groq' },
  { key: 'Mistral', piId: 'mistral' },
  { key: 'xAI', piId: 'xai' },
  { key: 'OpenRouter', piId: 'openrouter' },
  { key: 'MiniMax', piId: 'minimax' },
  { key: 'Hugging Face', piId: 'huggingface' },
  { key: 'Azure OpenAI', piId: 'azure-openai-responses' },
  { key: 'OpenCode Zen', piId: 'zai' },
  { key: 'Cerebras', piId: 'cerebras' },
];

type ProbeState =
  | { step: 'url' }
  | { step: 'probing' }
  | { step: 'auth'; probe: ProbeResult }
  | { step: 'saving' };

export const AIHarnessesSection = () => {
  const client = useClient();
  const queryClient = useQueryClient();
  const [installing, setInstalling] = useState(false);

  const { data: piMonoVersion, isLoading: piMonoLoading } = useQuery({
    queryKey: ['PI_MONO_VERSION'],
    queryFn: () => client.get<VersionInfo>('/server-settings/pi-mono/version'),
  });

  const { data: piMonoKeys } = useQuery({
    queryKey: ['PI_MONO_API_KEYS'],
    queryFn: () => client.get<StoredApiKeys>('/server-settings/pi-mono/api-keys'),
    enabled: !!piMonoVersion?.version,
  });

  const { data: localProviders = [] as LocalProviderEntry[] } = useQuery({
    queryKey: ['PI_MONO_LOCAL_PROVIDERS'],
    queryFn: () => client.get<LocalProviderEntry[]>('/server-settings/pi-mono/local-providers'),
    enabled: !!piMonoVersion?.version,
  });

  const { data: localHealth = {} as Record<string, boolean> } = useQuery({
    queryKey: ['PI_MONO_LOCAL_HEALTH'],
    queryFn: () => client.get<Record<string, boolean>>('/server-settings/pi-mono/local-providers/health'),
    enabled: localProviders.length > 0,
    refetchInterval: 15_000,
  });

  const [keyInputs, setKeyInputs] = useState<Record<string, string>>({});
  const [savingKey, setSavingKey] = useState<string | null>(null);
  const [commandOpen, setCommandOpen] = useState(false);
  const [editingProvider, setEditingProvider] = useState<string | null>(null);

  // Local provider connection flow
  const [addingLocal, setAddingLocal] = useState(false);
  const [localName, setLocalName] = useState('');
  const [localUrl, setLocalUrl] = useState('');
  const [probeState, setProbeState] = useState<ProbeState>({ step: 'url' });
  const [authApiKey, setAuthApiKey] = useState('');
  const [authUsername, setAuthUsername] = useState('');
  const [authPassword, setAuthPassword] = useState('');

  const resetLocalForm = () => {
    setAddingLocal(false);
    setLocalName('');
    setLocalUrl('');
    setProbeState({ step: 'url' });
    setAuthApiKey('');
    setAuthUsername('');
    setAuthPassword('');
  };

  const storedKeys: StoredApiKeys['keys'] = piMonoKeys?.keys ?? [];
  const storedPiIds = new Set(storedKeys.map((k) => k.provider));
  const connectedProviders = PI_PROVIDERS.filter((p) => storedPiIds.has(p.piId));
  const unconnectedProviders = PI_PROVIDERS.filter((p) => !storedPiIds.has(p.piId));

  const getStoredMasked = (piId: string) =>
    storedKeys.find((k) => k.provider === piId)?.value ?? '';

  const saveApiKey = async (piId: string) => {
    const value = keyInputs[piId];
    if (value === undefined) return;
    setSavingKey(piId);
    try {
      await client.put('/server-settings/pi-mono/api-keys', { provider: piId, value });
      queryClient.invalidateQueries({ queryKey: ['PI_MONO_API_KEYS'] });
      queryClient.invalidateQueries({ queryKey: ['PI_MODELS'] });
      setKeyInputs((prev) => {
        const next = { ...prev };
        delete next[piId];
        return next;
      });
    } finally {
      setSavingKey(null);
    }
  };

  const disconnectProvider = async (provider: typeof PI_PROVIDERS[number]) => {
    await client.put('/server-settings/pi-mono/api-keys', { provider: provider.piId, value: '' });
    queryClient.invalidateQueries({ queryKey: ['PI_MONO_API_KEYS'] });
    queryClient.invalidateQueries({ queryKey: ['PI_MODELS'] });
    if (editingProvider === provider.key) setEditingProvider(null);
  };

  const installPiMono = async () => {
    setInstalling(true);
    try {
      const result = await client.post<VersionInfo>('/server-settings/pi-mono/install');
      queryClient.setQueryData(['PI_MONO_VERSION'], result);
    } finally {
      setInstalling(false);
    }
  };

  const handleProbe = async (auth?: { type: 'api-key'; apiKey: string } | { type: 'basic'; username: string; password: string }) => {
    setProbeState({ step: 'probing' });
    try {
      const result = await client.post<ProbeResult>('/server-settings/pi-mono/local-providers/probe', {
        url: localUrl.trim(),
        auth,
      });
      if (!result.success) {
        toast.error(result.error ?? 'Could not detect API type');
        setProbeState({ step: 'url' });
        return;
      }
      if (result.needsAuth) {
        setProbeState({ step: 'auth', probe: result });
        return;
      }
      // No auth needed — save directly
      await saveLocalProvider(result, auth);
    } catch {
      toast.error('Failed to connect');
      setProbeState({ step: 'url' });
    }
  };

  const handleAuthSubmit = async (probe: ProbeResult) => {
    const auth = probe.authType === 'basic'
      ? { type: 'basic' as const, username: authUsername, password: authPassword }
      : { type: 'api-key' as const, apiKey: authApiKey };

    // Re-probe with credentials to verify they work
    setProbeState({ step: 'probing' });
    try {
      const result = await client.post<ProbeResult>('/server-settings/pi-mono/local-providers/probe', {
        url: localUrl.trim(),
        auth,
      });
      if (!result.success) {
        toast.error(result.error ?? 'Could not connect with provided credentials');
        setProbeState({ step: 'auth', probe });
        return;
      }
      if (result.needsAuth) {
        toast.error('Authentication failed');
        setProbeState({ step: 'auth', probe });
        return;
      }
      await saveLocalProvider(result, auth);
    } catch {
      toast.error('Failed to connect');
      setProbeState({ step: 'auth', probe });
    }
  };

  const saveLocalProvider = async (
    probe: ProbeResult,
    auth?: { type: 'api-key'; apiKey: string } | { type: 'basic'; username: string; password: string },
  ) => {
    setProbeState({ step: 'saving' });
    try {
      await client.post('/server-settings/pi-mono/local-providers', {
        url: localUrl.trim(),
        name: localName.trim() || probe.name,
        apiType: probe.apiType,
        auth,
      });
      queryClient.invalidateQueries({ queryKey: ['PI_MONO_LOCAL_PROVIDERS'] });
      queryClient.invalidateQueries({ queryKey: ['PI_MONO_LOCAL_HEALTH'] });
      queryClient.invalidateQueries({ queryKey: ['PI_MODELS'] });
      toast.success(`Connected to ${probe.name}`);
      resetLocalForm();
    } catch {
      toast.error('Failed to save provider');
      setProbeState({ step: 'url' });
    }
  };

  const removeLocalProvider = async (id: string) => {
    await client.delete(`/server-settings/pi-mono/local-providers/${id}`);
    queryClient.invalidateQueries({ queryKey: ['PI_MONO_LOCAL_PROVIDERS'] });
    queryClient.invalidateQueries({ queryKey: ['PI_MONO_LOCAL_HEALTH'] });
    queryClient.invalidateQueries({ queryKey: ['PI_MODELS'] });
  };

  return (
    <div className="flex flex-col gap-4">
      <div>
        <div className="flex items-center gap-3">
          <span className="text-sm font-medium text-duck-dark">Pi Mono</span>
          {piMonoLoading ? (
            <span className="text-xs text-duck-dark/50">Checking...</span>
          ) : piMonoVersion?.version ? (
            <div className="flex items-center gap-2">
              <span className="text-xs text-duck-dark/50">{piMonoVersion.version}</span>
              <button
                type="button"
                onClick={installPiMono}
                disabled={installing}
                className="p-0.5 rounded hover:bg-duck-teal/10 cursor-pointer transition-colors disabled:opacity-30"
                title="Update Pi Mono"
              >
                <RefreshCw className={`h-3 w-3 text-duck-teal ${installing ? 'animate-spin' : ''}`} />
              </button>
            </div>
          ) : (
            <Button
              size="sm"
              className="h-6 text-xs bg-duck-teal text-duck-yellow hover:bg-duck-teal/90"
              onClick={installPiMono}
              disabled={installing}
            >
              {installing ? 'Installing...' : 'Install'}
            </Button>
          )}
        </div>
        {piMonoVersion?.version && (
          <>
          {/* Local Providers */}
          <div className="mt-6 text-xs text-duck-dark/50">
            <span className="text-sm font-medium text-duck-dark dark:text-foreground">Local Providers</span>
            <div className="mt-3 flex flex-col gap-2">
              {localProviders.map((lp: LocalProviderEntry) => (
                <div key={lp.id} className="flex items-center gap-2">
                  <span
                    className={`inline-block h-2 w-2 shrink-0 rounded-full ${lp.id in localHealth ? (localHealth[lp.id] ? 'bg-green-500' : 'bg-red-500') : 'bg-duck-dark/20'}`}
                    title={lp.id in localHealth ? (localHealth[lp.id] ? 'Online' : 'Offline') : 'Checking...'}
                  />
                  <span className="text-duck-dark/70 dark:text-foreground/70 font-medium">{lp.name}</span>
                  <span className="text-duck-dark/40 dark:text-foreground/40 truncate">{lp.url}</span>
                  <button
                    type="button"
                    onClick={() => removeLocalProvider(lp.id)}
                    className="shrink-0 p-0.5 rounded hover:bg-red-500/10 cursor-pointer transition-colors"
                    title="Remove provider"
                  >
                    <Trash2 className="h-3 w-3 text-duck-dark/30 hover:text-red-500" />
                  </button>
                </div>
              ))}
              {addingLocal ? (
                <div className="flex flex-col gap-2 max-w-sm">
                  {/* Step 1: Name + URL inputs stacked */}
                  <Input
                    type="text"
                    className="h-8 text-xs"
                    placeholder="Name (e.g. My Ollama)"
                    value={localName}
                    onChange={(ev) => setLocalName(ev.target.value)}
                    onKeyDown={(ev) => {
                      if (ev.key === 'Escape') resetLocalForm();
                    }}
                    disabled={probeState.step !== 'url'}
                    autoFocus
                  />
                  <Input
                    type="url"
                    className="h-8 text-xs"
                    placeholder="http://localhost:11434"
                    value={localUrl}
                    onChange={(ev) => setLocalUrl(ev.target.value)}
                    onKeyDown={(ev) => {
                      if (ev.key === 'Escape') resetLocalForm();
                      if (ev.key === 'Enter' && localUrl.trim() && probeState.step === 'url') handleProbe();
                    }}
                    disabled={probeState.step !== 'url'}
                  />
                  <div className="flex items-center gap-2">
                    {probeState.step === 'url' && (
                      <Button
                        type="button"
                        size="sm"
                        className="h-7 text-xs bg-duck-dark text-white hover:bg-duck-dark/90 dark:bg-foreground dark:text-background dark:hover:bg-foreground/90 cursor-pointer disabled:opacity-40"
                        disabled={!localUrl.trim()}
                        onClick={() => handleProbe()}
                      >
                        Connect
                      </Button>
                    )}
                    {(probeState.step === 'probing' || probeState.step === 'saving') && (
                      <Loader2 className="h-4 w-4 animate-spin text-duck-teal shrink-0" />
                    )}
                    <button
                      type="button"
                      onClick={resetLocalForm}
                      className="text-duck-dark/40 hover:text-duck-dark/70 dark:text-foreground/40 dark:hover:text-foreground/70 cursor-pointer transition-colors text-xs"
                    >
                      Cancel
                    </button>
                  </div>

                  {/* Step 2: Auth form (if needed) */}
                  {probeState.step === 'auth' && (
                    <div className="flex flex-col gap-2 rounded-md border border-duck-dark/10 dark:border-foreground/10 p-3">
                      <span className="text-duck-dark/70 dark:text-foreground/70 font-medium">
                        {probeState.probe.name} requires authentication
                      </span>
                      {probeState.probe.authType === 'basic' ? (
                        <>
                          <Input
                            type="text"
                            className="h-8 text-xs"
                            placeholder="Username"
                            value={authUsername}
                            onChange={(ev) => setAuthUsername(ev.target.value)}
                            autoFocus
                          />
                          <Input
                            type="password"
                            className="h-8 text-xs"
                            placeholder="Password"
                            value={authPassword}
                            onChange={(ev) => setAuthPassword(ev.target.value)}
                            onKeyDown={(ev) => {
                              if (ev.key === 'Enter' && authUsername && authPassword) handleAuthSubmit(probeState.probe);
                            }}
                          />
                          <Button
                            type="button"
                            size="sm"
                            className="h-7 text-xs w-fit bg-duck-dark text-white hover:bg-duck-dark/90 dark:bg-foreground dark:text-background dark:hover:bg-foreground/90 cursor-pointer disabled:opacity-40"
                            disabled={!authUsername || !authPassword}
                            onClick={() => handleAuthSubmit(probeState.probe)}
                          >
                            Authenticate
                          </Button>
                        </>
                      ) : (
                        <>
                          <Input
                            type="password"
                            className="h-8 text-xs"
                            placeholder="API Key"
                            value={authApiKey}
                            onChange={(ev) => setAuthApiKey(ev.target.value)}
                            onKeyDown={(ev) => {
                              if (ev.key === 'Enter' && authApiKey) handleAuthSubmit(probeState.probe);
                            }}
                            autoFocus
                          />
                          <Button
                            type="button"
                            size="sm"
                            className="h-7 text-xs w-fit bg-duck-dark text-white hover:bg-duck-dark/90 dark:bg-foreground dark:text-background dark:hover:bg-foreground/90 cursor-pointer disabled:opacity-40"
                            disabled={!authApiKey}
                            onClick={() => handleAuthSubmit(probeState.probe)}
                          >
                            Authenticate
                          </Button>
                        </>
                      )}
                    </div>
                  )}
                </div>
              ) : (
                <Button
                  type="button"
                  variant="outline"
                  size="sm"
                  className="w-fit cursor-pointer"
                  onClick={() => setAddingLocal(true)}
                >
                  <Plus className="h-3.5 w-3.5 mr-1.5" />
                  Connect Local Provider
                </Button>
              )}
            </div>
          </div>

          {/* Remote Providers */}
          <div className="mt-6 text-xs text-duck-dark/50">
            <span className="text-sm font-medium text-duck-dark dark:text-foreground">Remote Providers</span>
            <div className="mt-3 flex flex-col gap-2">
              {connectedProviders.map((provider) => (
                <div key={provider.piId} className="flex items-center gap-2">
                  <label className="w-36 text-duck-dark/70 dark:text-foreground/70 shrink-0 truncate font-medium text-[11px]">{provider.key}</label>
                  <Input
                    type="password"
                    className="h-7 text-xs flex-1"
                    placeholder={getStoredMasked(provider.piId) || 'Not set'}
                    value={keyInputs[provider.piId] ?? ''}
                    onChange={(ev) => setKeyInputs((prev) => ({ ...prev, [provider.piId]: ev.target.value }))}
                  />
                  <button
                    type="button"
                    disabled={keyInputs[provider.piId] === undefined || savingKey === provider.piId}
                    onClick={() => saveApiKey(provider.piId)}
                    className="shrink-0 p-1 rounded hover:bg-duck-teal/10 cursor-pointer transition-colors disabled:opacity-30 disabled:cursor-default"
                    title="Save key"
                  >
                    <Save className="h-3.5 w-3.5 text-duck-teal" />
                  </button>
                  <button
                    type="button"
                    onClick={() => disconnectProvider(provider)}
                    className="shrink-0 p-0.5 rounded hover:bg-red-500/10 cursor-pointer transition-colors"
                    title="Disconnect provider"
                  >
                    <Trash2 className="h-3 w-3 text-duck-dark/30 hover:text-red-500" />
                  </button>
                </div>
              ))}
              {editingProvider && (() => {
                const provider = PI_PROVIDERS.find((p) => p.key === editingProvider);
                if (!provider || connectedProviders.includes(provider)) return null;
                return (
                  <div key={provider.piId} className="flex items-center gap-2">
                    <label className="w-36 text-duck-dark/70 dark:text-foreground/70 shrink-0 truncate font-medium text-[11px]">{provider.key}</label>
                    <Input
                      type="password"
                      className="h-7 text-xs flex-1"
                      placeholder="Enter API key"
                      value={keyInputs[provider.piId] ?? ''}
                      onChange={(ev) => setKeyInputs((prev) => ({ ...prev, [provider.piId]: ev.target.value }))}
                      autoFocus
                    />
                    <button
                      type="button"
                      disabled={keyInputs[provider.piId] === undefined || savingKey === provider.piId}
                      onClick={() => saveApiKey(provider.piId)}
                      className="shrink-0 p-1 rounded hover:bg-duck-teal/10 cursor-pointer transition-colors disabled:opacity-30 disabled:cursor-default"
                      title="Save key"
                    >
                      <Save className="h-3.5 w-3.5 text-duck-teal" />
                    </button>
                  </div>
                );
              })()}
              {unconnectedProviders.length > 0 && (
                <Button
                  type="button"
                  variant="outline"
                  size="sm"
                  className="mt-1 w-fit cursor-pointer"
                  onClick={() => setCommandOpen(true)}
                >
                  <Plus className="h-3.5 w-3.5 mr-1.5" />
                  Connect Provider
                </Button>
              )}
            </div>
            <CommandDialog open={commandOpen} onOpenChange={setCommandOpen}>
              <CommandInput placeholder="Search providers..." />
              <CommandList>
                <CommandEmpty>No providers found.</CommandEmpty>
                <CommandGroup heading="Available Providers">
                  {unconnectedProviders.map((provider) => (
                    <CommandItem
                      key={provider.key}
                      onSelect={() => {
                        setEditingProvider(provider.key);
                        setCommandOpen(false);
                      }}
                    >
                      <span>{provider.key}</span>
                    </CommandItem>
                  ))}
                </CommandGroup>
              </CommandList>
            </CommandDialog>
          </div>
          </>
        )}
      </div>
    </div>
  );
};