import _ from 'lodash';
import type { Server } from 'http';
import type { Socket } from 'net';
import path from 'path';
import { Request, Router } from 'express';
import { createProxyMiddleware } from 'http-proxy-middleware';
import winston from 'winston';
import logger from '../services/logger';
import consts from '../consts';
import { createProxyModifyResponseMiddleware, isHtml, injectScript } from '@stackbit/dev-common/dist/utils/proxy-utils';
import { renderLoadingPage } from '@stackbit/dev-common/dist/utils/pages';
import socketService from '@stackbit/dev-common/dist/services/socket-service';
import config from '../config';
import Runner from '../runner';

export default function proxyMiddleware(httpServer: Server, router: Router, ssgHost: string, ssgPort: number, runner: Runner) {
    router.use((req, res, next) => {
        _.unset(req, 'headers.x-forwarded-proto');
        const rawHeaders = _.get(req, 'rawHeaders');
        const headerIndex = _.findIndex(rawHeaders, (headerValue) => headerValue.toLowerCase() === 'x-forwarded-proto');
        if (headerIndex > -1) {
            // remove the header key and value
            rawHeaders.splice(headerIndex, 2);
        }

        next();
    });
    router.use((req, res, next) => {
        res.set('Content-Security-Policy', `frame-ancestors *; `);
        next();
    });

    const proxyTargetUrl = ssgHost.toLowerCase().startsWith('http') ? `${ssgHost}:${ssgPort}` : `http://${ssgHost}:${ssgPort}`;

    const handlerProxyError = (err: any, req: any, res: any) => {
        if (err?.code === 'ECONNREFUSED' && res?.status) {
            renderLoadingPage(req, res, {
                assetsDir: path.join(__dirname, '../assets'),
                proxyUrl: proxyTargetUrl,
                proxyPath: req?.path || '',
                htmlMessage: '',
                status: 504,
                statusMessage: 'aa'
            });
        }
    };

    const winstonLogger = winston.createLogger({
        level: config.logLevel === 'debug' ? 'info' : 'silent',
        format: winston.format.combine(winston.format.splat(), winston.format.simple()),
        transports: [new winston.transports.Console()]
    });

    // visit https://github.com/chimurai/http-proxy-middleware#context-matching
    // for docs on how path matching works, it is different from express matching
    const directPaths = _.uniq(_.concat(consts.DEFAULT_DIRECT_PROXY_PATH, runner.getDirectPaths()));
    router.use(
        createProxyMiddleware({
            pathFilter: directPaths,
            target: proxyTargetUrl,
            changeOrigin: runner.getDirectChangeOrigin(),
            ws: true,
            secure: false,
            logger: winstonLogger,
            router: runner.getDirectRoutes() ?? {},
            on: {
                error: (err: any, req: any, res: any) => {
                    logger.debug(`onError in direct proxy, request: ${req.method} ${req.url}, error: ${err.message}`, { error: err });
                    handlerProxyError(err, req, res);
                }
            }
        })
    );

    const proxyWithWs = createProxyMiddleware({
        target: proxyTargetUrl,
        changeOrigin: true,
        ws: false,
        secure: false,
        logger: winstonLogger,
        selfHandleResponse: true,
        on: {
            error: (err: any, req: any, res: any) => {
                logger.debug(`onError in main proxy, request: ${req.method} ${req.url}, error: ${err.message}`, { error: err });
                handlerProxyError(err, req, res);
            },
            proxyRes: createProxyModifyResponseMiddleware({
                logTagName: 'proxy',
                shouldModifyBody: (proxyRes: Request, req: Request) => {
                    return req.method === 'GET' && isHtml(proxyRes);
                },
                modifyBody: (body: string) => {
                    body = injectScript(body, config.snippetUrl);
                    return body;
                }
            })
        }
    });
    router.use(proxyWithWs);

    if (runner.shouldProxyWebsockets()) {
        const allListeners = httpServer.listeners('upgrade');
        httpServer.removeAllListeners('upgrade');
        httpServer.on('upgrade', (req, socket, head) => {
            if (req.url?.startsWith(socketService.socketPath)) {
                (socketService.socket?.engine as any | undefined)?.handleUpgrade(req, socket, head);
            } else {
                allListeners.forEach((listener) => listener(req, socket, head));
                proxyWithWs.upgrade(req, socket as Socket, head);
            }
        });
    }

    return router;
}