Atomizer/atomizer-dashboard/frontend/src/components/dashboard/ControlPanel.tsx

import React, { useState, useEffect } from 'react';
import {
  Play,
  Pause,
  CheckCircle,
  Settings,
  AlertTriangle,
  Loader2,
  ExternalLink,
  Sliders,
  Skull,
  Info
} from 'lucide-react';
import { apiClient, ProcessStatus } from '../../api/client';
import { useStudy } from '../../context/StudyContext';

interface ControlPanelProps {
  onStatusChange?: () => void;
  horizontal?: boolean;
}

export const ControlPanel: React.FC<ControlPanelProps> = ({ onStatusChange, horizontal = false }) => {
  const { selectedStudy, refreshStudies } = useStudy();
  const [processStatus, setProcessStatus] = useState<ProcessStatus | null>(null);
  const [actionInProgress, setActionInProgress] = useState<string | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [showSettings, setShowSettings] = useState(false);
  const [optunaAvailable, setOptunaAvailable] = useState<boolean | null>(null);
  const [optunaInstallHint, setOptunaInstallHint] = useState<string | null>(null);

  // Settings for starting optimization
  const [settings, setSettings] = useState({
    freshStart: false,
    maxIterations: 100,
    trials: 300,  // For SAT scripts
    feaBatchSize: 5,
    tuneTrials: 30,
    ensembleSize: 3,
    patience: 5,
  });

  // Validate top N
  const [validateTopN, setValidateTopN] = useState(5);

  useEffect(() => {
    if (selectedStudy) {
      fetchProcessStatus();
      const interval = setInterval(fetchProcessStatus, 5000);
      return () => clearInterval(interval);
    }
  }, [selectedStudy]);

  // Check optuna-dashboard availability on mount
  useEffect(() => {
    const checkOptuna = async () => {
      try {
        const result = await apiClient.checkOptunaAvailable();
        setOptunaAvailable(result.available);
        if (!result.available && result.install_instructions) {
          setOptunaInstallHint(result.install_instructions);
        }
      } catch {
        setOptunaAvailable(false);
        setOptunaInstallHint('pip install optuna-dashboard');
      }
    };
    checkOptuna();
  }, []);

  const fetchProcessStatus = async () => {
    if (!selectedStudy) return;
    try {
      const status = await apiClient.getProcessStatus(selectedStudy.id);
      setProcessStatus(status);
    } catch (err) {
      // Process status endpoint might not exist yet
      setProcessStatus(null);
    }
  };

  const handleStart = async () => {
    if (!selectedStudy) return;
    setActionInProgress('start');
    setError(null);

    try {
      await apiClient.startOptimization(selectedStudy.id, {
        freshStart: settings.freshStart,
        maxIterations: settings.maxIterations,
        trials: settings.trials,
        feaBatchSize: settings.feaBatchSize,
        tuneTrials: settings.tuneTrials,
        ensembleSize: settings.ensembleSize,
        patience: settings.patience,
      });

      await fetchProcessStatus();
      await refreshStudies();
      onStatusChange?.();
    } catch (err: any) {
      setError(err.message || 'Failed to start optimization');
    } finally {
      setActionInProgress(null);
    }
  };

  const handleStop = async () => {
    if (!selectedStudy) return;
    setActionInProgress('stop');
    setError(null);

    try {
      await apiClient.stopOptimization(selectedStudy.id);
      await fetchProcessStatus();
      await refreshStudies();
      onStatusChange?.();
    } catch (err: any) {
      setError(err.message || 'Failed to stop optimization');
    } finally {
      setActionInProgress(null);
    }
  };

  const handlePause = async () => {
    if (!selectedStudy) return;
    setActionInProgress('pause');
    setError(null);

    try {
      await apiClient.pauseOptimization(selectedStudy.id);
      await fetchProcessStatus();
      onStatusChange?.();
    } catch (err: any) {
      setError(err.message || 'Failed to pause optimization');
    } finally {
      setActionInProgress(null);
    }
  };

  const handleResume = async () => {
    if (!selectedStudy) return;
    setActionInProgress('resume');
    setError(null);

    try {
      await apiClient.resumeOptimization(selectedStudy.id);
      await fetchProcessStatus();
      onStatusChange?.();
    } catch (err: any) {
      setError(err.message || 'Failed to resume optimization');
    } finally {
      setActionInProgress(null);
    }
  };

  const handleValidate = async () => {
    if (!selectedStudy) return;
    setActionInProgress('validate');
    setError(null);

    try {
      await apiClient.validateOptimization(selectedStudy.id, { topN: validateTopN });
      await fetchProcessStatus();
      await refreshStudies();
      onStatusChange?.();
    } catch (err: any) {
      setError(err.message || 'Failed to start validation');
    } finally {
      setActionInProgress(null);
    }
  };

  const handleLaunchOptuna = async () => {
    if (!selectedStudy) return;
    setActionInProgress('optuna');
    setError(null);

    try {
      const result = await apiClient.launchOptunaDashboard(selectedStudy.id);
      window.open(result.url, '_blank');
    } catch (err: any) {
      setError(err.message || 'Failed to launch Optuna dashboard');
    } finally {
      setActionInProgress(null);
    }
  };

  const isRunning = processStatus?.is_running || selectedStudy?.status === 'running';
  const isPaused = processStatus?.is_paused || false;

  // Horizontal layout for top of page
  if (horizontal) {
    return (
      <div className="bg-dark-800 rounded-xl border border-dark-600 overflow-hidden">
        <div className="px-4 py-3 flex items-center gap-6">
          {/* Status */}
          <div className="flex items-center gap-3">
            <div className="flex items-center gap-2">
              {isRunning && isPaused ? (
                <>
                  <div className="w-3 h-3 bg-yellow-500 rounded-full" />
                  <span className="text-yellow-400 font-medium text-sm">Paused</span>
                </>
              ) : isRunning ? (
                <>
                  <div className="w-3 h-3 bg-green-500 rounded-full animate-pulse" />
                  <span className="text-green-400 font-medium text-sm">Running</span>
                </>
              ) : (
                <>
                  <div className="w-3 h-3 bg-dark-500 rounded-full" />
                  <span className="text-dark-400 text-sm">Stopped</span>
                </>
              )}
            </div>
            {/* Trial counts */}
            {processStatus?.completed_trials !== undefined && (
              <span className="text-xs text-dark-400">
                Trials: <span className="text-white font-mono">{processStatus.completed_trials}</span>
                {processStatus.total_trials && (
                  <span className="text-dark-500">/{processStatus.total_trials}</span>
                )}
              </span>
            )}
            {/* ETA and Rate */}
            {isRunning && processStatus?.eta_formatted && (
              <span className="text-xs text-dark-400">
                ETA: <span className="text-primary-400 font-mono">{processStatus.eta_formatted}</span>
              </span>
            )}
            {processStatus?.rate_per_hour && (
              <span className="text-xs text-dark-400">
                Rate: <span className="text-dark-300 font-mono">{processStatus.rate_per_hour.toFixed(1)}/hr</span>
              </span>
            )}
          </div>

          {/* Main Actions */}
          <div className="flex items-center gap-2">
            {isRunning ? (
              <>
                {/* Pause/Resume Button */}
                {isPaused ? (
                  <button
                    onClick={handleResume}
                    disabled={actionInProgress !== null}
                    className="flex items-center gap-2 px-3 py-1.5 bg-green-600 hover:bg-green-500
                               disabled:opacity-50 text-white rounded-lg text-sm font-medium"
                  >
                    {actionInProgress === 'resume' ? <Loader2 className="w-4 h-4 animate-spin" /> : <Play className="w-4 h-4" />}
                    Resume
                  </button>
                ) : (
                  <button
                    onClick={handlePause}
                    disabled={actionInProgress !== null}
                    className="flex items-center gap-2 px-3 py-1.5 bg-yellow-600 hover:bg-yellow-500
                               disabled:opacity-50 text-white rounded-lg text-sm font-medium"
                  >
                    {actionInProgress === 'pause' ? <Loader2 className="w-4 h-4 animate-spin" /> : <Pause className="w-4 h-4" />}
                    Pause
                  </button>
                )}
                {/* Kill Button */}
                <button
                  onClick={handleStop}
                  disabled={actionInProgress !== null}
                  className="flex items-center gap-2 px-3 py-1.5 bg-red-600 hover:bg-red-500
                             disabled:opacity-50 text-white rounded-lg text-sm font-medium"
                >
                  {actionInProgress === 'stop' ? <Loader2 className="w-4 h-4 animate-spin" /> : <Skull className="w-4 h-4" />}
                  Kill
                </button>
              </>
            ) : (
              <button
                onClick={handleStart}
                disabled={actionInProgress !== null}
                className="flex items-center gap-2 px-3 py-1.5 bg-green-600 hover:bg-green-500
                           disabled:opacity-50 text-white rounded-lg text-sm font-medium"
              >
                {actionInProgress === 'start' ? <Loader2 className="w-4 h-4 animate-spin" /> : <Play className="w-4 h-4" />}
                Start
              </button>
            )}

            <button
              onClick={handleValidate}
              disabled={actionInProgress !== null || isRunning}
              className="flex items-center gap-2 px-3 py-1.5 bg-primary-600 hover:bg-primary-500
                         disabled:opacity-50 text-white rounded-lg text-sm font-medium"
            >
              {actionInProgress === 'validate' ? <Loader2 className="w-4 h-4 animate-spin" /> : <CheckCircle className="w-4 h-4" />}
              Validate
            </button>

            <div className="relative group">
              <button
                onClick={handleLaunchOptuna}
                disabled={actionInProgress !== null || optunaAvailable === false}
                className={`flex items-center gap-2 px-3 py-1.5 bg-dark-700 hover:bg-dark-600
                           border border-dark-600 disabled:opacity-50 text-dark-300 hover:text-white rounded-lg text-sm
                           ${optunaAvailable === false ? 'cursor-not-allowed' : ''}`}
              >
                <ExternalLink className="w-4 h-4" />
                Optuna
              </button>
              {optunaAvailable === false && optunaInstallHint && (
                <div className="absolute bottom-full left-0 mb-2 px-2 py-1 bg-dark-900 border border-dark-600
                                rounded text-xs text-dark-400 whitespace-nowrap opacity-0 group-hover:opacity-100
                                transition-opacity pointer-events-none z-10">
                  <Info className="w-3 h-3 inline mr-1" />
                  Install with: {optunaInstallHint}
                </div>
              )}
            </div>
          </div>

          {/* Settings Toggle */}
          <button
            onClick={() => setShowSettings(!showSettings)}
            className={`p-1.5 rounded-lg transition-colors ${
              showSettings ? 'bg-primary-600 text-white' : 'bg-dark-700 text-dark-300 hover:text-white'
            }`}
          >
            <Settings className="w-4 h-4" />
          </button>

          {/* Error */}
          {error && (
            <div className="flex items-center gap-2 text-red-400 text-sm">
              <AlertTriangle className="w-4 h-4" />
              {error}
            </div>
          )}
        </div>

        {/* Collapsible Settings */}
        {showSettings && (
          <div className="px-4 py-3 border-t border-dark-700 bg-dark-750">
            <div className="flex items-center gap-4 flex-wrap">
              <div className="flex items-center gap-2">
                <label className="text-xs text-dark-400">Trials (SAT):</label>
                <input
                  type="number"
                  value={settings.trials}
                  onChange={(e) => setSettings({ ...settings, trials: parseInt(e.target.value) || 300 })}
                  className="w-16 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm"
                />
              </div>
              <div className="flex items-center gap-2">
                <label className="text-xs text-dark-400">FEA Batch:</label>
                <input
                  type="number"
                  value={settings.feaBatchSize}
                  onChange={(e) => setSettings({ ...settings, feaBatchSize: parseInt(e.target.value) || 5 })}
                  className="w-12 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm"
                />
              </div>
              <div className="flex items-center gap-2">
                <label className="text-xs text-dark-400">Patience:</label>
                <input
                  type="number"
                  value={settings.patience}
                  onChange={(e) => setSettings({ ...settings, patience: parseInt(e.target.value) || 5 })}
                  className="w-12 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm"
                />
              </div>
              <div className="flex items-center gap-2">
                <label className="text-xs text-dark-400">Tune:</label>
                <input
                  type="number"
                  value={settings.tuneTrials}
                  onChange={(e) => setSettings({ ...settings, tuneTrials: parseInt(e.target.value) || 30 })}
                  className="w-12 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm"
                />
              </div>
              <div className="flex items-center gap-2">
                <label className="text-xs text-dark-400">Ensemble:</label>
                <input
                  type="number"
                  value={settings.ensembleSize}
                  onChange={(e) => setSettings({ ...settings, ensembleSize: parseInt(e.target.value) || 3 })}
                  className="w-12 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm"
                />
              </div>
              <div className="flex items-center gap-2">
                <label className="text-xs text-dark-400">Validate Top:</label>
                <input
                  type="number"
                  min={1}
                  max={20}
                  value={validateTopN}
                  onChange={(e) => setValidateTopN(parseInt(e.target.value) || 5)}
                  className="w-12 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm"
                />
              </div>
              <label className="flex items-center gap-2 cursor-pointer">
                <input
                  type="checkbox"
                  checked={settings.freshStart}
                  onChange={(e) => setSettings({ ...settings, freshStart: e.target.checked })}
                  className="w-4 h-4 rounded border-dark-600 bg-dark-700 text-primary-600"
                />
                <span className="text-xs text-dark-300">Fresh Start</span>
              </label>
            </div>
          </div>
        )}
      </div>
    );
  }

  // Vertical layout (original sidebar layout)
  return (
    <div className="bg-dark-800 rounded-xl border border-dark-600 overflow-hidden">
      {/* Header */}
      <div className="px-6 py-4 border-b border-dark-600 flex items-center justify-between">
        <h2 className="text-lg font-semibold text-white flex items-center gap-2">
          <Sliders className="w-5 h-5 text-primary-400" />
          Optimization Control
        </h2>
        <button
          onClick={() => setShowSettings(!showSettings)}
          className={`p-2 rounded-lg transition-colors ${
            showSettings ? 'bg-primary-600 text-white' : 'bg-dark-700 text-dark-300 hover:text-white'
          }`}
        >
          <Settings className="w-4 h-4" />
        </button>
      </div>

      {/* Status */}
      <div className="px-6 py-4 border-b border-dark-700">
        <div className="flex items-center justify-between">
          <div>
            <div className="text-sm text-dark-400 mb-1">Status</div>
            <div className="flex items-center gap-2">
              {isRunning && isPaused ? (
                <>
                  <div className="w-3 h-3 bg-yellow-500 rounded-full" />
                  <span className="text-yellow-400 font-medium">Paused</span>
                </>
              ) : isRunning ? (
                <>
                  <div className="w-3 h-3 bg-green-500 rounded-full animate-pulse" />
                  <span className="text-green-400 font-medium">Running</span>
                </>
              ) : (
                <>
                  <div className="w-3 h-3 bg-dark-500 rounded-full" />
                  <span className="text-dark-400">Stopped</span>
                </>
              )}
            </div>
          </div>

          {processStatus && (
            <div className="text-right">
              {processStatus.completed_trials !== undefined && (
                <div className="text-sm text-dark-400">
                  Trials: <span className="text-white font-mono">{processStatus.completed_trials}</span>
                  {processStatus.total_trials && (
                    <span className="text-dark-500">/{processStatus.total_trials}</span>
                  )}
                </div>
              )}
              {isRunning && !isPaused && processStatus.eta_formatted && (
                <div className="text-sm text-dark-400">
                  ETA: <span className="text-primary-400 font-mono">{processStatus.eta_formatted}</span>
                </div>
              )}
              {processStatus.rate_per_hour && (
                <div className="text-sm text-dark-400">
                  Rate: <span className="text-dark-300 font-mono">{processStatus.rate_per_hour.toFixed(1)}/hr</span>
                </div>
              )}
            </div>
          )}
        </div>
      </div>

      {/* Settings Panel */}
      {showSettings && (
        <div className="px-6 py-4 border-b border-dark-700 bg-dark-750">
          <div className="grid grid-cols-3 gap-4">
            <div>
              <label className="block text-xs text-dark-400 mb-1">Trials (SAT)</label>
              <input
                type="number"
                value={settings.trials}
                onChange={(e) => setSettings({ ...settings, trials: parseInt(e.target.value) || 300 })}
                className="w-full px-3 py-2 bg-dark-700 border border-dark-600 rounded-lg text-white text-sm"
              />
            </div>
            <div>
              <label className="block text-xs text-dark-400 mb-1">FEA Batch Size</label>
              <input
                type="number"
                value={settings.feaBatchSize}
                onChange={(e) => setSettings({ ...settings, feaBatchSize: parseInt(e.target.value) || 5 })}
                className="w-full px-3 py-2 bg-dark-700 border border-dark-600 rounded-lg text-white text-sm"
              />
            </div>
            <div>
              <label className="block text-xs text-dark-400 mb-1">Patience</label>
              <input
                type="number"
                value={settings.patience}
                onChange={(e) => setSettings({ ...settings, patience: parseInt(e.target.value) || 5 })}
                className="w-full px-3 py-2 bg-dark-700 border border-dark-600 rounded-lg text-white text-sm"
              />
            </div>
            <div>
              <label className="block text-xs text-dark-400 mb-1">Tuning Trials</label>
              <input
                type="number"
                value={settings.tuneTrials}
                onChange={(e) => setSettings({ ...settings, tuneTrials: parseInt(e.target.value) || 30 })}
                className="w-full px-3 py-2 bg-dark-700 border border-dark-600 rounded-lg text-white text-sm"
              />
            </div>
            <div>
              <label className="block text-xs text-dark-400 mb-1">Ensemble Size</label>
              <input
                type="number"
                value={settings.ensembleSize}
                onChange={(e) => setSettings({ ...settings, ensembleSize: parseInt(e.target.value) || 3 })}
                className="w-full px-3 py-2 bg-dark-700 border border-dark-600 rounded-lg text-white text-sm"
              />
            </div>
            <div className="flex items-end">
              <label className="flex items-center gap-2 cursor-pointer">
                <input
                  type="checkbox"
                  checked={settings.freshStart}
                  onChange={(e) => setSettings({ ...settings, freshStart: e.target.checked })}
                  className="w-4 h-4 rounded border-dark-600 bg-dark-700 text-primary-600"
                />
                <span className="text-sm text-dark-300">Fresh Start</span>
              </label>
            </div>
          </div>
        </div>
      )}

      {/* Error Message */}
      {error && (
        <div className="px-6 py-3 bg-red-900/20 border-b border-red-800/30">
          <div className="flex items-center gap-2 text-red-400 text-sm">
            <AlertTriangle className="w-4 h-4" />
            {error}
          </div>
        </div>
      )}

      {/* Actions */}
      <div className="p-6">
        <div className="grid grid-cols-2 gap-3">
          {/* Start / Pause / Resume / Kill Buttons */}
          {isRunning ? (
            <>
              {/* Pause/Resume Button */}
              {isPaused ? (
                <button
                  onClick={handleResume}
                  disabled={actionInProgress !== null}
                  className="flex items-center justify-center gap-2 px-4 py-3 bg-green-600 hover:bg-green-500
                             disabled:opacity-50 disabled:cursor-not-allowed
                             text-white rounded-lg transition-colors font-medium"
                >
                  {actionInProgress === 'resume' ? (
                    <Loader2 className="w-5 h-5 animate-spin" />
                  ) : (
                    <Play className="w-5 h-5" />
                  )}
                  Resume
                </button>
              ) : (
                <button
                  onClick={handlePause}
                  disabled={actionInProgress !== null}
                  className="flex items-center justify-center gap-2 px-4 py-3 bg-yellow-600 hover:bg-yellow-500
                             disabled:opacity-50 disabled:cursor-not-allowed
                             text-white rounded-lg transition-colors font-medium"
                >
                  {actionInProgress === 'pause' ? (
                    <Loader2 className="w-5 h-5 animate-spin" />
                  ) : (
                    <Pause className="w-5 h-5" />
                  )}
                  Pause
                </button>
              )}
              {/* Kill Button */}
              <button
                onClick={handleStop}
                disabled={actionInProgress !== null}
                className="flex items-center justify-center gap-2 px-4 py-3 bg-red-600 hover:bg-red-500
                           disabled:opacity-50 disabled:cursor-not-allowed
                           text-white rounded-lg transition-colors font-medium"
                title="Force kill the optimization process and all child processes"
              >
                {actionInProgress === 'stop' ? (
                  <Loader2 className="w-5 h-5 animate-spin" />
                ) : (
                  <Skull className="w-5 h-5" />
                )}
                Kill Process
              </button>
            </>
          ) : (
            <>
              <button
                onClick={handleStart}
                disabled={actionInProgress !== null}
                className="flex items-center justify-center gap-2 px-4 py-3 bg-green-600 hover:bg-green-500
                           disabled:opacity-50 disabled:cursor-not-allowed
                           text-white rounded-lg transition-colors font-medium"
              >
                {actionInProgress === 'start' ? (
                  <Loader2 className="w-5 h-5 animate-spin" />
                ) : (
                  <Play className="w-5 h-5" />
                )}
                Start Optimization
              </button>

              {/* Validate Button */}
              <button
                onClick={handleValidate}
                disabled={actionInProgress !== null}
                className="flex items-center justify-center gap-2 px-4 py-3 bg-primary-600 hover:bg-primary-500
                           disabled:opacity-50 disabled:cursor-not-allowed
                           text-white rounded-lg transition-colors font-medium"
              >
                {actionInProgress === 'validate' ? (
                  <Loader2 className="w-5 h-5 animate-spin" />
                ) : (
                  <CheckCircle className="w-5 h-5" />
                )}
                Validate Top {validateTopN}
              </button>
            </>
          )}
        </div>

        {/* Validation Settings */}
        <div className="mt-3 flex items-center gap-2">
          <span className="text-sm text-dark-400">Validate top</span>
          <input
            type="number"
            min={1}
            max={20}
            value={validateTopN}
            onChange={(e) => setValidateTopN(parseInt(e.target.value) || 5)}
            className="w-16 px-2 py-1 bg-dark-700 border border-dark-600 rounded text-white text-sm text-center"
          />
          <span className="text-sm text-dark-400">NN predictions with FEA</span>
        </div>

        {/* Optuna Dashboard Button */}
        <div className="mt-4">
          <button
            onClick={handleLaunchOptuna}
            disabled={actionInProgress !== null || optunaAvailable === false}
            className={`w-full flex items-center justify-center gap-2 px-4 py-2
                       bg-dark-700 hover:bg-dark-600 border border-dark-600
                       disabled:opacity-50 disabled:cursor-not-allowed
                       text-dark-300 hover:text-white rounded-lg transition-colors text-sm
                       ${optunaAvailable === false ? 'cursor-not-allowed' : ''}`}
          >
            {actionInProgress === 'optuna' ? (
              <Loader2 className="w-4 h-4 animate-spin" />
            ) : (
              <ExternalLink className="w-4 h-4" />
            )}
            Launch Optuna Dashboard
          </button>
          {optunaAvailable === false && optunaInstallHint && (
            <div className="mt-1 flex items-center gap-1 text-xs text-dark-500">
              <Info className="w-3 h-3" />
              Install with: {optunaInstallHint}
            </div>
          )}
        </div>
      </div>
    </div>
  );
};

export default ControlPanel;