// src/components/EngineManager.tsx
import { useState, useEffect } from 'react';
import apiClient from '../services/api';

interface ModelConfig {
  model_id: string;
  display_name: string;
  provider?: string;
  model_name?: string;
  max_tokens?: number;
  temperature?: number;
  timeout?: number;
  is_default?: boolean;
  supported_domains?: string[];
  description?: string;
}

interface EngineStatus extends ModelConfig {
  status: 'master' | 'apprentice' | 'available' | 'retired';
  unique_id?: string; // For apprentice engines
}

const EngineManager = () => {
  const [allEngines, setAllEngines] = useState<EngineStatus[]>([]);
  const [currentMasterId, setCurrentMasterId] = useState('');
  const [currentApprenticeId, setCurrentApprenticeId] = useState('none');
  const [isLoading, setIsLoading] = useState(true);
  const [error, setError] = useState('');

  const [newModelProvider, setNewModelProvider] = useState('ollama');
  const [newModelId, setNewModelId] = useState('');
  const [newModelDisplayName, setNewModelDisplayName] = useState('');
  const [newModelModelName, setNewModelModelName] = useState('');
  const [newModelSupportedDomains, setNewModelSupportedDomains] = useState(''); // Comma separated
  const [newModelDescription, setNewModelDescription] = useState('');
  const [newModelFile, setNewModelFile] = useState<File | null>(null); // For GGUF file upload

  const fetchModelsAndActiveEngines = async () => {
    try {
      const [allEnginesResponse, activeEnginesResponse] = await Promise.all([
        apiClient.get('/config/engines'), // Fetch all engines with their statuses
        apiClient.get('/config/engines/active') // Fetch current active master/apprentice
      ]);
      setAllEngines(allEnginesResponse.data);
      
      const { master_engine_id, apprentice_engine_id } = activeEnginesResponse.data;
      if (master_engine_id) setCurrentMasterId(master_engine_id);
      if (apprentice_engine_id) setCurrentApprenticeId(apprentice_engine_id);
      
    } catch (err) {
      setError('Failed to load engines or active engines.');
      console.error(err);
    } finally {
      setIsLoading(false);
    }
  };

  useEffect(() => {
    fetchModelsAndActiveEngines();
  }, []);

  const handleSetActiveEngines = async () => {
    try {
      await apiClient.post('/config/engines/active', {
        master_engine_id: currentMasterId,
        apprentice_engine_id: currentApprenticeId === 'none' ? null : currentApprenticeId,
      });
      alert('Active engines updated successfully!');
      fetchModelsAndActiveEngines(); // Re-fetch to confirm and update UI
    } catch (err: any) {
      const errorMessage = err.response?.data?.detail || err.message;
      alert(`Failed to update active engines: ${errorMessage}`);
      console.error(err);
    }
  };


  const handleSubmitNewModel = async () => {
    if (!newModelId || !newModelDisplayName || !newModelModelName) {
      alert('Please fill in Model ID, Display Name, and Model Name.');
      return;
    }

    let actualModelName = newModelModelName; // Default to user-provided model name

    // Handle GGUF file upload first if applicable
    if (newModelProvider === 'gguf') {
      if (!newModelFile) {
        alert('Please select a GGUF file for the GGUF provider.');
        return;
      }
      try {
        const formData = new FormData();
        formData.append('file', newModelFile);
        const uploadResponse = await apiClient.post('/config/upload-gguf', formData, {
          headers: {
            'Content-Type': 'multipart/form-data',
          },
        });
        actualModelName = uploadResponse.data.path; // Use the path returned by the backend
        alert(`GGUF file uploaded successfully to: ${actualModelName}`);
      } catch (err: any) {
        const errorMessage = err.response?.data?.detail || err.message;
        alert(`Failed to upload GGUF file: ${errorMessage}`);
        console.error(err);
        return; // Stop processing if file upload fails
      }
    }

    const modelData: ModelConfig = {
      model_id: newModelId,
      display_name: newModelDisplayName,
      provider: newModelProvider,
      model_name: actualModelName, // Use the actualModelName (might be local file path)
      max_tokens: 4096, // Default values
      temperature: 0.7,
      timeout: 120,
      is_default: false,
      supported_domains: newModelSupportedDomains.split(',').map(s => s.trim()).filter(s => s),
      description: newModelDescription,
    };

    try {
      await apiClient.post('/config/models', modelData);
      alert('Model registered successfully!');
      fetchModelsAndActiveEngines(); // Refresh the model list
      // Clear form
      setNewModelId('');
      setNewModelDisplayName('');
      setNewModelModelName('');
      setNewModelSupportedDomains('');
      setNewModelDescription('');
      setNewModelFile(null);
    } catch (err: any) {
      const errorMessage = err.response?.data?.detail || err.message;
      alert(`Failed to register model: ${errorMessage}`);
      console.error(err);
    }
  };

  if (isLoading) {
    return <div className="panel-section"><h3>Engine Management</h3><p>Loading models...</p></div>;
  }

  if (error) {
    return <div className="panel-section"><h3>Engine Management</h3><p style={{color: 'red'}}>{error}</p></div>;
  }

  const handleSwapEngines = async (apprenticeModelId: string) => {
    try {
      await apiClient.post('/config/engines/swap', { apprentice_model_id: apprenticeModelId });
      alert('Engines swapped successfully!');
      fetchModelsAndActiveEngines(); // Refresh UI
    } catch (err: any) {
      const errorMessage = err.response?.data?.detail || err.message;
      alert(`Failed to swap engines: ${errorMessage}`);
      console.error(err);
    }
  };

  const handlePromoteApprentice = async (apprenticeModelId: string) => {
    try {
      await apiClient.post('/config/engines/promote', { apprentice_model_id: apprenticeModelId });
      alert('Apprentice promoted to Master successfully!');
      fetchModelsAndActiveEngines(); // Refresh UI
    } catch (err: any) {
      const errorMessage = err.response?.data?.detail || err.message;
      alert(`Failed to promote apprentice: ${errorMessage}`);
      console.error(err);
    }
  };

  const handleCreateNewApprentice = async () => {
    try {
      const response = await apiClient.post('/config/engines/apprentice/new');
      alert(response.data.message);
      fetchModelsAndActiveEngines(); // Refresh UI
    } catch (err: any) {
      const errorMessage = err.response?.data?.detail || err.message;
      alert(`Failed to create new apprentice: ${errorMessage}`);
      console.error(err);
    }
  };

  return (
    <div className="panel-section">
      <h3>Engine Management</h3>

      {/* Current Active Engines Display */}
      <div className="mb-4 p-3 border rounded">
        <p className="text-sm font-medium">
          Master Engine:{' '}
          <span className="font-semibold text-blue-600">
            {allEngines.find(e => e.model_id === currentMasterId)?.display_name || 'Not Set'}
          </span>
        </p>
        <p className="text-sm font-medium">
          Apprentice Engine:{' '}
          <span className="font-semibold text-green-600">
            {allEngines.find(e => e.model_id === currentApprenticeId)?.display_name || 'None'}
          </span>
        </p>
      </div>

      {/* Set Active Engines (existing functionality, but updated) */}
      <div className="mb-4">
        <h4>Set Current Master/Apprentice</h4>
        <div className="form-group">
          <label>Master (Teacher) Engine</label>
          <select value={currentMasterId} onChange={(e) => setCurrentMasterId(e.target.value)}>
            <option value="">Select Master</option>
            {allEngines.filter(e => e.status !== 'retired').map(engine => (
              <option key={engine.model_id} value={engine.model_id}>
                {engine.display_name} ({engine.status === 'master' ? 'Current Master' : engine.status === 'apprentice' ? 'Current Apprentice' : engine.status})
              </option>
            ))}
          </select>
        </div>

        <div className="form-group">
          <label>Apprentice (Student) Engine</label>
          <select value={currentApprenticeId} onChange={(e) => setCurrentApprenticeId(e.target.value)}>
            <option value="none">None</option>
            {allEngines.filter(e => e.status !== 'master' && e.status !== 'retired').map(engine => (
              <option key={engine.model_id} value={engine.model_id}>
                {engine.display_name} ({engine.status === 'apprentice' ? 'Current Apprentice' : engine.status}) {engine.unique_id ? `[${engine.unique_id.substring(0, 4)}]` : ''}
              </option>
            ))}
          </select>
        </div>

        <button onClick={handleSetActiveEngines} className="btn btn-primary">
          Set Active Engines
        </button>
      </div>

      {/* New: Advanced Engine Management */}
      <div style={{ marginTop: '1.5rem', borderTop: '1px solid #eee', paddingTop: '1.5rem' }}>
        <h4>Advanced Engine Operations</h4>

        {/* Swap Master and Apprentice */}
        <div className="mb-3">
          <label className="block text-sm font-medium text-gray-700">Swap Master and Apprentice</label>
          <select 
            value={""} // Temporary value
            onChange={(e) => {
              const selectedApprenticeId = e.target.value;
              if (selectedApprenticeId) handleSwapEngines(selectedApprenticeId);
              e.target.value = ""; // Reset dropdown
            }}
            className="mt-1 block w-full pl-3 pr-10 py-2 text-base border-gray-300 focus:outline-none focus:ring-indigo-500 focus:border-indigo-500 sm:text-sm rounded-md"
          >
            <option value="">Select Apprentice to Swap with Master</option>
            {allEngines.filter(e => e.status === 'apprentice' || (e.status === 'available' && e.unique_id)).map(engine => (
              <option key={engine.model_id} value={engine.model_id}>
                {engine.display_name} {engine.unique_id ? `[${engine.unique_id.substring(0, 4)}]` : ''}
              </option>
            ))}
          </select>
        </div>

        {/* Promote Apprentice */}
        <div className="mb-3">
          <label className="block text-sm font-medium text-gray-700">Promote Apprentice to Master</label>
          <select 
            value={""} // Temporary value
            onChange={(e) => {
              const selectedApprenticeId = e.target.value;
              if (selectedApprenticeId) handlePromoteApprentice(selectedApprenticeId);
              e.target.value = ""; // Reset dropdown
            }}
            className="mt-1 block w-full pl-3 pr-10 py-2 text-base border-gray-300 focus:outline-none focus:ring-indigo-500 focus:border-indigo-500 sm:text-sm rounded-md"
          >
            <option value="">Select Apprentice to Promote</option>
            {allEngines.filter(e => e.status === 'apprentice' || (e.status === 'available' && e.unique_id)).map(engine => (
              <option key={engine.model_id} value={engine.model_id}>
                {engine.display_name} {engine.unique_id ? `[${engine.unique_id.substring(0, 4)}]` : ''}
              </option>
            ))}
          </select>
        </div>

        {/* Create New Apprentice */}
        <div className="mb-6">
          <button onClick={handleCreateNewApprentice} className="btn btn-secondary">
            Create New Empty Apprentice
          </button>
        </div>
      </div>

      {/* Register New Model Section (existing functionality) */}
      <div style={{marginTop: '1.5rem', borderTop: '1px solid #eee', paddingTop: '1.5rem'}}>
        <h4>Register New Model</h4>
        {/* Existing new model registration form content */}
        <div className="form-group">
          <label>Provider</label>
          <select value={newModelProvider} onChange={(e) => setNewModelProvider(e.target.value)}>
            <option value="ollama">Ollama</option>
            <option value="huggingface">HuggingFace (Transformers)</option>
            <option value="mlx">MLX (Apple Silicon)</option>
            <option value="gguf">GGUF (Local File)</option>
          </select>
        </div>
        <div className="form-group">
          <label>Model ID (Unique)</label>
          <input type="text" value={newModelId} onChange={(e) => setNewModelId(e.target.value)} placeholder="e.g., 'my-custom-model'" />
        </div>
        <div className="form-group">
          <label>Display Name</label>
          <input type="text" value={newModelDisplayName} onChange={(e) => setNewModelDisplayName(e.target.value)} placeholder="e.g., 'My Custom Model'" />
        </div>
        <div className="form-group">
          <label>Model Name (for Provider)</label>
          <input type="text" value={newModelModelName} onChange={(e) => setNewModelModelName(e.target.value)} placeholder="e.g., 'gemma2:9b' or 'kofdai/my-hf-model'" />
        </div>
        <div className="form-group">
          <label>Supported Domains (comma-separated)</label>
          <input type="text" value={newModelSupportedDomains} onChange={(e) => setNewModelSupportedDomains(e.target.value)} placeholder="e.g., 'medical,general'" />
        </div>
        <div className="form-group">
          <label>Description</label>
          <textarea value={newModelDescription} onChange={(e) => setNewModelDescription(e.target.value)} placeholder="Optional description"></textarea>
        </div>
        {newModelProvider === 'gguf' && (
          <div className="form-group">
              <label>GGUF File (Select local .gguf file)</label>
              <input type="file" accept=".gguf" onChange={(e) => setNewModelFile(e.target.files ? e.target.files[0] : null)} />
          </div>
        )}
        <button onClick={handleSubmitNewModel} className="btn btn-secondary">
          Register Model
        </button>
      </div>
    </div>
  );
};

export default EngineManager;