import crypto from "crypto";
import axios from "axios";
import { minutesToMilliseconds } from 'date-fns';
import { URL } from 'url';
import { config } from "./config";
import { logger } from "./logger";

const CHECK_INTERVAL = minutesToMilliseconds(3);

/** Runtime information about a proxy. */
export type Proxy = {
  url: string;
  isGpt4: boolean;
  /** Whether this proxy is currently disabled. We set this if we get a 429 or 401 response from OpenAI. */
  isDisabled?: boolean;
  /** The number of prompts that have been sent with this proxy. */
  promptCount: number;
  /** The time at which this proxy was last used. */
  lastUsed: number;
  /** Proxy hash for displaying usage in the dashboard. */
  hash: string;
  remaining: number;
  isLogging?: boolean;
};

const proxyPool: Proxy[] = [];

async function init() {
  setTimeout(checker, CHECK_INTERVAL);

  const proxiestring = config.openaiProxy;

  const proxyList = parse(proxiestring);
  for (const proxy of proxyList) {
    await add(proxy);
  }
}

async function checker() {
  for (const proxy of proxyPool) {
    await check(proxy, true);
  }

  setTimeout(checker, CHECK_INTERVAL);
}

function list() {
  return proxyPool.map((proxy) => ({
    ...proxy,
    url: undefined,
  }));
}

function parse(proxiestring: string | undefined): string[] {
  if (!proxiestring) {
    return [];
  }

  try {
    const decoded = Buffer.from(proxiestring, "base64").toString();
    return JSON.parse(decoded);
  } catch (err) {
    logger.info("Proxy is not base64-encoded JSON, assuming bare url");
    return [proxiestring];
  }
}

async function check(proxy: Proxy, silent = false) {
  let isDisabled = false;
  let isGpt4 = false;
  let rateLimit = 0;
  let isLogging = false;
  try {
    const href = new URL(proxy.url);
    href.pathname = '';
    href.search = '';
    const { data } = await axios.get(href.toString());

    const startIdx = data.indexOf('<pre>{');
    const endIdx = data.indexOf('}</pre>');
    if (startIdx >= 0 && endIdx >= 0) {
      const slice = data.slice(startIdx + 5, endIdx + 1);
      try {
        const config = JSON.parse(slice);
        if (config['gpt-4']?.active > 0 && config['gpt-4']?.remaining !== '0%') {
          isGpt4 = true;
        }
        if (config.keys?.gpt4 > 0 && config.keys?.quotaLeft !== '0%') {
          isGpt4 = true;
        }
        if (config.keyInfo?.gpt4 > 0 && config.keyInfo?.quotaLeft?.gpt4 !== '0%') {
          isGpt4 = true;
        }
        if (config.keys?.active === 0) {
          isDisabled = true;
        }
        if (config.keyInfo?.active === 0) {
          isDisabled = true;
        }
        if (config.config?.proxyKey) {
          isDisabled = true;
        }
        if (config.config?.promptLogging === 'true') {
          isLogging = true;
        }
        const quotaLeft = isGpt4
          ? parseInt(config.keys?.quotaLeft.gpt4 || config.keyInfo?.quotaLeft?.gpt4 || config['gpt-4']?.remaining)
          : parseInt(config.keys?.quotaLeft.all || config.keyInfo?.quotaLeft?.all || config['gpt-3.5-turbo']?.remaining);

        if (!isNaN(quotaLeft)) {
          proxy.remaining = quotaLeft;
        }

        rateLimit = +config.config?.modelRateLimit || 0;
      } catch (e) { }
    }

    if (!isDisabled && rateLimit > 2 && !isLogging) {
      try {
        await axios.post(
          `${proxy.url}/v1/chat/completions`,
          {
            model: "gpt-3.5-turbo", max_tokens: 1,
            messages: [{ role: "user", content: "" }],
          },
          { headers: { "Content-Type": "application/json" } },
        );
      } catch (e) {
        if (!silent) logger.error(e, `Proxy is not active ${proxy.url}`);
        isDisabled = true;
      }
    }

    if (!isDisabled && isGpt4 && rateLimit > 2 && !isLogging) {
      try {
        await axios.post(
          `${proxy.url}/v1/chat/completions`,
          {
            model: "gpt-4", max_tokens: 1,
            messages: [{ role: "user", content: "" }],
          },
          { headers: { "Content-Type": "application/json" } },
        );
      } catch (e) {
        if (!silent) logger.error(e, `Proxy is not GPT-4 ${proxy.url}`);
        isGpt4 = false;
      }
    }
  } catch (e) {
    isDisabled = true;
    if (!silent) logger.error(e, `Error checking proxy ${proxy.url}`);
  }

  proxy.isDisabled = isDisabled;
  proxy.isGpt4 = isGpt4;
  proxy.isLogging = isLogging;
}

