import 'dotenv/config'
import express from 'express';
import bodyParser from 'body-parser';
import cmd from 'cmd-promise';
import cors from 'cors';
import cheerio from 'cheerio';
import llamaTokenizer from 'llama-tokenizer-js';
import googleIt from 'google-it';

const prompt = process.env.PROMPT;
const app = express();
const port = 3000;

app.use(cors({
   origin: '*',
   allowedHeaders: ['Content-Type', 'x-forwarded-for-id', 'x-forwarded-for-name']
}));
app.use(bodyParser.json());

let isProcessing = false;
let conversationHistory = {};

// Helper function to get current timestamp
const getTimestamp = () => {
   const now = new Date();
   const date = now.toLocaleDateString('en-US');
   const time = now.toLocaleTimeString('en-US');
   return `${date} [${time}]`;
};

// Middleware to track conversation history by CF-Connecting-IP
app.use((req, res, next) => {
   const ip = req.headers['x-forwarded-for-id'] || req.headers['cf-connecting-ip'] || req.headers['x-forwarded-for'] || req.headers['x-real-ip'] || req.ip;
   const name = req.headers['x-forwarded-for-name'];
   const guild = req.headers['x-forwarded-for-guild'];

   req.clientIp = ip; // Store the IP in a request property

   if (!conversationHistory[req.clientIp]) {
      if (name) {
         if (guild) {
            console.log(`${getTimestamp()} [INFO] Incoming request from: ${req.clientIp} | ${name} within ${guild}`); // Log the IP address
            console.log(`${getTimestamp()} [INFO] Initializing conversation history for: ${req.clientIp}`);

            conversationHistory[req.clientIp] = [{
               role: 'system',
               content: `My name is: ${name}, We are chatting inside ${guild} a Discord Server. ` + prompt
            }];

         } else {
            conversationHistory[req.clientIp] = [{
               role: 'system',
               content: `My name is: ${name} ` + prompt
            }];
         }
      } else {
         console.log(`${getTimestamp()} [INFO] Incoming request from: ${req.clientIp}`); // Log the IP address
         console.log(`${getTimestamp()} [INFO] Initializing conversation history for new IP: ${req.clientIp}`);

         conversationHistory[req.clientIp] = [{
            role: 'system',
            content: prompt
         }];
      }
   }
   next();
});

function countLlamaTokens(messages) {
   let totalTokens = 0;
   for (const message of messages) {
      if (message.role === 'user' || message.role === 'assistant') {
         const encodedTokens = llamaTokenizer.encode(message.content);
         totalTokens += encodedTokens.length;
      }
   }
   return totalTokens;
}

function trimConversationHistory(messages, maxLength, tolerance) {
   let tokenLength = countLlamaTokens(messages);
   if (tokenLength > maxLength + tolerance) {
      const diff = tokenLength - (maxLength + tolerance);
      let removedTokens = 0;

      // Iterate over the messages in reverse order
      for (let i = messages.length - 1; i >= 0; i--) {
         const message = messages[i];
         const messageTokens = countLlamaTokens([message]);

         if (removedTokens + messageTokens <= diff) {
            messages.splice(i, 1);
            removedTokens += messageTokens;
            console.log(`${getTimestamp()} [CLEANUP] ${removedTokens} removed | After Resize: ${countLlamaTokens(messages)}`);
         } else {
            const messagesToRemove = Math.floor(diff / messageTokens);
            for (let j = 0; j < messagesToRemove; j++) {
               messages.splice(i, 1);
               removedTokens += messageTokens;
            }
            break;
         }
      }
   }
}

