const { Configuration, OpenAIApi } = require("openai");
const express = require('express')
const bodyParser = require('body-parser')
const cookieParser = require("cookie-parser")
const cors = require('cors')
const { createClient } = require('redis')
require('dotenv').config()
const rateLimit = require('express-rate-limit')
const fetch = require('node-fetch');
const anchorme = require("anchorme").default;
const axios = require('axios');
const { encodingForModel } = require('js-tiktoken');
const tiktokenModels = [
  'text-davinci-003',
  'text-davinci-002',
  'text-davinci-001',
  'text-curie-001',
  'text-babbage-001',
  'text-ada-001',
  'davinci',
  'curie',
  'babbage',
  'ada',
  'code-davinci-002',
  'code-davinci-001',
  'code-cushman-002',
  'code-cushman-001',
  'davinci-codex',
  'cushman-codex',
  'text-davinci-edit-001',
  'code-davinci-edit-001',
  'text-embedding-ada-002',
  'text-similarity-davinci-001',
  'text-similarity-curie-001',
  'text-similarity-babbage-001',
  'text-similarity-ada-001',
  'text-search-davinci-doc-001',
  'text-search-curie-doc-001',
  'text-search-babbage-doc-001',
  'text-search-ada-doc-001',
  'code-search-babbage-code-001',
  'code-search-ada-code-001',
  'gpt2',
  'gpt-4',
  'gpt-4-0314',
  'gpt-4-32k',
  'gpt-4-32k-0314',
  'gpt-3.5-turbo',
  'gpt-3.5-turbo-0301'
];

let client;
let filteredModels = {};
const allowedEndpoints = ["openAI", "Opensource", "Llama"];
const excludedModels = ["o1-preview", "o1-mini"];

// Open AI Configuration
// console.log(process.env.OPENAI_API_ORG)
const configuration = new Configuration({
  organization: process.env.OPENAI_API_ORG,
  apiKey: process.env.OPENAI_API_KEY,
});
const openai = new OpenAIApi(configuration);

const rateLimiter = rateLimit({
  windowMs: 1000 * 60 * 1, // 1 minute (refreshTime)
  max: 3000, // limit each IP to x requests per windowMs (refreshTime)
  message: 'Sorry, too many requests. Please try again in a bit!',
});

// Express Configuration
const app = express()
const port = 3080

app.use(bodyParser.json())
app.use(cors())
app.use(require('morgan')('dev'))
app.use(rateLimiter)
app.use(cookieParser());

