import React, { useEffect, useState } from 'react';

// Sample modelList definition; replace this with your actual modelList
const modelList = {
  'Offline Fast': {
    model_type: 'local',
    Math: 'qwen2-math:1.5b',
    Code: 'starcoder2',
    Language: 'llama3.2',
    Character: 'dolphin-phi',
    Finance: 'qwen2-math:1.5b',
    Weather: 'llama3.2',
    Time: 'llama3.2',
    Image: 'llava-phi3'
  },
  'Offline Slow': {
    model_type: 'local',
    Math: 'wizard-math',
    Code: 'starcoder2:7b',
    Language: 'llama3.1',
    Character: 'dolphin-llama3',
    Finance: 'wizard-math',
    Weather: 'llama3.1',
    Time: 'llama3.1',
    Image: 'llava'
  },
  'Offline Fast (FOSS)': {
    model_type: 'local',
    Math: 'qwen2-math:1.5b',
    Code: 'qwen2.5-coder:1.5b',
    Language: 'phi3.5',
    Character: 'dolphin-mistral',
    Finance: 'qwen2-math:1.5b',
    Weather: 'phi3.5',
    Time: 'phi3.5',
    Image: 'llava'
  },
  'Offline Slow (FOSS)': {
    model_type: 'local',
    Math: 'mathstral',
    Code: 'qwen2.5-coder',
    Language: 'qwen2.5',
    Character: 'dolphin-mistral',
    Finance: 'mathstral',
    Weather: 'qwen2.5',
    Time: 'qwen2.5',
    Image: 'llava'
  },
  'Online Cheap (OpenAI)': {
    model_type: 'openai',
    Math: 'gpt-4o-mini',
    Code: 'gpt-4o-mini',
    Language: 'gpt-4o-mini',
    Character: 'gpt-4o-mini',
    Finance: 'gpt-4o-mini',
    Weather: 'gpt-4o-mini',
    Time: 'gpt-4o-mini',
    Image: 'gpt-4o-mini'
  },
  'Online Expensive (OpenAI)': {
    model_type: 'openai',
    Math: 'gpt-4o',
    Code: 'gpt-4o',
    Language: 'gpt-4o',
    Character: 'gpt-4o',
    Finance: 'gpt-4o',
    Weather: 'gpt-4o',
    Time: 'gpt-4o',
    Image: 'gpt-4o'
  },
  'Online Cheap (Anthropic)': {
    model_type: 'anthropic',
    Math: 'claude-3-haiku',
    Code: 'claude-3-haiku',
    Language: 'claude-3-haiku',
    Character: 'claude-3-haiku',
    Finance: 'claude-3-haiku',
    Weather: 'claude-3-haiku',
    Time: 'claude-3-haiku',
    Image: 'claude-3-haiku'
  },
  'Online Expensive (Anthropic)': {
    model_type: 'anthropic',
    Math: 'claude-3-5-sonnet',
    Code: 'claude-3-5-sonnet',
    Language: 'claude-3-5-sonnet',
    Character: 'claude-3-5-sonnet',
    Finance: 'claude-3-5-sonnet',
    Weather: 'claude-3-5-sonnet',
    Time: 'claude-3-5-sonnet',
    Image: 'claude-3-5-sonnet'
  },
  'Online Cheap (Google)': {
    model_type: 'google',
    Math: 'gemini-1.5-flash-latest',
    Code: 'gemini-1.5-flash-latest',
    Language: 'gemini-1.5-flash-latest',
    Character: 'gemini-1.5-flash-latest',
    Finance: 'gemini-1.5-flash-latest',
    Weather: 'gemini-1.5-flash-latest',
    Time: 'gemini-1.5-flash-latest',
    Image: 'gemini-1.5-flash-latest'
  },
  'Online Expensive (Google)': {
    model_type: 'google',
    Math: 'gemini-1.5-pro-latest',
    Code: 'gemini-1.5-pro-latest',
    Language: 'gemini-1.5-pro-latest',
    Character: 'gemini-1.5-pro-latest',
    Finance: 'gemini-1.5-pro-latest',
    Weather: 'gemini-1.5-pro-latest',
    Time: 'gemini-1.5-pro-latest',
    Image: 'gemini-1.5-pro-latest'
  },
  'Online (La Plateforme)': {
    model_type: 'mistral',
    Math: 'open-mistral-nemo',
    Code: 'codestral-latest',
    Language: 'mistral-small-latest',
    Character: 'mistral-large-latest',
    Finance: 'open-mistral-nemo',
    Weather: 'mistral-small-latest',
    Time: 'mistral-small-latest',
    Image: 'pixtral-12b-2409'
  },
  'Online (FOSS) (La Plateforme)': {
    model_type: 'mistral',
    Math: 'open-mistral-nemo',
    Code: 'open-codestral-mamba',
    Language: 'open-mistral-nemo',
    Character: 'open-mixtral-8x22b',
    Finance: 'open-mixtral-8x22b',
    Weather: 'open-mistral-nemo',
    Time: 'open-mistral-nemo',
    Image: 'pixtral-12b-2409'
  }
}
  // Define the keys of modelList
  type ModelKeys = keyof typeof modelList;

  const ModelSection: React.FC = () => {
    // Initialize state with value from localStorage or default to ''
    const [selectedModelDropdown, setSelectedModelDropdown] = useState<ModelKeys>('Offline Fast');
    const [radioSelection, setRadioSelection] = useState<string | null>("");
    const [activeSelectedAIFunction, setActiveSelectedAIFunction] = useState('');
    const [currentSelectedAIFunction, setCurrentSelectedAIFunction] = useState<string | null>("");

    useEffect(() => {
      // Load initial values from localStorage
      const temp = localStorage.getItem("activeSelectedAIFunction") || "";
      setActiveSelectedAIFunction(temp);

      // Check if localStorage keys exist, else set defaults
      if (!localStorage.getItem('selectedModelDropdown')) {
        localStorage.setItem("selectedModelDropdown", "Offline Fast");
      }

      // Set the dropdown and selected AI function from localStorage
      const storedSelectedModel = localStorage.getItem('selectedModelDropdown') as ModelKeys;
      const storedActiveFunction = localStorage.getItem("activeSelectedAIFunction") || 'Code';

      setSelectedModelDropdown(storedSelectedModel);
      setCurrentSelectedAIFunction(storedActiveFunction);

      if (!localStorage.getItem("model")) {
        localStorage.setItem("model", 'starcoder2');
      }

      if (!localStorage.getItem("radioSelection")) {
        localStorage.setItem("radioSelection", 'None');
      }

      // Load radio selection from localStorage
      setRadioSelection(localStorage.getItem('radioSelection') || '');
    }, []);

    const filteredModels = Object.keys(modelList).filter(model => modelList[model as ModelKeys]);

    const handleModelChange = (event: React.ChangeEvent<HTMLSelectElement>) => {
      const selectedModel = event.target.value as ModelKeys;
      setSelectedModelDropdown(selectedModel);
      localStorage.setItem('selectedModelDropdown', selectedModel);

      // Update model and type in local storage
      const modelType = modelList[selectedModel].model_type;
      const model = modelList[selectedModel][activeSelectedAIFunction as keyof typeof modelList[typeof selectedModel]];

      localStorage.setItem('model', model);
      localStorage.setItem('type', modelType);
    };

    const isOfflineModel = (model: string) => {
      return modelList[model as ModelKeys].model_type === 'local';
    };

    const modelClicked = (model: string) => {
      const selectedAIFunction = currentSelectedAIFunction as keyof typeof modelList;

      if (modelList[selectedModelDropdown] && modelList[selectedModelDropdown][model as keyof typeof modelList[typeof selectedModelDropdown]]) {
        const newModel = modelList[selectedModelDropdown][model as keyof typeof modelList[typeof selectedModelDropdown]];
        const modelType = modelList[selectedModelDropdown]['model_type' as keyof typeof modelList[typeof selectedModelDropdown]];

        localStorage.setItem('activeSelectedAIFunction', model);
        setActiveSelectedAIFunction(model);
        localStorage.setItem('model', newModel);
        localStorage.setItem('type', modelType);
        setCurrentSelectedAIFunction(model); // Ensure to set current function when model is clicked
      } else {
        console.error(`Model ${model} not found for function ${selectedModelDropdown}`);
      }
    };

    return (
      <div className="model-background">
        <div className="models">
          <div className="title">
            <h3>Different AI Models</h3>
          </div>

          {/* Model Selection Dropdown */}
          <div className="model-dropdown">
            <label htmlFor="model-select">Select AI Model:</label>
            <select id="model-select" value={selectedModelDropdown} onChange={handleModelChange}>
              {filteredModels.map((model) => (
                <option key={model} value={model}>
                  {model}
                </option>
              ))}
            </select>
          </div>

          {/* Model Grid with Cards */}
          <div className="grid">
            {Object.keys(modelList[selectedModelDropdown])
              .filter((key) => key !== 'model_type') // Exclude model_type from display
              .map((displayedCategory) => (
                <button
                  key={displayedCategory}
                  className={`${displayedCategory.toLowerCase()}-model model-box ${currentSelectedAIFunction === displayedCategory ? 'selected' : ''}`} 
                  onClick={() => modelClicked(displayedCategory)}
                >
                  <div className="overlay">
                    <h3>{displayedCategory}</h3>
                    {isOfflineModel(selectedModelDropdown) && <img src="/img/nowifi.svg" alt="No Wi-Fi" />}
                  </div>
                </button>
              ))
            }
          </div>
        </div>
      </div>
    );
  };

export default ModelSection;