// Function to scrape web page
async function scrapeWebPage(url, length) {
   console.log(`${getTimestamp()} [INFO] Starting to scrape URL: ${url}`);
   try {
      const res = await fetch(url);
      const html = await res.text();
      const $ = cheerio.load(html);

      // Extract page title, meta description and content
      const pageTitle = $('head title').text().trim();
      const pageDescription = $('head meta[name="description"]').attr('content');
      const pageContent = $('body').text().trim();

      // Construct response message with page details
      let response = `Title: ${pageTitle}\n`;
      if (pageDescription) {
         response += `Description: ${pageDescription}\n`;
      }
      if (pageContent) {
         const MAX_CONTENT_LENGTH = length || process.env.MAX_CONTENT_LENGTH || 2000;
         let plainTextContent = $('<div>').html(pageContent).text().trim().replace(/[\r\n\t]+/g, ' ');
         const codePattern = /\/\/|\/\*|\*\/|\{|\}|\[|\]|\bfunction\b|\bclass\b|\b0x[0-9A-Fa-f]+\b|\b0b[01]+\b/;
         const isCode = codePattern.test(plainTextContent);

         if (isCode) {
            plainTextContent = plainTextContent.replace(codePattern, '');
         }
         plainTextContent = plainTextContent.replace(/ *\([^)]*\) */g, '');
         if (plainTextContent.length > MAX_CONTENT_LENGTH) {
            plainTextContent = plainTextContent.substring(0, MAX_CONTENT_LENGTH) + '...';
         }
         response += `Content: ${plainTextContent.trim()}`;
      }
      response += `\nURL: ${url}`;

      // console.log(`${getTimestamp()} [INFO] Successfully scraped URL: ${url}`);
      return response;
   } catch (err) {
      console.error(`${getTimestamp()} [ERROR] Failed to scrape URL: ${url}`, err);
      return null;
   }
}

