stable-diffusion.cpp-rest/webui/components/enhanced-model-select.tsx

'use client';

import React, { useState } from 'react';
import { Label } from '@/components/ui/label';
import { Button } from '@/components/ui/button';
import { Badge } from '@/components/ui/badge';
import { 
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from '@/components/ui/select';
import { 
  CheckCircle2, 
  Zap, 
  RotateCcw, 
  AlertCircle,
  Info,
  Loader2
} from 'lucide-react';
import { cn } from '@/lib/utils';
import { ModelInfo } from '@/lib/api';
import { ModelSelectionIndicator, AutoSelectionStatus } from './model-selection-indicator';

interface EnhancedModelSelectProps {
  modelType: string;
  label: string;
  description?: string;
  value: string | null;
  availableModels: ModelInfo[];
  isAutoSelected: boolean;
  isUserOverride: boolean;
  isLoading?: boolean;
  onValueChange: (value: string) => void;
  onSetUserOverride: (value: string) => void;
  onClearOverride: () => void;
  onRevertToAuto?: () => void;
  placeholder?: string;
  className?: string;
  disabled?: boolean;
}

export function EnhancedModelSelect({
  modelType,
  label,
  description,
  value,
  availableModels,
  isAutoSelected,
  isUserOverride,
  isLoading = false,
  onValueChange,
  onSetUserOverride,
  onClearOverride,
  onRevertToAuto,
  placeholder = "Select a model...",
  className,
  disabled = false,
}: EnhancedModelSelectProps) {
  const [isOpen, setIsOpen] = useState(false);

  const selectedModel = value ? availableModels.find(m => m.name === value) : null;
  const isLoaded = selectedModel?.loaded || false;

  const handleValueChange = (newValue: string) => {
    if (newValue !== value) {
      onValueChange(newValue);
      onSetUserOverride(newValue);
    }
  };

  const handleClearOverride = () => {
    onClearOverride();
  };

  const handleRevertToAuto = () => {
    if (onRevertToAuto) {
      onRevertToAuto();
    }
  };

  const getModelIcon = (type: string) => {
    switch (type.toLowerCase()) {
      case 'vae':
        return <Zap className="h-4 w-4" />;
      case 'checkpoint':
      case 'stable-diffusion':
        return <CheckCircle2 className="h-4 w-4" />;
      default:
        return <Info className="h-4 w-4" />;
    }
  };

  const getModelStatusColor = (model: ModelInfo) => {
    if (model.loaded) {
      return 'text-green-600 dark:text-green-400';
    }
    return 'text-muted-foreground';
  };

  return (
    <div className={cn("space-y-2", className)}>
      <div className="flex items-center justify-between">
        <Label htmlFor={`${modelType}-select`} className="text-sm font-medium">
          {label}
        </Label>
        
        <ModelSelectionIndicator
          modelName={value}
          isAutoSelected={isAutoSelected}
          isUserOverride={isUserOverride}
          isLoaded={isLoaded}
          onClearOverride={isUserOverride ? handleClearOverride : undefined}
          onRevertToAuto={isUserOverride && onRevertToAuto ? handleRevertToAuto : undefined}
        />
      </div>

      {description && (
        <p className="text-xs text-muted-foreground">{description}</p>
      )}

      <div className="relative">
        <Select
          value={value || ''}
          onValueChange={handleValueChange}
          disabled={disabled || isLoading}
          open={isOpen}
          onOpenChange={setIsOpen}
        >
          <SelectTrigger
            id={`${modelType}-select`}
            className={cn(
              "w-full",
              isAutoSelected && !isUserOverride && "border-green-500 dark:border-green-600",
              isUserOverride && "border-blue-500 dark:border-blue-600"
            )}
          >
            <div className="flex items-center justify-between w-full">
              <SelectValue placeholder={placeholder} />
              {isLoading && (
                <Loader2 className="h-4 w-4 animate-spin ml-2" />
              )}
            </div>
          </SelectTrigger>
          
          <SelectContent>
            {availableModels.length === 0 ? (
              <div className="p-2 text-sm text-muted-foreground text-center">
                No {modelType} models available
              </div>
            ) : (
              <>
                {availableModels.map((model) => (
                  <SelectItem key={model.id || model.name} value={model.name}>
                    <div className="flex items-center justify-between w-full">
                      <div className="flex items-center gap-2">
                        {getModelIcon(model.type)}
                        <span className={cn(getModelStatusColor(model))}>
                          {model.name}
                        </span>
                      </div>
                      
                      <div className="flex items-center gap-2">
                        {model.loaded && (
                          <Badge variant="secondary" className="text-xs bg-green-100 text-green-800 dark:bg-green-900 dark:text-green-200">
                            Loaded
                          </Badge>
                        )}
                        
                        {model.file_size_mb && (
                          <span className="text-xs text-muted-foreground">
                            {model.file_size_mb.toFixed(1)} MB
                          </span>
                        )}
                      </div>
                    </div>
                  </SelectItem>
                ))}
              </>
            )}
          </SelectContent>
        </Select>

        {/* Auto-selection indicator */}
        {isAutoSelected && !isUserOverride && (
          <div className="absolute -top-1 -right-1">
            <div className="bg-green-500 rounded-full p-1">
              <Zap className="h-3 w-3 text-white" />
            </div>
          </div>
        )}

        {/* User override indicator */}
        {isUserOverride && (
          <div className="absolute -top-1 -right-1">
            <div className="bg-blue-500 rounded-full p-1">
              <CheckCircle2 className="h-3 w-3 text-white" />
            </div>
          </div>
        )}
      </div>

      {/* Model info display */}
      {selectedModel && (
        <div className="p-2 bg-muted/50 rounded-md">
          <div className="flex items-center justify-between text-xs">
            <div className="flex items-center gap-2">
              <span className="text-muted-foreground">Type:</span>
              <Badge variant="outline" className="text-xs">
                {selectedModel.type}
              </Badge>
            </div>
            
            {selectedModel.file_size_mb && (
              <div className="flex items-center gap-2">
                <span className="text-muted-foreground">Size:</span>
                <span>{selectedModel.file_size_mb.toFixed(1)} MB</span>
              </div>
            )}
          </div>
          
          {selectedModel.architecture && (
            <div className="flex items-center gap-2 text-xs mt-1">
              <span className="text-muted-foreground">Architecture:</span>
              <Badge variant="secondary" className="text-xs">
                {selectedModel.architecture}
              </Badge>
            </div>
          )}
        </div>
      )}

      {/* No models warning */}
      {availableModels.length === 0 && !isLoading && (
        <div className="flex items-center gap-2 p-2 rounded-md bg-yellow-50 dark:bg-yellow-950/20 border border-yellow-200 dark:border-yellow-800">
          <AlertCircle className="h-4 w-4 text-yellow-500" />
          <p className="text-sm text-yellow-700 dark:text-yellow-300">
            No {modelType} models found. Please add {modelType} models to your models directory.
          </p>
        </div>
      )}
    </div>
  );
}

interface EnhancedModelSelectGroupProps {
  title: string;
  description?: string;
  children: React.ReactNode;
  isLoading?: boolean;
  className?: string;
}

export function EnhancedModelSelectGroup({
  title,
  description,
  children,
  isLoading = false,
  className,
}: EnhancedModelSelectGroupProps) {
  return (
    <div className={cn("space-y-4", className)}>
      <div className="space-y-2">
        <h3 className="text-lg font-semibold">{title}</h3>
        {description && (
          <p className="text-sm text-muted-foreground">{description}</p>
        )}
      </div>
      
      {isLoading ? (
        <div className="flex items-center justify-center py-8">
          <Loader2 className="h-6 w-6 animate-spin mr-2" />
          <span>Loading models...</span>
        </div>
      ) : (
        <div className="space-y-4">
          {children}
        </div>
      )}
    </div>
  );
}