const max_tokens = process.env.MAX_TOKENS_chatbot_plus ? parseInt(process.env.MAX_TOKENS_chatbot_plus) : 512;
// Routing
const hostapi = process.env.API_URL || "https://api.ai-pro.org";
const user_secret_id = process.env.USER_SECRET_ID || "aiwp_logged_in";
const aiwp_app_id = "chatbot+";
// Primary Open AI Route
app.post('/api', async (req, res) => {
  // if(!req.get('origin') || (!req.get('origin').includes(req.get('host')))) {
  //   res.status(401);
  //   res.send('Method Not Allowed');
  //   return;
  // }
  const { message, currentModel, temperature } = req.body;
  const getEndpoint = (modelName) => {
      const model = Object.values(filteredModels).find(m => m.model === modelName);
      return model ? model.endpoint : 'Endpoint not found';
  };

  const endpoint = getEndpoint(currentModel);

  if (endpoint == "openAI") {
    runGPTTurbo(req, res);
    return;
  } 
  
  if (endpoint == "Llama" || endpoint == "Opensource") {
    runOpensource(req, res);
    return;
  }

  const validate = await validation(aiwp_app_id, req, res);
  if(!validate) return;
  const { IS_FREE_USER, aiwp_logged_in, TRIED_USAGE} = validate;

  let greetingPrompt = 'Hello, how can I assist you?'
  const greetings = ['hi', 'hello', 'hey']
  if (greetings.some((greeting) => message.toLowerCase().includes(greeting))) {
    greetingPrompt = 'Hello, how can I help you today?'
  }
  let query_prompt = `${greetingPrompt}\n${message}`;
  str_length = req.body.message.split(' ').length;
  if (str_length >= 800) {
    arr_body = req.body.message.split("\n");
    if (arr_body.length >= 4) {
      var i = arr_body.length - 2
      while (i--) {
        arr_body.splice(i, 1);
      }
      query_prompt = arr_body.join("\n")
    }
  }
  const moderation = await axios.post("https://api.openai.com/v1/moderations", {
    input: query_prompt
  }, { headers: { 'content-type': 'application/json', 'Authorization': `Bearer ${process.env.OPENAI_API_KEY}` } });

  if (moderation.data.results[0].flagged) {
    res.json({
      success: false,
      message: "I'm sorry, but I can't assist with that. We want everyone to use our tool safely and responsibly.\nIf you have any other questions or need advice on a different topic, feel free to ask."
    });
    res.end();
    return;
  }

  try {
    const response = await openai.createCompletion({
      model: `${currentModel}`,// "text-davinci-003",
      prompt: query_prompt,
      max_tokens: max_tokens,
      temperature,
    });
    let input = response.data.choices[0].text;
    let usage = {};
    let enc = null;
    try {
      enc = encodingForModel(tiktokenModels.includes(currentModel) ? currentModel : 'gpt-3.5-turbo');
      usage.prompt_tokens = (enc.encode(query_prompt)).length;
      usage.completion_tokens = (enc.encode(input)).length;
      usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;

    } catch (e) {
      console.log('Error encoding prompt text', e);
    }
    let usage_params = {
      "aiwp_logged_in": aiwp_logged_in, "app": "chatbot+", "prompt_token": usage.prompt_tokens, "total_token": usage.total_tokens, "aiwp_app_id":aiwp_app_id , "usage_tries": TRIED_USAGE, ipdflu: getClientIP(req)
    };
    if(IS_FREE_USER) {
			await setUsage(usage_params);
		} else {
      await setChatUsage(usage_params);
    }
    res.json({
      usage: usage,
      message: anchorme({
        input,
        options: {
          attributes: {
            target: "_blank"
          },
        }
      })
    })
  } catch (e) {
    let error_msg = e.response.data.error.message ? e.response.data.error.message : '';
    if (error_msg.indexOf('maximum context length') >= 0) {
      res.json({
        message: "The output for your prompt is too long for us to process. Please reduce your prompt and try again.",
      })
    } else {
      // console.log(e.response);
    }
  } finally {
    // console.log('We do cleanup here');
  }
});

async function runGPTTurbo(req, res) {
  // "gpt-3.5-turbo"
  const { message, currentModel, temperature } = req.body;
  var input = '';
  const message_history = JSON.parse(message);
  const query_prompt = message_history.length ? message_history[message_history.length - 1].content : "";
  const moderation = await axios.post("https://api.openai.com/v1/moderations", {
    input: query_prompt
  }, { headers: { 'content-type': 'application/json', 'Authorization': `Bearer ${process.env.OPENAI_API_KEY}` } });

  const validate = await validation(aiwp_app_id, req, res);
  if(!validate) return;
  const { IS_FREE_USER, aiwp_logged_in, TRIED_USAGE} = validate;

  if (moderation.data.results[0].flagged) {
    res.json({
      success: false,
      message: "I'm sorry, but I can't assist with that. We want everyone to use our tool safely and responsibly.\nIf you have any other questions or need advice on a different topic, feel free to ask."
    });
    res.end();
    return;
  }
  try {
    const response = await openai.createChatCompletion({
      model: `${currentModel}`,
      messages: JSON.parse(message),
      max_tokens: max_tokens,
      temperature
    });
    input = response.data.choices[0].message.content
  } catch (e) {
    let error_msg = e.response.data.error.message ? e.response.data.error.message : '';
    if (error_msg.indexOf('maximum context length') >= 0) {
      input = "The output for your prompt is too long for us to process. Please reduce your prompt and try again.";
    } else {
      // console.log(e.response);
    }
  } finally {

    let usage = {};
    let enc = null;
    try {
      enc = encodingForModel(tiktokenModels.includes(currentModel) ? currentModel : 'gpt-3.5-turbo');
      usage.prompt_tokens = (enc.encode(query_prompt)).length;
      usage.completion_tokens = (enc.encode(input)).length;
      usage.total_tokens = usage.prompt_tokens + usage.completion_tokens;
    } catch (e) {
      console.log('Error encoding prompt text', e);
    }

    let usage_params = {
      "aiwp_logged_in": aiwp_logged_in,
      "app": "chatbot+",
      "prompt_token": usage.prompt_tokens,
      "total_token": usage.total_tokens,
      "aiwp_app_id": aiwp_app_id,
      "usage_tries": TRIED_USAGE,
      ipdflu: getClientIP(req)
    };

    if (IS_FREE_USER) {
			await setUsage(usage_params);
		} else {
      await setChatUsage(usage_params);
    }
    res.json({
      prompt: JSON.parse(message),
      usage: usage,
      message: anchorme({
        input,
        options: {
          attributes: {
            target: "_blank"
          },
        }
      })
    });
    return;
  }
}