app.post('/api/v1/chat', async (req, res) => {
   const startTime = Date.now(); // Start time tracking
   const ip = req.clientIp;
   console.log(`${getTimestamp()} [INFO] Handling chat request from IP: ${ip}`); // Log the IP address

   if (isProcessing) {
      console.log(`${getTimestamp()} [WARN] System is busy processing another request`);
      return res.status(429).json({
         message: "Sorry, I am working on another request, try again later"
      });
   }

   isProcessing = true;

   try {
      let userMessage = req.body.message;
      console.log(`${getTimestamp()} [INFO] Received user message: ${userMessage}`);
      userMessage = req.body.message + `\nDate/Time:${getTimestamp()}`;

      if (!conversationHistory[ip]) {
         console.log(`${getTimestamp()} [INFO] Initializing conversation history for new IP: ${ip}`);
         conversationHistory[ip] = [{
            role: 'system',
            content: prompt
         }];
      }
      conversationHistory[ip].push({
         role: 'user',
         content: userMessage
      });

      // Trim conversation history if it exceeds the token limit
      const maxLength = 7800;
      const tolerance = 25;
      trimConversationHistory(conversationHistory[ip], maxLength, tolerance);

      // Start Plugins ---
      const ipRegex = /(\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b)/g;
      const ipAddresses = userMessage.match(ipRegex);

      const urlRegex = /(https?:\/\/[^\s]+)/g;
      const urls = userMessage.match(urlRegex);

      const newPersonRegex = /\bnew\s+person\b/gi;
      const newPersons = userMessage.match(newPersonRegex);

      const whatServersRegex = /\bwhat\s+servers\b/gi;
      const myMcRegex = /\bmy-mc\b/gi;
      const dlinuxRegex = /\bdlinux\b/gi;
      const searchRegex = /\b[Ss]earch\s+(.+)\b/;
      const searchMatch = userMessage.match(searchRegex);

      // Handle scraping and other plugins asynchronously
      const pluginTasks = [];

      if (ipAddresses) {
         console.log(`${getTimestamp()} [INFO] Detected IP addresses in user message: ${ipAddresses}`);
         ipAddresses.forEach(ipAddr => {
            pluginTasks.push(handleIPPlugin(ipAddr, ip, conversationHistory));
         });
      }

      if (urls) {
         console.log(`${getTimestamp()} [INFO] Detected URLs in user message: ${urls}`);
         urls.forEach(url => {
            pluginTasks.push(handleURLPlugin(url, ip, conversationHistory));
         });
      }

      if (newPersons) {
         console.log(`${getTimestamp()} [INFO] Detected new person request in user message`);
         newPersons.forEach(() => {
            pluginTasks.push(handleNewPersonPlugin(ip, conversationHistory));
         });
      }

      if (whatServersRegex.test(userMessage)) {
         console.log(`${getTimestamp()} [INFO] Detected what servers request in user message`);
         pluginTasks.push(handleWhatServersPlugin(ip, conversationHistory));
      }

      if (myMcRegex.test(userMessage)) {
         console.log(`${getTimestamp()} [INFO] Detected My-MC.Link request in user message`);
         pluginTasks.push(handleMyMcPlugin(ip, conversationHistory));
      }

      if (dlinuxRegex.test(userMessage)) {
         console.log(`${getTimestamp()} [INFO] Detected dlinux request in user message`);
         pluginTasks.push(handleDlinuxPlugin(ip, conversationHistory));
      }

      if (searchMatch) {
         const searchQuery = searchMatch[1];
         console.log(`${getTimestamp()} [INFO] Detected search query in user message: ${searchQuery}`);
         pluginTasks.push(handleSearchPlugin(searchQuery, ip, conversationHistory));
      }

      // Wait for all plugin tasks to complete before sending the request to llama
      await Promise.all(pluginTasks);

      console.log(`${getTimestamp()} [INFO] Sending request to llama API for response`);

      const sent = {
         model: 'model',
         messages: conversationHistory[ip]
      };

      console.log(sent);
      const llamaResponse = await fetch(`http://127.0.0.1:8002/v1/chat/completions`, {
         method: 'POST',
         headers: {
            'Content-Type': 'application/json'
         },
         body: JSON.stringify(sent)
      });

      const response = await llamaResponse.json();
      const assistantMessage = response.choices[0].message;
      conversationHistory[ip].push(assistantMessage);

      console.log(`${getTimestamp()} [INFO] Received response from llama API`);
      console.log(`${getTimestamp()} [DEBUG] Finish Reason: ${response.choices[0].finish_reason}`);
      console.log(`${getTimestamp()} [STATS] Usage: prompt_tokens=${response.usage.prompt_tokens}, completion_tokens=${response.usage.completion_tokens}, total_tokens=${response.usage.total_tokens}`);

      res.json(assistantMessage);
   } catch (error) {
      console.error(`${getTimestamp()} [ERROR] An error occurred while handling chat request`, error);
      res.status(500).json({
         message: "An error occurred",
         error: error.message
      });
   } finally {
      isProcessing = false;
      const endTime = Date.now(); // End time tracking
      const processingTime = ((endTime - startTime) / 1000).toFixed(2); // Calculate processing time in seconds
      console.log(`${getTimestamp()} [STATS] Processing Time: ${processingTime} seconds`);
      console.log(`${getTimestamp()} [INFO] Finished processing chat request for: ${ip}`);
   }
});

// Define plugin handling functions
async function handleIPPlugin(ipAddr, ip, conversationHistory) {
   try {
      const url = new URL('https://api.abuseipdb.com/api/v2/check');
      url.searchParams.append('ipAddress', ipAddr);
      url.searchParams.append('maxAgeInDays', '90');
      url.searchParams.append('verbose', '');

      const options = {
         method: 'GET',
         headers: {
            'Key': process.env.ABUSE_KEY,
            'Accept': 'application/json'
         }
      };

      const response = await fetch(url, options);
      const data = await response.json();

      let abuseResponse = `IP: ${ipAddr}\n`;
      abuseResponse += `Abuse Score: ${data.data.abuseConfidenceScore}\n`;
      abuseResponse += `Country: ${data.data.countryCode}\n`;
      abuseResponse += `Usage Type: ${data.data.usageType}\n`;
      abuseResponse += `ISP: ${data.data.isp}\n`;
      abuseResponse += `Domain: ${data.data.domain}\n`;

      if (data.data.totalReports) {
         abuseResponse += `Total Reports: ${data.data.totalReports}\n`;
         abuseResponse += `Last Reported: ${data.data.lastReportedAt}\n`;
      }

      const lastMessageIndex = conversationHistory[ip].length - 1;
      if (lastMessageIndex >= 0) {
         conversationHistory[ip][lastMessageIndex].content += "\n" + abuseResponse;
         console.log(`${getTimestamp()} [INFO] Processed IP address: ${ipAddr}, response: ${abuseResponse}`);
      } else {
         console.error(`${getTimestamp()} [ERROR] Conversation history is unexpectedly empty for: ${ip}`);
      }
   } catch (err) {
      console.error(`${getTimestamp()} [ERROR] Failed to process IP address: ${ipAddr}`, err);
   }
}

