// Proxy.cpp
#include "Proxy.h"
#include "HostnameResolver.h"
#include <iostream>
#include <thread>
#include <format>
#include <boost/asio/ssl.hpp>
#include "Utils.h"
#include "Configfile.h"

extern Configfile cfg;

void Proxy::run(boost::asio::ssl::stream<boost::asio::ip::tcp::socket>* outside) {
    boost::asio::ssl::stream<boost::asio::ip::tcp::socket>* inside;
    // Original comments and variables retained
    std::string from = "";
    std::string to = "";
    std::string ehlostr = "";
    std::string resolvedname = "";
    unsigned char last_state = SMTP_STATE_WAIT_FOR_HELO;
    long unimplemented_requests = 0;

    #ifdef HAVE_SPF
    SpfChecker spf_checker;
    #endif

    try {
        // Resolve hostname using the Boost resolver
        try {
            resolvedname = HostnameResolver::resolveHostname(peer_address);
        }
        catch (const std::exception& e) {
            std::cerr << std::format("Hostname resolution error: {}", e.what()) << std::endl;
            resolvedname = "";
        }

        // Configure SSL contexts
        boost::asio::ssl::context ssl_context(boost::asio::ssl::context::tlsv12);
        ssl_context.set_verify_mode(boost::asio::ssl::verify_none);

        bool throttled = cfg.getThrottle(); // Start with a throttled connection
        bool authenticated = false; // Start with a non-authenticated connection
        bool esmtp = false;
        std::string strtemp;
        std::string hermes_status = "unknown";

        // Check whitelist
        if (!cfg.getDnsWhitelistDomains().empty() &&
            Utils::listed_on_dns_lists(cfg.getDnsWhitelistDomains(), cfg.getDnsWhitelistPercentage(), peer_address)) {
            authenticated = true;
            hermes_status = "whitelisted";
            if (cfg.getWhitelistedDisablesEverything()) {
                throttled = false;
            }
        }

        if (cfg.getWhitelistedDisablesEverything() && Utils::whitelisted(cfg.getDatabaseFile(), peer_address)) {
            throttled = false;
            authenticated = true;
        } else {
            if (!cfg.getAllowDataBeforeBanner()) {
                std::this_thread::sleep_for(std::chrono::seconds(cfg.getBannerDelayTime()));

                // Check if data is waiting before server banner
                boost::system::error_code ec;
                size_t available = outside->lowest_layer().available(ec);
                if (ec || available > 0) {
                    std::cout << std::format("421 (data_before_banner) (ip:{})\n", peer_address);
                    std::this_thread::sleep_for(std::chrono::seconds(20));

                    // Write rejection message
                    std::string rejection_msg = "421 Stop sending data before we show you the banner\r\n";
                    boost::asio::write(outside->lowest_layer(), boost::asio::buffer(rejection_msg), ec);
                    return;
                }
            }
        }

        // Connect to the inside server
        boost::asio::ip::tcp::resolver resolver(io_service_);
        boost::asio::ip::tcp::resolver::results_type endpoints = 
            resolver.resolve(cfg.getServerHost(), std::to_string(cfg.getServerPort()));

        boost::asio::connect(inside.lowest_layer(), endpoints);

        // SSL setup
        if (cfg.getOutgoingSsl()) {
            inside.set_verify_mode(boost::asio::ssl::verify_none);
            inside.handshake(boost::asio::ssl::stream_base::client);
        }

        if (cfg.getIncomingSsl()) {
            outside->set_verify_mode(boost::asio::ssl::verify_none);
            outside->handshake(boost::asio::ssl::stream_base::server);
        }

        // Communication buffers
        std::vector<char> read_buffer(4096);
        boost::system::error_code ec;

        // Main loop for communication
        while (!outside->lowest_layer().is_open() || !inside.lowest_layer().is_open()) {
            // Check if the client wants to send something to the server
            size_t client_available = outside->lowest_layer().available(ec);
            if (client_available > 0) {
                size_t bytes_read = outside->read_some(boost::asio::buffer(read_buffer), ec);
                strtemp = std::string(read_buffer.begin(), read_buffer.begin() + bytes_read);

                if (strtemp.length() > 10 && "mail from:" == Utils::strtolower(strtemp.substr(0, 10))) {
                    from = Utils::getmail(strtemp);
                    last_state = SMTP_STATE_WAIT_FOR_RCPTTO;
                }

                if ("ehlo" == Utils::strtolower(strtemp.substr(0, 4))) esmtp = true;

                if (strtemp.length() > 4 && ("ehlo" == Utils::strtolower(strtemp.substr(0, 4)) ||
                        "helo" == Utils::strtolower(strtemp.substr(0, 4)))) {
                    ehlostr = Utils::trim(strtemp.substr(5));
                    last_state = SMTP_STATE_WAIT_FOR_MAILFROM;
                }

                // RCPT TO handling with comprehensive checks
                if (strtemp.length() > 8 && "rcpt to:" == Utils::strtolower(strtemp.substr(0, 8))) {
                    std::string mechanism = "";
                    std::string message = "";
                    to = Utils::getmail(strtemp);

                    // Construct log string using std::format
                    std::string strlog = std::format(
                        "from {} (ip:{}, hostname:{}, {} {}:{}) -> to {}",
                        from, 
                        peer_address, 
                        resolvedname,
                        (esmtp ? "ehlo" : "helo"),
                        ehlostr,
                        to
                    );

                    // Greylisting check
                    std::string code = "250";
                    if (cfg.getGreylist() && !authenticated && 
                        Utils::greylist(cfg.getDatabaseFile(), peer_address, from, to)) {
                        code = "421";
                        mechanism = "greylist";
                        message = std::format("{} Greylisted!! Please try again in a few minutes.", code);
                        std::cout << std::format("checking {}\n", mechanism);
                    }
                    // SPF Check
                    #ifdef HAVE_SPF
                    else if (cfg.getQuerySpf() && !authenticated && 
                             !spf_checker.query(peer_address, ehlostr, from)) {
                        code = cfg.getAddStatusHeader() ? "250" : 
                            (cfg.getReturnTempErrorOnReject() ? "421" : "550");
                        mechanism = "spf";
                        message = std::format(
                            "{} You do not seem to be allowed to send email for that particular domain.", 
                            code
                        );
                        std::cout << std::format("checking {}\n", mechanism);
                    }
                    #endif
                    // Blacklist check
                    else if (!authenticated && 
                             Utils::blacklisted(cfg.getDatabaseFile(), peer_address, to)) {
                        code = cfg.getReturnTempErrorOnReject() ? "421" : "550";
                        mechanism = "allowed-domain-per-ip";
                        message = std::format(
                            "{} You do not seem to be allowed to send email to that particular domain from that address.", 
                            code
                        );
                        std::cout << std::format("checking {}\n", mechanism);
                    }
                    // DNS Blacklist check
                    else if (!cfg.getDnsBlacklistDomains().empty() && !authenticated &&
                             Utils::listed_on_dns_lists(cfg.getDnsBlacklistDomains(), 
                                                        cfg.getDnsBlacklistPercentage(), 
                                                        peer_address)) {
                        code = cfg.getAddStatusHeader() ? "250" : 
                            (cfg.getReturnTempErrorOnReject() ? "421" : "550");
                        mechanism = "dnsbl";
                        message = std::format(
                            "{} You are listed on some DNS blacklists. Get delisted before trying to send us email.", 
                            code
                        );
                        std::cout << std::format("checking {}\n", mechanism);
                    }
                    // Reverse DNS check
                    else if (cfg.getRejectNoReverseResolution() && !authenticated && resolvedname.empty()) {
                        code = cfg.getReturnTempErrorOnReject() ? "421" : "550";
                        mechanism = "no reverse resolution";
                        message = std::format(
                            "{} Your IP address does not resolve to a hostname.", 
                            code
                        );
                        std::cout << std::format("checking {}\n", mechanism);
                    }
                    // HELO/Reverse name check
                    else if (cfg.getCheckHeloAgainstReverse() && !authenticated && ehlostr != resolvedname) {
                        code = cfg.getReturnTempErrorOnReject() ? "421" : "550";
                        mechanism = "helo differs from resolved name";
                        message = std::format(
                            "{} Your IP hostname doesn't match your envelope hostname.", 
                            code
                        );
                        std::cout << std::format("checking {}\n", mechanism);
                    }

                    // Prepare log message
                    if (!mechanism.empty()) {
                        strlog = std::format("({}) {}", mechanism, strlog);
                    }
                    strlog = std::format("{} {}", code, strlog);
                    std::cout << strlog << "\n";

                    // Handle rejection
                    if (code != "250") {
                        // Close inside connection
                        inside.lowest_layer().close();
                        
                        // Delay to annoy spammers
                        std::this_thread::sleep_for(std::chrono::seconds(20));
                        
                        // Send rejection message
                        std::string rejection_msg = message + "\r\n";
                        boost::asio::write(outside->lowest_layer(), boost::asio::buffer(rejection_msg), ec);
                        return;
                    }

                    last_state = SMTP_STATE_WAIT_FOR_DATA;
                }

                // Send to inside server
                boost::asio::write(inside.lowest_layer(), boost::asio::buffer(strtemp), ec);
            }

            // Check if the server wants to send something to the client
            size_t server_available = inside.lowest_layer().available(ec);
            if (server_available > 0) {
                size_t bytes_read = inside.read_some(boost::asio::buffer(read_buffer), ec);
                strtemp = std::string(read_buffer.begin(), read_buffer.begin() + bytes_read);

                // Send to outside socket
                boost::asio::write(outside->lowest_layer(), boost::asio::buffer(strtemp), ec);
            }

            // Throttling
            if (throttled) {
                std::this_thread::sleep_for(std::chrono::seconds(cfg.getThrottlingTime()));
            }
        }
    }
    catch (const boost::system::system_error& e) {
        std::cerr << std::format("Boost.Asio error: {}", e.what()) << std::endl;
    }
    catch (const std::exception& e) {
        std::cerr << std::format("Standard exception: {}", e.what()) << std::endl;
    }
}