async function runOpensource(req, res) {
  const { message, currentModel, temperature, P6XcW47o: together_ai_response = null } = req.body;
  var input = '';
  const message_history = JSON.parse(message);
  const query_prompt = message_history.length ? message_history[message_history.length - 1].content : "";

  const validate = await validation(aiwp_app_id, req, res);
  if(!validate) return;
  const { IS_FREE_USER, aiwp_logged_in, TRIED_USAGE} = validate;
  let usage = {};
  let result_response = {}
  let together_ai = null
  
  try {
    let error_msg = "";
    const endpoint_api_url = process.env.OPENSOURCE_ENDPOINTS;
    const api_key = OPENSOURCE_API_KEY;
    const response = await axios.post(endpoint_api_url, {
      model: currentModel,
      messages: JSON.parse(message),
      max_tokens: 2048,
      temperature,
      top_p: 0.7,
      top_k: 50,
      repetition_penalty: 1
    }, {
      headers: {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ' + api_key
      },
    }).catch(error => {
      if(error.response?.data?.error?.param === 'max_tokens') {
        input = "The output for your prompt is too long for us to process. Please reduce your prompt and try again.";
      } else {
        error_msg = error.response.statusText ? error.response.statusText : '';
      }
      console.log("err",error.response.data.error);
    });

    if (error_msg !== '') {
      input = "Error: " + error_msg;
    } else {
      input = response.data.choices[0].message.content
    }

    usage = response.data.usage

    if (together_ai_response) {
      together_ai = response.data
    }
  } catch (e) {
    let error_msg = e.response.data.error.message ? e.response.data.error.message : '';
    if (error_msg.indexOf('maximum context length') >= 0) {
      input = "The output for your prompt is too long for us to process. Please reduce your prompt and try again.";
    } else {
      // console.log(e.response);
    }
  } finally {
    let usage_params = {
      "aiwp_logged_in": aiwp_logged_in,
      "app": "chatbot+",
      "prompt_token": usage.prompt_tokens,
      "total_token": usage.total_tokens,
      "aiwp_app_id": aiwp_app_id,
      "usage_tries": TRIED_USAGE,
      ipdflu: getClientIP(req)
    };

    if (IS_FREE_USER) {
			await setUsage(usage_params);
		} else {
      await setChatUsage(usage_params);
    }

    result_response = {
      prompt: JSON.parse(message),
      usage: usage,
      message: anchorme({
        input,
        options: {
          attributes: {
            target: "_blank"
          },
        }
      })
    }

    if (together_ai !== null) {
      result_response['together_ai'] = together_ai
    }

    res.json(result_response);
    return;
  }
}

async function authenticate(params) {

	let data = await fetch(`${hostapi}/e/authenticate/v2`, {
			method: "POST",
			headers: {
					"Content-Type": "application/json"
			},
			body: JSON.stringify(params),
			referrer: "https://chatgpt.ai-pro.org"
	});

	return await data.json();
}

async function getLimitedUsage(params) {

	let data = await fetch(`${hostapi}/e/get-usage`, {
			method: "POST",
			headers: {
					"Content-Type": "application/json"
			},
			body: JSON.stringify(params),
			referrer: "https://chatgpt.ai-pro.org"
	});

	return await data.json();
}
async function getUsage(params) {

	let data = await fetch(`${hostapi}/e/get-chat-usage`, {
			method: "POST",
			headers: {
					"Content-Type": "application/json"
			},
			body: JSON.stringify(params),
			referrer: "https://chatgpt.ai-pro.org"
	});

	return await data.json();
}
async function setUsage(params) {

	fetch(`${hostapi}/e/set-usage`, {
			method: "POST",
			headers: {
					"Content-Type": "application/json"
			},
			body: JSON.stringify(params),
			referrer: "https://chatgpt.ai-pro.org"
	});
}
async function setChatUsage(params) {

	fetch(`${hostapi}/e/set-chat-usage`, {
			method: "POST",
			headers: {
					"Content-Type": "application/json"
			},
			body: JSON.stringify(params),
			referrer: "https://chatgpt.ai-pro.org"
	});
}