async function handleURLPlugin(url, ip, conversationHistory) {
   const scrapedContent = await scrapeWebPage(url);
   if (scrapedContent) {
      conversationHistory[ip].push({
         role: 'assistant',
         content: scrapedContent
      });
      console.log(`${getTimestamp()} [INFO] Added scraped content to conversation history for: ${ip}`);
   }
}

async function handleNewPersonPlugin(ip, conversationHistory) {
   try {
      const randomUser = await fetchRandomUser();

      if (randomUser) {
         let response = `New Person:\n`;
         response += `Name: ${randomUser.name.first} ${randomUser.name.last}\n`;
         response += `Gender: ${randomUser.gender}\n`;
         response += `Location: ${randomUser.location.city}, ${randomUser.location.state}, ${randomUser.location.country}\n`;
         response += `Email: ${randomUser.email}\n`;
         response += `Phone: ${randomUser.phone}\n`;
         // Add more details as needed

         conversationHistory[ip].push({
            role: 'assistant',
            content: response
         });

         console.log("A request for a new person was made. Response: " + response);
      } else {
         console.log('Failed to fetch random user.');
      }
   } catch (err) {
      console.error(err);
   }
}

async function handleWhatServersPlugin(ip, conversationHistory) {
   try {
      const response = await fetch(`https://api.my-mc.link/list_all_servers/${process.env.PATH_KEY}/`, {
         headers: {
            'x-my-mc-auth': process.env.API_KEY
         }
      });

      const data = await response.json();

      if (data.success) {
         let responseMessage = `Information provided by the system not me: The Current Minecraft Servers online within the My-MC.link P2P JUMP Node System are listed below. These servers are also listed within My-MC Realms.`;

         for (const server of data.servers) {
            responseMessage += `\nName: ${server.serverName}\n`;
            responseMessage += `Game Version: ${server.gameVersion}\n`;
            responseMessage += `MOTD: ${server.motd}\n`;
            responseMessage += `Online: ${server.online}\n`;
            // Add more details as needed
         }

         conversationHistory[ip].push({
            role: 'assistant',
            content: responseMessage
         });

         console.log("A request for server information was made. Response: " + responseMessage);
      } else {
         console.log('Failed to fetch server information.');
      }
   } catch (error) {
      console.error('Error fetching server information:', error);
   }
}

async function handleMyMcPlugin(ip, conversationHistory) {
   try {
      const response = await fetch('https://my-mc.link/wiki.json');
      const data = await response.json();

      if (data) {
         let responseMessage = `My-MC.Link Wiki:\n`;

         // Append all data fields to the response message
         Object.keys(data).forEach(key => {
            if (typeof data[key] === 'object') {
               // Handle nested objects
               responseMessage += `${key}:\n`;
               Object.keys(data[key]).forEach(innerKey => {
                  responseMessage += `${innerKey}: ${data[key][innerKey]}\n`;
               });
            } else {
               responseMessage += `${key}: ${data[key]}\n`;
            }
         });

         conversationHistory[ip].push({
            role: 'assistant',
            content: responseMessage
         });
         console.log("A request for My-MC.Link wiki information was made.");
      } else {
         console.log('Failed to fetch My-MC.Link wiki information.');
      }
   } catch (error) {
      console.error('Error fetching My-MC.Link wiki information:', error);
   }
}

