import React, { useState, useEffect, ReactElement } from 'react';
import FormGroup from '@mui/material/FormGroup';
import FormControlLabel from '@mui/material/FormControlLabel';
import Switch from '@mui/material/Switch';
import Divider from '@mui/material/Divider';
import TextField from '@mui/material/TextField';
import Accordion from '@mui/material/Accordion';
import AccordionActions from '@mui/material/AccordionActions';
import AccordionSummary from '@mui/material/AccordionSummary';
import AccordionDetails from '@mui/material/AccordionDetails';
import Typography from '@mui/material/Typography';
import Button from '@mui/material/Button';
import Box from '@mui/material/Box';
import ResetIcon from '@mui/icons-material/History';
import ExpandMoreIcon from '@mui/icons-material/ExpandMore';

import { SetSnackType } from './Snack';

type Tool = {
  type: string,
  function?: {
    name: string,
    description: string,
    parameters?: any,
    returns?: any
  },
  name?: string,
  description?: string,
  enabled: boolean
};

interface ControlsParams {
  connectionBase: string,
  sessionId: string | undefined,
  setSnack: SetSnackType,
};

type GPUInfo = {
  name: string,
  memory: number,
  discrete: boolean
}


type SystemInfo = {
  "Installed RAM": string,
  "Graphics Card": GPUInfo[],
  "CPU": string
};

const SystemInfoComponent: React.FC<{ systemInfo: SystemInfo | undefined }> = ({ systemInfo }) => {
  const [systemElements, setSystemElements] = useState<ReactElement[]>([]);

  const convertToSymbols = (text: string) => {
    return text
      .replace(/\(R\)/g, '®')  // Replace (R) with the ® symbol
      .replace(/\(C\)/g, '©')  // Replace (C) with the © symbol
      .replace(/\(TM\)/g, '™'); // Replace (TM) with the ™ symbol
  };

  useEffect(() => {
    if (systemInfo === undefined) {
      return;
    }
    const elements = Object.entries(systemInfo).flatMap(([k, v]) => {
      // If v is an array, repeat for each card
      if (Array.isArray(v)) {
        return v.map((card, index) => (
          <div key={index} className="SystemInfoItem">
            <div>{convertToSymbols(k)} {index}</div>
            <div>{convertToSymbols(card.name)} {card.discrete ? `w/ ${Math.round(card.memory / (1024 * 1024 * 1024))}GB RAM` : "(integrated)"}</div>
          </div>
        ));
      }

      // If it's not an array, handle normally
      return (
        <div key={k} className="SystemInfoItem">
          <div>{convertToSymbols(k)}</div>
          <div>{convertToSymbols(String(v))}</div>
        </div>
      );
    });

    setSystemElements(elements);
  }, [systemInfo]);

  return <div className="SystemInfo">{systemElements}</div>;
};

