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';
import Groq from 'groq-sdk';

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

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

//    if (isProcessing) {
//       return res.status(429).json({ message: "System is busy processing another request, try again later" });
//    }

   isProcessing = true;

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

      const pluginTasks = [];

      // 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) {
            pluginTasks.push(handleIPPlugin(ipAddr, ip, conversationHistory));
         }
      }

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

      await Promise.all(pluginTasks);

      const completion = await groq.chat.completions.create({
         messages: conversationHistory[ip],
         //model: "llama3-8b-8192"
         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`);
   }
});

// Endpoint to 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
      });
   });
});

// Endpoint to 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 purposes
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}`);
});