const { Configuration, OpenAIApi } = require("openai");
const express = require('express')
const bodyParser = require('body-parser')
const cookieParser = require("cookie-parser")
const cors = require('cors')
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'
];

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

  if (currentModel == "gpt-3.5-turbo" || currentModel == "gpt-3.5-turbo-0301") {
    runGPTTurbo(req, res);

    return;
  }

  if (currentModel == "openchat_3.5" || currentModel == "zephyr-7B-beta"
    || currentModel == "google/gemma-2-9b-it" || currentModel == "meta-llama/Llama-3-8b-chat-hf"
  ) {
    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
    };
    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
    };
    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;
  }
}

const get_endpoint_api_url = (currentModel) => {
  const OPENSOURCE_ENDPOINTS = process.env.OPENSOURCE_ENDPOINTS;
  const endpoints = JSON.parse(OPENSOURCE_ENDPOINTS);
  const endpoint_api_url = endpoints?.[currentModel];
  return endpoint_api_url
}
const get_endpoint_api_key = (currentModel) => {
  const OPENSOURCE_API_KEY = process.env.OPENSOURCE_API_KEY;
  const api_keys = JSON.parse(OPENSOURCE_API_KEY);
  const key = api_keys?.[currentModel];
  return key
}
async function runOpensource(req, res) {
  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 validate = await validation(aiwp_app_id, req, res);
  if(!validate) return;
  const { IS_FREE_USER, aiwp_logged_in, TRIED_USAGE} = validate;

  try {
    let error_msg = "";
    const endpoint_api_url = get_endpoint_api_url(currentModel);
    const api_key = get_endpoint_api_key(currentModel);
    const response = await axios.post(endpoint_api_url + '/chat/completions', {
      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
    }

  } 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('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
    };
    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 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, '')) : 3;
    let IS_FREE_USER = false;
    let TRIED_USAGE = 0;

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

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



// Get Models Route
app.get('/models', async (req, res) => {
  const openai_models = process.env.OPENAI_MODELS ? JSON.parse(process.env.OPENAI_MODELS) : [{"value": "gpt-3.5-turbo", "label": "GPT-3.5"}];
  const opensource_models = process.env.OPENSOURCE_MODELS ? JSON.parse(process.env.OPENSOURCE_MODELS) : [];

  const models = {
    data: []
  };

  openai_models.forEach((model) => {
    models.data.push({
      id: model.value,
      label: model.label,
      name: model.label,
      beta: false,
    });
  })

  opensource_models.forEach((model) => {
    models.data.push({
      id: model.value,
      label: model.label,
      name: model.label,
      beta: true,
    });
  })

  res.json({
    models
  })
});

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