async function add(url: string) {
  if (!url) {
    return;
  }
  if (proxyPool.some(k => k.url === url)) {
    logger.warn('Proxy already exists');
    return;
  }
  const newProxy = {
    url,
    isGpt4: false,
    isDisabled: false,
    lastUsed: 0,
    promptCount: 0,
    remaining: 100,
    hash: crypto
      .createHash("sha256")
      .update(url)
      .digest("hex")
      .slice(0, 6),
  };

  await check(newProxy);

  proxyPool.push(newProxy);

  logger.info({ proxy: newProxy.hash }, "Proxy added");

  return newProxy;
}

function disable(proxy: Proxy) {
  const proxyFromPool = proxyPool.find((k) => k.url === proxy.url)!;
  if (proxyFromPool.isDisabled) return;
  proxyFromPool.isDisabled = true;
  logger.warn({ proxy: proxy.hash, url: proxy.url }, "Proxy disabled");
}

function anyAvailable() {
  return proxyPool.some((proxy) => !proxy.isDisabled);
}

function get(model: string) {
  const needsGpt4Proxy = model.startsWith("gpt-4");
  const enabledProxies = proxyPool.filter(
    (proxy) => !proxy.isDisabled
  );

  let availableProxies = enabledProxies.filter(k => k.isGpt4 === needsGpt4Proxy);
  if (availableProxies.length === 0 && !needsGpt4Proxy) {
    availableProxies = enabledProxies;
  }

  if (availableProxies.length === 0) {
    let message = "No proxies available. Please add more proxies.";
    if (needsGpt4Proxy) {
      message =
        "No GPT-4 proxies available. Please add more proxies or use a non-GPT-4 model.";
    }
    logger.error(message);
    return;
  }

  const notLoggingProxies = availableProxies.filter((proxy) => !proxy.isLogging);
  const selectionProxies = notLoggingProxies.length > 0 ? notLoggingProxies : availableProxies;
  const msg = notLoggingProxies.length > 0 ? "Assigning proxy to request." : "Using logging proxy";

  // Return the oldest proxy
  const oldestProxy = selectionProxies.sort((a, b) => a.lastUsed - b.lastUsed)[0];
  logger.info({ proxy: oldestProxy.hash }, msg);
  oldestProxy.lastUsed = Date.now();
  return { ...oldestProxy };
}

function incrementPrompt(proxyHash?: string) {
  if (!proxyHash) return;
  const proxy = proxyPool.find((k) => k.hash === proxyHash)!;
  proxy.promptCount++;
}

function downgradeProxy(proxyHash?: string) {
  if (!proxyHash) return;
  logger.warn({ proxy: proxyHash }, "Downgrading proxy to GPT-3.5.");
  const proxy = proxyPool.find((k) => k.hash === proxyHash)!;
  proxy.isGpt4 = false;
}

function getRemaining(onlyActive = false) {
  let remaining = 0;
  let remainingCount = 0;
  for (const proxy of proxyPool) {
    if (onlyActive && proxy.isDisabled) continue;
    remainingCount++;
    remaining += proxy.remaining;
  }

  if (remainingCount === 0) return '0%';
  return `${Math.round(remaining / remainingCount)}%`;
}

export const proxies = {
  init,
  list,
  get,
  anyAvailable,
  parse,
  add,
  check,
  getRemaining,
  disable,
  incrementPrompt,
  downgradeProxy,
};