async function handleDlinuxPlugin(ip, conversationHistory) {
   try {
      const response = await fetch('https://my-mc.link/dwiki.json');
      const data = await response.json();

      if (data) {
         let responseMessage = `dlinux Wiki:\n`;

         // Append all data fields to the response message
         Object.keys(data).forEach(key => {
            if (typeof data[key] === 'object') {
               // Handle nested objects
               responseMessage += `${key}:\n`;
               Object.keys(data[key]).forEach(innerKey => {
                  responseMessage += `${innerKey}: ${data[key][innerKey]}\n`;
               });
            } else {
               responseMessage += `${key}: ${data[key]}\n`;
            }
         });

         conversationHistory[ip].push({
            role: 'assistant',
            content: responseMessage
         });

         console.log("A request for dlinux wiki information was made.");
      } else {
         console.log('Failed to fetch dlinux wiki information.');
      }
   } catch (error) {
      console.error('Error fetching dlinux wiki information:', error);
   }
}

async function handleSearchPlugin(searchQuery, ip, conversationHistory) {
   const options = {
      query: searchQuery,
      limit: 5,
      disableConsole: true
   };

   try {
      const results = await googleIt(options);
      let searchResponse = `Search Query: ${searchQuery}\n`;
      searchResponse += `Top Google search results:\n`;
      let scrapedContent = '';

      for (let i = 0; i < results.length; i++) {
         const result = results[i];
         searchResponse += `${i + 1}. ${result.title} - ${result.link}\n`;

         try {
            const scrapeResult = await scrapeWebPage(result.link, 400);
            searchResponse += `Scraped Data: ${scrapeResult}\n`;
            scrapedContent += `Scraped Data from ${result.link}:\n${scrapeResult}\n`;
         } catch (scrapeErr) {
            console.error(`${getTimestamp()} [ERROR] Failed to scrape URL: ${result.link}`, scrapeErr);
            searchResponse += `Failed to scrape URL: ${result.link}\n`;
            scrapedContent += `Failed to scrape URL: ${result.link}\n`;
         }
      }

      const lastMessageIndex = conversationHistory[ip].length - 1;
      if (lastMessageIndex >= 0) {
         conversationHistory[ip][lastMessageIndex].content += "\nYou scraped these results, review them provide links: " + searchResponse;
         console.log(`${getTimestamp()} [INFO] Processed search query: ${searchQuery}, response: ${searchResponse}`);
      } else {
         console.error(`${getTimestamp()} [ERROR] Conversation history is unexpectedly empty for: ${ip}`);
      }

      if (scrapedContent) {
         conversationHistory[ip].push({
            role: 'assistant',
            content: scrapedContent
         });
         console.log(`${getTimestamp()} [INFO] Added scraped content to conversation history for: ${ip}`);
      }
   } catch (err) {
      console.error(`${getTimestamp()} [ERROR] Failed to perform Google search: ${searchQuery}`, err);
   }
}

app.get('/api/v1/conversation-history', (req, res) => {
   const ip = req.clientIp;
   console.log(`${getTimestamp()} [INFO] Fetching conversation history for: ${ip}`); // Log the IP address
   res.json(conversationHistory[ip]);
});

app.post('/api/v1/restart-core', (req, res) => {
   console.log(`${getTimestamp()} [INFO] Restarting core service`);
   cmd(`docker restart llama-gpu-server`).then(out => {
      console.log(`${getTimestamp()} [INFO] Core service restarted`);
      res.json(out.stdout);
   }).catch(err => {
      console.error(`${getTimestamp()} [ERROR] Failed to restart core service`, err);
      res.status(500).json({
         message: "An error occurred while restarting the core service",
         error: err.message
      });
   });
});

app.post('/api/v1/reset-conversation', (req, res) => {
   const ip = req.clientIp;
   console.log(`${getTimestamp()} [INFO] Resetting conversation history for: ${ip}`); // Log the IP address

   conversationHistory[ip] = [{
      role: 'system',
      content: prompt
   }];
   console.log(`${getTimestamp()} [INFO] Conversation history reset for: ${ip}`);
   res.json({
      message: "Conversation history reset for: " + ip
   });
});

app.listen(port, () => {
   console.log(`${getTimestamp()} [INFO] Server running at http://localhost:${port}`);
});