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 fetch from 'node-fetch';
import Groq from 'groq-sdk';
import googleIt from 'google-it';

// Constants and initialization
const app = express();
const port = 3000;
const prompt = process.env.PROMPT;

const groq = new Groq({ apiKey: process.env.GROQ });

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]) {
      console.log(`${getTimestamp()} [INFO] Initializing conversation history for: ${req.clientIp}`);
      conversationHistory[req.clientIp] = [{
         role: 'system',
         content: `My name is: ${name || 'Unknown'}, my Discord ID is: ${req.clientIp}.` + (guild ? ` We are chatting inside ${guild} a Discord Server.` : '') + 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(ip, maxLength = 14000, tolerance = 25) {
   const messages = conversationHistory[ip];
   let totalTokens = countLlamaTokens(messages);

   while (totalTokens > maxLength + tolerance && messages.length > 1) {
      messages.shift(); // Remove the oldest messages first
      totalTokens = countLlamaTokens(messages);
   }
}

// Function to scrape web page
async function scrapeWebPage(url, length = 2000) {
   try {
      const res = await fetch(url);
      const html = await res.text();
      const $ = cheerio.load(html);

      const pageTitle = $('head title').text().trim();
      const pageDescription = $('head meta[name="description"]').attr('content');
      let plainTextContent = $('body').text().trim().replace(/[\r\n\t]+/g, ' ');

      if (plainTextContent.length > length) {
         plainTextContent = plainTextContent.substring(0, length) + '...';
      }

      return `Title: ${pageTitle}\nDescription: ${pageDescription || 'N/A'}\nContent: ${plainTextContent}\nURL: ${url}`;
   } catch (err) {
      console.error(`${getTimestamp()} [ERROR] Failed to scrape URL: ${url}`, err);
      return null;
   }
}

// Function to handle IP plugin
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);
   }
}

// Function to handle What Servers plugin
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 = 'The Current Minecraft Servers online within the My-MC.link P2P JUMP Node System are listed below:\n';

         for (const server of data.servers) {
            responseMessage += `Name: ${server.serverName}\nGame Version: ${server.gameVersion}\nMOTD: ${server.motd}\nOnline: ${server.online}\n`;
         }

         conversationHistory[ip].push({ role: 'assistant', content: responseMessage });
         console.log(`${getTimestamp()} [INFO] Processed server information request.`);
      } else {
         console.error(`${getTimestamp()} [ERROR] Failed to fetch server information.`);
      }
   } catch (error) {
      console.error(`${getTimestamp()} [ERROR] Failed to fetch server information: `, error);
   }
}

// Function to handle My-MC.Plugin
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';
         for (const key in data) {
            responseMessage += `${key}: ${data[key]}\n`;
         }

         conversationHistory[ip].push({ role: 'assistant', content: responseMessage });
         console.log(`${getTimestamp()} [INFO] Processed My-MC.Link wiki request.`);
      } else {
         console.error(`${getTimestamp()} [ERROR] Failed to fetch My-MC.Link wiki information.`);
      }
   } catch (error) {
      console.error(`${getTimestamp()} [ERROR] Error fetching My-MC.Link wiki: `, error);
   }
}

// Search plugin function
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}\nTop 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, 800);
            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, generate a detailed report, Also provide the list of the scraped URLs in your report. \n" + searchResponse;
         console.log(`${getTimestamp()} [INFO] Processed search query: ${searchQuery}.`);
      } 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);
   }
}

// Main chat handler
app.post('/api/v1/chat', async (req, res) => {
   const startTime = Date.now();
   const ip = req.clientIp;

   isProcessing = true;

   try {
      const userMessage = req.body.message + `\nDate/Time: ${getTimestamp()}`;
      conversationHistory[ip].push({ role: 'user', content: userMessage });
      trimConversationHistory(ip);

      const pluginTasks = [];
      const processedIPs = new Set();  // To avoid duplicate IP processing
      const processedURLs = new Set(); // To avoid duplicate URL processing

      // Check for IPs in user message and process them
      const ipRegex = /(\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b)/g;
      const ipAddresses = userMessage.match(ipRegex);
      if (ipAddresses) {
         for (const ipAddr of ipAddresses) {
            if (!processedIPs.has(ipAddr)) {
               pluginTasks.push(handleIPPlugin(ipAddr, ip, conversationHistory));
               processedIPs.add(ipAddr);  // Mark IP as processed
            }
         }
      }

      // Check for URLs and scrape them
      const urlRegex = /(https?:\/\/[^\s]+)/g;
      const urls = userMessage.match(urlRegex);
      if (urls) {
         for (const url of urls) {
            if (!processedURLs.has(url)) {
               pluginTasks.push(scrapeWebPage(url).then(content => {
                  if (content) {
                     conversationHistory[ip].push({ role: 'assistant', content });
                  }
               }));
               processedURLs.add(url);  // Mark URL as processed
            }
         }
      }

      // Search Plugin
      const searchRegex = /\b[Ss]earch\s+(.+)\b/;
      const searchMatch = userMessage.match(searchRegex);
      if (searchMatch) {
         const searchQuery = searchMatch[1];
         console.log(`${getTimestamp()} [INFO] Detected search query in user message: ${searchQuery}`);
         pluginTasks.push(handleSearchPlugin(searchQuery, ip, conversationHistory));
      }

      await Promise.all(pluginTasks);

      const completion = await groq.chat.completions.create({
         messages: conversationHistory[ip],
         model: "llama-3.2-3b-preview"
      });

      const assistantMessage = completion.choices[0].message.content;
      conversationHistory[ip].push({ role: 'assistant', content: assistantMessage });

      res.json(assistantMessage);
   } catch (error) {
      console.error(`${getTimestamp()} [ERROR] An error occurred: `, error);
      res.status(500).json({ message: "An error occurred", error: error.message });
   } finally {
      isProcessing = false;
      const endTime = Date.now();
      const processingTime = ((endTime - startTime) / 1000).toFixed(2);
      console.log(`${getTimestamp()} [STATS] Processing Time: ${processingTime} seconds`);
   }
});

// Restart core service
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
      });
   });
});

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

   conversationHistory[ip] = [{
      role: 'system',
      content: prompt
   }];

   console.log(`${getTimestamp()} [INFO] Conversation history reset for: ${ip}`);
   res.json({ message: "Conversation history reset for: " + ip });
});

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

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