const Controls = ({ sessionId, setSnack, connectionBase }: ControlsParams) => {
  const [editSystemPrompt, setEditSystemPrompt] = useState<string>("");
  const [systemInfo, setSystemInfo] = useState<SystemInfo | undefined>(undefined);
  const [tools, setTools] = useState<Tool[]>([]);
  const [rags, setRags] = useState<Tool[]>([]);
  const [systemPrompt, setSystemPrompt] = useState<string>("");
  const [serverSystemPrompt, setServerSystemPrompt] = useState<string>("");
  const [messageHistoryLength, setMessageHistoryLength] = useState<number>(5);

    useEffect(() => {
      if (systemPrompt === serverSystemPrompt || !systemPrompt.trim() || sessionId === undefined) {
        return;
      }
      const sendSystemPrompt = async (prompt: string) => {
        try {
          const response = await fetch(connectionBase + `/api/tunables/${sessionId}`, {
            method: 'PUT',
            headers: {
              'Content-Type': 'application/json',
              'Accept': 'application/json',
            },
            body: JSON.stringify({ "system_prompt": prompt }),
          });
  
          const data = await response.json();
          const newPrompt = data["system_prompt"];
          if (newPrompt !== serverSystemPrompt) {
            setServerSystemPrompt(newPrompt);
            setSystemPrompt(newPrompt)
            setSnack("System prompt updated", "success");
          }
        } catch (error) {
          console.error('Fetch error:', error);
          setSnack("System prompt update failed", "error");
        }
      };
  
      sendSystemPrompt(systemPrompt);
  
    }, [systemPrompt, setServerSystemPrompt, serverSystemPrompt, connectionBase, sessionId, setSnack]);
  
    useEffect(() => {
      if (sessionId === undefined) {
        return;
      }
      const sendMessageHistoryLength = async (length: number) => {
        try {
          const response = await fetch(connectionBase + `/api/tunables/${sessionId}`, {
            method: 'PUT',
            headers: {
              'Content-Type': 'application/json',
              'Accept': 'application/json',
            },
            body: JSON.stringify({ "message_history_length": length }),
          });
  
          const data = await response.json();
          const newLength = data["message_history_length"];
          if (newLength !== messageHistoryLength) {
            setMessageHistoryLength(newLength);
            setSnack("Message history length updated", "success");
          }
        } catch (error) {
          console.error('Fetch error:', error);
          setSnack("Message history length update failed", "error");
        }
      };
  
      sendMessageHistoryLength(messageHistoryLength);
  
    }, [messageHistoryLength, setMessageHistoryLength, connectionBase, sessionId, setSnack]);

  const reset = async (types: ("rags" | "tools" | "history" | "system_prompt" | "message_history_length")[], message: string = "Update successful.") => {
      try {
        const response = await fetch(connectionBase + `/api/reset/${sessionId}`, {
          method: 'PUT',
          headers: {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
          },
          body: JSON.stringify({ "reset": types }),
        });
  
        if (response.ok) {
          const data = await response.json();
          if (data.error) {
            throw Error()
          }
          for (const [key, value] of Object.entries(data)) {
            switch (key) {
              case "rags":
                setRags(value as Tool[]);
                break;
              case "tools":
                setTools(value as Tool[]);
                break;
              case "system_prompt":
                setServerSystemPrompt((value as any)["system_prompt"].trim());
                setSystemPrompt((value as any)["system_prompt"].trim());
                break;
              case "history":
                console.log('TODO: handle history reset');
                break;
            }
          }
          setSnack(message, "success");
        } else {
          throw Error(`${{ status: response.status, message: response.statusText }}`);
        }
      } catch (error) {
        console.error('Fetch error:', error);
        setSnack("Unable to restore defaults", "error");
      }
    };
  
    
  // Get the system information
  useEffect(() => {
    if (systemInfo !== undefined || sessionId === undefined) {
      return;
    }
    fetch(connectionBase + `/api/system-info/${sessionId}`, {
      method: 'GET',
      headers: {
        'Content-Type': 'application/json',
      },
    })
      .then(response => response.json())
      .then(data => {
        setSystemInfo(data);
      })
      .catch(error => {
        console.error('Error obtaining system information:', error);
        setSnack("Unable to obtain system information.", "error");
      });
  }, [systemInfo, setSystemInfo, connectionBase, setSnack, sessionId])
    
  useEffect(() => {
    setEditSystemPrompt(systemPrompt);
  }, [systemPrompt, setEditSystemPrompt]);


  const toggleRag = async (tool: Tool) => {
    tool.enabled = !tool.enabled
    try {
      const response = await fetch(connectionBase + `/api/rags/${sessionId}`, {
        method: 'PUT',
        headers: {
          'Content-Type': 'application/json',
          'Accept': 'application/json',
        },
        body: JSON.stringify({ "tool": tool?.name, "enabled": tool.enabled }),
      });

      const rags = await response.json();
      setRags([...rags])
      setSnack(`${tool?.name} ${tool.enabled ? "enabled" : "disabled"}`);
    } catch (error) {
      console.error('Fetch error:', error);
      setSnack(`${tool?.name} ${tool.enabled ? "enabling" : "disabling"} failed.`, "error");
      tool.enabled = !tool.enabled
    }
  };

  const toggleTool = async (tool: Tool) => {
    tool.enabled = !tool.enabled
    try {
      const response = await fetch(connectionBase + `/api/tools/${sessionId}`, {
        method: 'PUT',
        headers: {
          'Content-Type': 'application/json',
          'Accept': 'application/json',
        },
        body: JSON.stringify({ "tool": tool?.function?.name, "enabled": tool.enabled }),
      });

      const tools = await response.json();
      setTools([...tools])
      setSnack(`${tool?.function?.name} ${tool.enabled ? "enabled" : "disabled"}`);
    } catch (error) {
      console.error('Fetch error:', error);
      setSnack(`${tool?.function?.name} ${tool.enabled ? "enabling" : "disabling"} failed.`, "error");
      tool.enabled = !tool.enabled
    }
  };

  // If the tools have not been set, fetch them from the server
  useEffect(() => {
    if (tools.length || sessionId === undefined) {
      return;
    }
    const fetchTools = async () => {
      try {
        // Make the fetch request with proper headers
        const response = await fetch(connectionBase + `/api/tools/${sessionId}`, {
          method: 'GET',
          headers: {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
          },
        });
        if (!response.ok) {
          throw Error();
        }
        const tools = await response.json();
        setTools(tools);
      } catch (error: any) {
        setSnack("Unable to fetch tools", "error");
        console.error(error);
      }
    }

    fetchTools();
  }, [sessionId, tools, setTools, setSnack, connectionBase]);

  // If the RAGs have not been set, fetch them from the server
  useEffect(() => {
    if (rags.length || sessionId === undefined) {
      return;
    }
    const fetchRags = async () => {
      try {
        // Make the fetch request with proper headers
        const response = await fetch(connectionBase + `/api/rags/${sessionId}`, {
          method: 'GET',
          headers: {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
          },
        });
        if (!response.ok) {
          throw Error();
        }
        const rags = await response.json();
        setRags(rags);
      } catch (error: any) {
        setSnack("Unable to fetch RAGs", "error");
        console.error(error);
      }
    }

    fetchRags();
  }, [sessionId, rags, setRags, setSnack, connectionBase]);

    // If the systemPrompt has not been set, fetch it from the server
    useEffect(() => {
      if (serverSystemPrompt !== "" || sessionId === undefined) {
        return;
      }
      const fetchTunables = async () => {
        // Make the fetch request with proper headers
        const response = await fetch(connectionBase + `/api/tunables/${sessionId}`, {
          method: 'GET',
          headers: {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
          },
        });
        const data = await response.json();
        const serverSystemPrompt = data["system_prompt"].trim();
        setServerSystemPrompt(serverSystemPrompt);
        setSystemPrompt(serverSystemPrompt);
        setMessageHistoryLength(data["message_history_length"]);
      }
  
      fetchTunables();
    }, [sessionId, serverSystemPrompt, setServerSystemPrompt, connectionBase]);
  
  



  const toggle = async (type: string, index: number) => {
    switch (type) {
      case "rag":
        toggleRag(rags[index])
        break;
      case "tool":
        toggleTool(tools[index]);
    }
  };

  const handleKeyPress = (event: any) => {
    if (event.key === 'Enter' && event.ctrlKey) {
      switch (event.target.id) {
        case 'SystemPromptInput':
          setSystemPrompt(editSystemPrompt);
          break;
      }
    }
  };

  return (<div className="Controls">
    <Typography component="span" sx={{ mb: 1 }}>
      You can change the information available to the LLM by adjusting the following settings:

    </Typography>
    <Accordion>
      <AccordionSummary expandIcon={<ExpandMoreIcon />}>
        <Typography component="span">System Prompt</Typography>
      </AccordionSummary>
      <AccordionActions style={{ flexDirection: "column" }}>
        <TextField
          variant="outlined"
          fullWidth
          multiline
          type="text"
          value={editSystemPrompt}
          onChange={(e) => setEditSystemPrompt(e.target.value)}
          onKeyDown={handleKeyPress}
          placeholder="Enter the new system prompt.."
          id="SystemPromptInput"
        />
        <div style={{ display: "flex", flexDirection: "row", gap: "8px", paddingTop: "8px" }}>
          <Button variant="contained" disabled={editSystemPrompt === systemPrompt} onClick={() => { setSystemPrompt(editSystemPrompt); }}>Set</Button>
          <Button variant="outlined" onClick={() => { reset(["system_prompt"], "System prompt reset."); }} color="error">Reset</Button>
        </div>
      </AccordionActions>
    </Accordion>

    <Accordion>
      <AccordionSummary expandIcon={<ExpandMoreIcon />}>
        <Typography component="span">Tunables</Typography>
      </AccordionSummary>
      <AccordionActions style={{ flexDirection: "column" }}>
        <TextField
          id="outlined-number"
          label="Message history"
          type="number"
          helperText="Only use this many messages as context. 0 = All. Keeping this low will reduce context growth and improve performance."
          value={messageHistoryLength}
          onChange={(e: any) => setMessageHistoryLength(e.target.value)}
          slotProps={{
            htmlInput: {
              min: 0
            },
            inputLabel: {
              shrink: true,
            },
          }}
        />
      </AccordionActions>
    </Accordion>

    <Accordion>
      <AccordionSummary expandIcon={<ExpandMoreIcon />}>
        <Typography component="span">Tools</Typography>
      </AccordionSummary>
      <AccordionDetails>
        These tools can be made available to the LLM for obtaining real-time information from the Internet. The description provided to the LLM is provided for reference.
      </AccordionDetails>
      <AccordionActions>
        <FormGroup sx={{ p: 1 }}>
          {
            tools.map((tool, index) =>
              <Box key={index}>
                <Divider />
                <FormControlLabel control={<Switch checked={tool.enabled} />} onChange={() => toggle("tool", index)} label={tool?.function?.name} />
                <Typography sx={{ fontSize: "0.8rem", mb: 1 }}>{tool?.function?.description}</Typography>
              </Box>
            )
          }</FormGroup>
      </AccordionActions>
    </Accordion>

    <Accordion>
      <AccordionSummary expandIcon={<ExpandMoreIcon />}>
        <Typography component="span">RAG</Typography>
      </AccordionSummary>
      <AccordionDetails>
        These RAG databases can be enabled / disabled for adding additional context based on the chat request.
      </AccordionDetails>
      <AccordionActions>
        <FormGroup sx={{ p: 1, flexGrow: 1, justifyContent: "flex-start" }}>
          {
            rags.map((rag, index) =>
              <Box key={index} sx={{ display: "flex", flexGrow: 1, flexDirection: "column" }}>
                <Divider />
                <FormControlLabel
                  control={<Switch checked={rag.enabled} />}
                  onChange={() => toggle("rag", index)} label={rag?.name}
                />
                <Typography>{rag?.description}</Typography>
              </Box>
            )
          }</FormGroup>
      </AccordionActions>
    </Accordion>
    <Accordion>
      <AccordionSummary expandIcon={<ExpandMoreIcon />}>
        <Typography component="span">System Information</Typography>
      </AccordionSummary>
      <AccordionDetails>
        The server is running on the following hardware:
      </AccordionDetails>
      <AccordionActions>
        <SystemInfoComponent systemInfo={systemInfo} />
      </AccordionActions>
    </Accordion>
    <Button startIcon={<ResetIcon />} onClick={() => { reset(["history"], "History cleared."); }}>Clear Backstory History</Button>
    <Button onClick={() => { reset(["rags", "tools", "system_prompt", "message_history_length"], "Default settings restored.") }}>Reset system prompt, tunables, and RAG to defaults</Button>
  </div>);
}


export type {
  ControlsParams
};

export { 
  Controls
};