"use client";
import React, { useState, useEffect } from 'react';

// Define the types for ModelType and ModelListType
type ModelType = {
  model_type: string;
  Math: string;
  Code: string;
  Language: string;
  Weather: string;
};

type ModelListType = {
  [key: string]: ModelType;
};

// Define the AI models list
const modelList: ModelListType = {
  'Offline Fast': {
    model_type: 'local',
    Math: 'qwen2-math:1.5b',
    Code: 'starcoder2',
    Language: 'llama3.2',
    Weather: 'llama3.2',
  },
  'Offline Slow': {
    model_type: 'local',
    Math: 'wizard-math',
    Code: 'starcoder2:7b',
    Language: 'llama3.1',
    Weather: 'llama3.1',
  },
  'Offline Fast (FOSS)': {
    model_type: 'local',
    Math: 'qwen2-math:1.5b',
    Code: 'qwen2.5-coder:1.5b',
    Language: 'phi3.5',
    Weather: 'phi3.5',
  },
  'Offline Slow (FOSS)': {
    model_type: 'local',
    Math: 'mathstral',
    Code: 'qwen2.5-coder',
    Language: 'qwen2.5',
    Weather: 'qwen2.5',
  },
  'Online Cheap (OpenAI)': {
    model_type: 'openai',
    Math: 'gpt-4o-mini',
    Code: 'gpt-4o-mini',
    Language: 'gpt-4o-mini',
    Weather: 'gpt-4o-mini',
  },
  'Online Expensive (OpenAI)': {
    model_type: 'openai',
    Math: 'gpt-4o',
    Code: 'gpt-4o',
    Language: 'gpt-4o',
    Weather: 'gpt-4o',
  },
  'Online Cheap (Anthropic)': {
    model_type: 'anthropic',
    Math: 'claude-3-haiku',
    Code: 'claude-3-haiku',
    Language: 'claude-3-haiku',
    Weather: 'claude-3-haiku',
  },
  'Online Expensive (Anthropic)': {
    model_type: 'anthropic',
    Math: 'claude-3-5-sonnet',
    Code: 'claude-3-5-sonnet',
    Language: 'claude-3-5-sonnet',
    Weather: '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',
    Weather: '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',
    Weather: 'gemini-1.5-pro-latest',
  },
  'Online (La Plateforme)': {
    model_type: 'mistral',
    Math: 'open-mistral-nemo',
    Code: 'codestral-latest',
    Language: 'mistral-small-latest',
    Weather: 'mistral-small-latest',
  },
  'Online (FOSS) (La Plateforme)': {
    model_type: 'mistral',
    Math: 'open-mistral-nemo',
    Code: 'open-codestral-mamba',
    Language: 'open-mistral-nemo',
    Weather: 'open-mistral-nemo',
  },
};

// AI Functions list
const aiFunctions = ['Math', 'Code', 'Language', 'Weather'] as const;
type AIFunction = typeof aiFunctions[number]; // Restrict to these exact function names

const ModelSection: React.FC = () => {
  const [selectedModelDropdown, setSelectedModelDropdown] = useState(() => "Offline Fast");
  const [activeSelectedAIFunction, setActiveSelectedAIFunction] = useState(() => "Code");
  const [, setSelectedModel] = useState<string>("");
  const [, setSelectedModelType] = useState<string>("");

  useEffect(() => {
    setSelectedModelDropdown(localStorage.getItem("selectedModelDropdown")|| "Offline Fast");
    setActiveSelectedAIFunction(localStorage.getItem("activeSelectedAIFunction") || "Code");
  }, []);

  // Update the model and type when the dropdown or function changes
  useEffect(() => {
    const aiFunctionsActiveSelectedAIFunction = activeSelectedAIFunction as AIFunction
    const newSelectedModel = modelList[selectedModelDropdown]?.[aiFunctionsActiveSelectedAIFunction] || "";
    const newModelType = modelList[selectedModelDropdown]?.model_type || "";

    setSelectedModel(newSelectedModel);
    setSelectedModelType(newModelType);

    localStorage.setItem("model", newSelectedModel);
    localStorage.setItem("type", newModelType);
  }, [selectedModelDropdown, activeSelectedAIFunction]);

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

  const modelClicked = (functionName: AIFunction) => {
    setActiveSelectedAIFunction(functionName);
    localStorage.setItem('activeSelectedAIFunction', functionName);
  };

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

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

        <div className="grid">
          {aiFunctions.map((functionName) => (
            <button
              key={functionName}
              className={`${functionName.toLowerCase()}-model model-box ${activeSelectedAIFunction === functionName ? 'selected' : ''}`}
              onClick={() => modelClicked(functionName)}
            >
              <div className="overlay">
                <h3>{functionName}</h3>
              </div>
            </button>
          ))}
        </div>
      </div>
    </div>
  );
};

export default ModelSection;