import { Request, Response, Router } from "express";
import * as http from "http";
import { createProxyMiddleware } from "http-proxy-middleware";
import { logger } from "../logger";
import {
  handleDownstreamErrors,
  handleInternalError,
  incrementKeyUsage,
  copyHttpHeaders,
} from "./common";
import { ipLimiter, keyLimiter } from "./rate-limit";
import { injectMDReq } from "../proxy/middleware/request/md-request";
import {
  addKey,
  languageFilter,
  finalizeBody,
  limitOutputTokens,
} from "./rewriters";

import { proxies } from '../proxies';

const rewriteRequest = (
  proxyReq: http.ClientRequest,
  req: Request,
  res: http.ServerResponse
) => {
  const rewriterPipeline = [
    addKey,
    languageFilter,
    limitOutputTokens,
    injectMDReq,
    finalizeBody,
  ];

  try {
    for (const rewriter of rewriterPipeline) {
      rewriter(proxyReq, req, res, {});
    }
  } catch (error) {
    logger.error(error, "Error while executing proxy rewriter");
    proxyReq.destroy(error as Error);
  }
};

const handleProxiedResponse = async (
  proxyRes: http.IncomingMessage,
  req: Request,
  res: Response
) => {
  try {
    await handleDownstreamErrors(proxyRes, req, res);
  } catch (error) {
    // Handler takes over the response, we're done here.
    return;
  }
  incrementKeyUsage(req);
  copyHttpHeaders(proxyRes, res);
  proxyRes.pipe(res);
};

const proxyRouter = (req: Request): string | undefined => {
  if (!req.body?.stream) {
    const proxy = proxies.get(req.body?.model || "gpt-3.5");
    req.proxy = proxy;
    return proxy?.url;
  }
};

const openaiProxy = createProxyMiddleware({
  target: "https://api.openai.com",
  router: proxyRouter,
  changeOrigin: true,
  on: {
    proxyReq: rewriteRequest,
    proxyRes: handleProxiedResponse,
    error: handleInternalError,
  },
  selfHandleResponse: true,
  logger,
});

const openaiRouter = Router();
openaiRouter.get("/v1/models", openaiProxy);
openaiRouter.post("/v1/chat/completions", ipLimiter, keyLimiter, openaiProxy);
openaiRouter.use((req, res) => {
  logger.warn(`Blocked openai proxy request: ${req.method} ${req.path}`);
  res.status(404).json({ error: "Not found" });
});

export const openai = openaiRouter;