async function validation (aiwp_app_id, req, res) {
    const aiwp_logged_in = req.cookies[user_secret_id] ? decodeURIComponent(req.cookies[user_secret_id]) : "";
    const limit = req.cookies["WcvYPABR"] ? parseInt(req.cookies["WcvYPABR"].replace(/\D/g, '')) : 9999999999999999999;
    let IS_FREE_USER = false;
    let TRIED_USAGE = 0;
    const ip_address = getClientIP(req);

    if (aiwp_logged_in) {
        let auth = await authenticate({ "aiwp_logged_in": aiwp_logged_in, "user_event_data": {}, "user_event": "endpoint" });

        if (!auth.success) {
          IS_FREE_USER = true;
          if (auth.is_restrict) {
              res.json({ status: "invalid", restrict: true, redirect: auth.redirect });
              res.end();
              return false;
          } else if (auth.subscription_type &&
            typeof auth.has_pro_access === "undefined" && !auth.has_pro_access) {
              res.json({ status: "invalid", restrict: true });
              res.end();
              return false;
          }
        } else {
            let data = await getUsage({
                aiwp_logged_in, app: 'chatbot+'
            });

            if (!(data.success === 1 && data.status === 'valid')) {
                res.json({ status: "invalid", data });
                res.status(200);
                return false;
            }
        }
    } else {
      IS_FREE_USER = true;
    }
    if(IS_FREE_USER) {
      let data = await getLimitedUsage({
        "aiwp_app_id": aiwp_app_id, ipdflu: ip_address
      });

      if (data.usage !== null) {
        TRIED_USAGE = parseInt(data.usage);
      }
      if (TRIED_USAGE >= limit) {
          res.json({ status: "invalid", limited: true });
          res.end();
          return false;
      }
      TRIED_USAGE++;
    }

		return { IS_FREE_USER, aiwp_logged_in, TRIED_USAGE };
};

const getClientIP = (req) => {
  let ipAddress = '';

  // Order of checking headers for client IP
  if (req.headers['x-forwarded-for']) {
      ipAddress = Array.isArray(req.headers['x-forwarded-for'])
          ? req.headers['x-forwarded-for'][0]
          : req.headers['x-forwarded-for'];
  } else if (req.headers['x-client-ip']) {
      ipAddress = req.headers['x-client-ip'];
  } else if (req.headers['x-cluster-client-ip']) {
      ipAddress = req.headers['x-cluster-client-ip'];
  } else if (req.headers['x-forwarded']) {
      ipAddress = req.headers['x-forwarded'];
  } else if (req.headers['forwarded-for']) {
      ipAddress = req.headers['forwarded-for'];
  } else if (req.headers['forwarded']) {
      ipAddress = req.headers['forwarded'];
  } else {
      ipAddress = req.socket?.remoteAddress || '127.0.0.2';
  }

  // Split the IP address in case of multiple entries in 'x-forwarded-for'
  const arrIp = ipAddress.split(',');
  return arrIp[0] || '127.0.0.2';
};




// Get Models Route
app.get('/models', async (req, res) => {
  await getModels() 
  const models = {
    data: Object.values(filteredModels).map((model) => ({
        id: model.model,
        label: model.name,
        name: model.name,
        beta: model.isBeta,
        endpoint: model.endpoint
    }))
  };

  res.json({
    models
  })
});

// Start the server
app.listen(port, () => {
  console.log(`Example app listening at http://localhost:${port}`)
});

const initializeRedisClient = () => {
  const newClient = createClient({
    socket: {
      host: process.env.REDIS_HOST,
      port: process.env.REDIS_PORT,
    },
    username: process.env.REDIS_USER,
    password: process.env.REDIS_PASS,
  });

  newClient.on('error', (err) => {
    console.error('Redis error:', err);
  });

  return newClient;
};

const fetchAndFilterModels = async (client) => {
  try {
    await client.select(process.env.REDIS_DB);
    const models = await client.get('model');
    return Object.entries(JSON.parse(models))
      .filter(([_, value]) => allowedEndpoints.includes(value.endpoint))
      .filter(([_, value]) => !excludedModels.includes(value.model))
      .reduce((acc, [key, value]) => {
        acc[key] = value;
        return acc;
      }, {});
  } catch (err) {
    console.error('Error fetching and filtering models:', err);
    throw err;
  }
};

const getModels = async () => {
  if (!client) {
    client = await initializeRedisClient();
    await client.connect();

    try {
      console.log('Connected to Redis successfully');
    } catch (err) {
      console.error('Error connecting to Redis:', err);
      throw err;
    }
  }
  if (Object.keys(filteredModels).length === 0) {
    filteredModels = await fetchAndFilterModels(client);
  }
};