#include <filesystem>
#include <iostream>
#include <sstream>
#include <fstream>
#include <random>     // For random number generation
#include <chrono>     // For seeding random number generator
#include <vector> // For getAllKeys
#include <string_view> // For litecask values
#include <stdexcept> // For std::runtime_error
#include <sstream>

#include "server.hpp"
#include "logger.hpp"
#include "hash.hpp"
#include "compress.hpp"
#include "string_utils.hpp" // Include the new utility header
#include "put_handler.hpp"
#include "update_handler.hpp" // Include the new update handler header
#include "utils.hpp"
#include "welcome_page.hpp"
#include "rate_limiter.hpp"
#include "validation.hpp"
#include "HttpController.hpp"
#include "bcrypt.hpp" // For secure token hashing

namespace simple_object_storage {

Server* Server::instance_ = nullptr;
std::mutex Server::instance_mutex_;


bool Server::init_db() {
    try {
        std::filesystem::path db_path = config_.object_store_path / "index.db";
        db_ = std::make_unique<Database>(db_path);
        return true;
    } catch (const std::runtime_error& e) {
        LOG_ERROR("Database initialization error: {}", e.what());
        return false;
    }
}

bool Server::validate_write_request(const drogon::HttpRequestPtr &req, drogon::HttpResponsePtr &res, const std::vector<std::string> &required_params, std::map<std::string, std::string> &params)
{
    std::string client_ip = req->getPeerAddr().toIp();

    // Check if the client is already over the limit (do NOT increment)
    if (auth_rate_limiter_->is_over_limit(client_ip)) {
        res->setStatusCode(drogon::k429TooManyRequests);
        nlohmann::json response = {{"result", "error"}, {"error", "Too many authentication attempts. Please try again later."}};
        res->setBody(response.dump());
        res->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        return false;
    }

    // Get token from Authorization header
    std::string token;
    auto auth_header = req->getHeader("Authorization");
    if (!auth_header.empty() && auth_header.substr(0, 7) == "Bearer ") {
        token = auth_header.substr(7);
    }

    if (token.empty()) {
        // Only count failed attempt (increment the limiter)
        auth_rate_limiter_->is_allowed(client_ip); // This will increment the count
        res->setStatusCode(drogon::k401Unauthorized);
        nlohmann::json response = {{"result", "error"}, {"error", "Missing or invalid Authorization header"}};
        res->setBody(response.dump());
        res->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        return false;
    }

    // Check if token is valid by comparing against stored bcrypt hashes
    bool write_token_valid = false;
    for (const auto& stored_hash : config_.write_tokens) {
        // Verify the token against the stored bcrypt hash
        if (BCrypt::verifyPassword(token, stored_hash)) {
            write_token_valid = true;
            break;
        }
    }
    
    if (!write_token_valid) {
        // Only count failed attempt (increment the limiter)
        auth_rate_limiter_->is_allowed(client_ip); // This will increment the count
        res->setStatusCode(drogon::k403Forbidden);
        nlohmann::json response = {{"result", "error"}, {"error", "Invalid write token"}};
        res->setBody(response.dump());
        res->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        return false;
    }

    // If authentication is successful, do not increment rate limiter

    auto req_params = req->getParameters();
    for (const auto& param : req_params) {
        params[param.first] = param.second;
    }

    for (const auto& param : required_params) {
        if (req->getParameter(param).empty()) {
            res->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Missing required query parameter: " + param}};
            res->setBody(response.dump());
            res->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            return false;
        }
    }

    return true;
}

Server::Server(const ServerConfig& config)
    : config_(config), running_(false) {
    Server::setInstance(this);

    if (!std::filesystem::exists(config_.object_store_path)) {
        LOG_ERROR("Object store directory does not exist: {}", config_.object_store_path.string());
        return;
    }
    
    // Initialize the database
    if (!init_db()) {
        // Error already printed in init_db
        // Consider throwing or setting an error state
    }

    // Initialize the put handler
    put_handler_ = std::make_unique<PutHandler>(*this);

    // Initialize the update handler
    update_handler_ = std::make_unique<UpdateHandler>(*this);

    // Initialize rate limiter
    auth_rate_limiter_ = std::make_unique<RateLimiter>(
        config_.auth_rate_limit,
        std::chrono::seconds(config_.auth_window_seconds)
    );
}

Server::~Server() {
    stop();
    if (Server::instance_ == this) {
        Server::setInstance(nullptr);
    }
}

bool Server::start() {
    if (!db_) { // Check if DB initialization failed
        LOG_ERROR("Database is not initialized. Cannot start server.");
        return false;
    }
    setup_routes();

    running_ = true;
    
    // Configure Drogon
    drogon::app().addListener(config_.host, config_.port);
    drogon::app().setThreadNum(16);
    
    // Set security limits (allowing for large files up to 6GB)
    drogon::app().setClientMaxBodySize(6ULL * 1024 * 1024 * 1024); // 6GB max body size
    drogon::app().setClientMaxMemoryBodySize(100 * 1024 * 1024); // 100MB max memory body (keep this lower)
    drogon::app().enableGzip(true); // Enable compression
    
    // Start the server
    try {
        drogon::app().run();
    } catch (const std::exception& e) {
        running_ = false;
        LOG_ERROR("Failed to start server: {}", e.what());
        return false;
    }
    return true;
}

void Server::stop() {
    if (running_) {
        drogon::app().quit();
        running_ = false;
        LOG_INFO("Server stopped.");
    }
}

void Server::setup_routes() {
    // Configure CORS only for requests with Origin header
    drogon::app().registerPostHandlingAdvice([this](const drogon::HttpRequestPtr &req, const drogon::HttpResponsePtr &resp) {
        if (!req->getHeader("Origin").empty()) {
            add_cors_headers(req, resp);
        }
    });
    
    // Handle OPTIONS requests for CORS preflight
    drogon::app().registerHandlerViaRegex(".*", [this](const drogon::HttpRequestPtr &req,
                                                          std::function<void(const drogon::HttpResponsePtr &)> &&callback) {
        handle_cors_preflight(req, std::move(callback));
    }, {}, "OPTIONS");
    
}

void Server::handle_cors_preflight(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback) {
    auto resp = drogon::HttpResponse::newHttpResponse();
    add_cors_headers(req, resp);
    add_security_headers(resp);
    resp->setStatusCode(drogon::k204NoContent);
    callback(resp);
}

void Server::add_cors_headers(const drogon::HttpRequestPtr& req, const drogon::HttpResponsePtr& res) {
    // Get the origin from the request
    std::string origin = req->getHeader("Origin");
    
    // If no origin header, no CORS headers needed
    if (origin.empty()) {
        return;
    }

    // Check if origin is allowed
    bool origin_allowed = false;
    if (config_.allowed_origins.empty() || 
        std::find(config_.allowed_origins.begin(), config_.allowed_origins.end(), "*") != config_.allowed_origins.end()) {
        origin_allowed = true;
    } else {
        origin_allowed = std::find(config_.allowed_origins.begin(), config_.allowed_origins.end(), origin) != config_.allowed_origins.end();
    }

    if (origin_allowed) {
        res->addHeader("Access-Control-Allow-Origin", origin);
        
        // Add other CORS headers
        std::string methods = join(config_.allowed_methods, ", ");
        res->addHeader("Access-Control-Allow-Methods", methods);
        
        std::string headers = join(config_.allowed_headers, ", ");
        res->addHeader("Access-Control-Allow-Headers", headers);
        
        if (config_.allow_credentials) {
            res->addHeader("Access-Control-Allow-Credentials", "true");
        }
        
        // Add max age for preflight requests
        res->addHeader("Access-Control-Max-Age", "86400");  // 24 hours
    }
}

void Server::add_security_headers(const drogon::HttpResponsePtr& res) {
    // Add security headers to prevent common web vulnerabilities
    
    // Prevent clickjacking attacks
    res->addHeader("X-Frame-Options", "DENY");
    
    // Prevent MIME type sniffing
    res->addHeader("X-Content-Type-Options", "nosniff");
    
    // Enable XSS filter in browsers (legacy, but still useful for older browsers)
    res->addHeader("X-XSS-Protection", "1; mode=block");
    
    // Enforce HTTPS (only add if we're sure the server uses HTTPS)
    // Note: Commented out by default as it requires HTTPS to be configured
    // res->addHeader("Strict-Transport-Security", "max-age=31536000; includeSubDomains");
    
    // Content Security Policy - restrictive by default
    // This prevents loading of external resources and inline scripts
    res->addHeader("Content-Security-Policy", "default-src 'self'; script-src 'none'; style-src 'self' 'unsafe-inline'; img-src 'self' data:; font-src 'self'; connect-src 'self'; frame-ancestors 'none'");
    
    // Referrer Policy - don't leak referrer information
    res->addHeader("Referrer-Policy", "strict-origin-when-cross-origin");
    
    // Permissions Policy (formerly Feature Policy) - disable unnecessary browser features
    res->addHeader("Permissions-Policy", "geolocation=(), microphone=(), camera=(), payment=(), usb=(), magnetometer=(), gyroscope=(), accelerometer=()");
}

std::string Server::join(const std::vector<std::string>& strings, const std::string& delimiter) {
    if (strings.empty()) return "";
    
    std::string result = strings[0];
    for (size_t i = 1; i < strings.size(); ++i) {
        result += delimiter + strings[i];
    }
    return result;
}

void Server::handle_get_object(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback, const std::string& key) {
    std::string hash_str = key;

    // first check if the key matches.
    dbEntry entry;
    if (!db_->get(key, entry)) {
        auto resp = drogon::HttpResponse::newHttpResponse();
        resp->setStatusCode(drogon::k404NotFound);
        nlohmann::json response = {{"result", "error"}, {"error", "Couldn't find: " + key}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        add_security_headers(resp);
        callback(resp);
        return;
    }

    // Construct the file path using the hash string
    std::filesystem::path file_path = config_.object_store_path / entry.hash;
    if (!std::filesystem::exists(file_path) || !std::filesystem::is_regular_file(file_path)) {
        auto resp = drogon::HttpResponse::newHttpResponse();
        resp->setStatusCode(drogon::k404NotFound);
        nlohmann::json response = {{"result", "error"}, {"error", "Hash recognised, but object missing for: " + entry.hash}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        add_security_headers(resp);
        callback(resp);
        return;
    }

    // Send file
    auto resp = drogon::HttpResponse::newFileResponse(file_path.string());
    add_security_headers(resp);
    callback(resp);
}

void Server::handle_get_hash(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback, const std::string& key) {
    const auto& labeltag = key;

    dbEntry entry;
    if (!db_->get(labeltag, entry)) {
        auto resp = drogon::HttpResponse::newHttpResponse();
        resp->setStatusCode(drogon::k404NotFound);
        nlohmann::json response = {{"result", "error"}, {"error", "Label:tag not found"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        add_security_headers(resp);
        callback(resp);
        return;
    }

    auto resp = drogon::HttpResponse::newHttpResponse();
    nlohmann::json response = {{"result", "success"}, {"hash", entry.hash}};
    resp->setBody(response.dump());
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    add_security_headers(resp);
    callback(resp);
}

void Server::handle_get_directory(const drogon::HttpRequestPtr& /*req*/, std::function<void(const drogon::HttpResponsePtr &)>&& callback) {
    std::vector<dbEntry> entries;
    
    if (!db_->list(entries)) {
        auto resp = drogon::HttpResponse::newHttpResponse();
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to retrieve directory listing"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        add_security_headers(resp);
        callback(resp);
        return;
    }
    
    // output the labeltags as an array of label:tag pairs
    nlohmann::json entries_array = nlohmann::json::array();
    for (const auto& entry : entries) {
        entries_array.push_back({{"labeltags", entry.labeltags}, {"hash", entry.hash}});
    }

    auto resp = drogon::HttpResponse::newHttpResponse();
    nlohmann::json response = {{"result", "success"}, {"entries", entries_array}};
    resp->setBody(response.dump());
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    add_security_headers(resp);
    callback(resp);
}

void Server::handle_get_metadata(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback, const std::string& key) {
    std::string hash_str = key;

    // Check if the key is a label:tag format
    dbEntry entry;
    nlohmann::json response;
    auto resp = drogon::HttpResponse::newHttpResponse();

    if (db_->get(key, entry)) {
        // Got it from label:tag, use the hash
        try {
            response = {{"result", "success"}, {"metadata", entry.metadata}};
            resp->setBody(response.dump());
        } catch (const nlohmann::json::exception& e) {
            LOG_ERROR("Error serializing metadata for hash {}: {}", hash_str, e.what());
            resp->setStatusCode(drogon::k500InternalServerError);
            response = {{"result", "error"}, {"error", "Internal server error: Failed to serialize metadata"}};
            resp->setBody(response.dump());
        }
    }
    else
    {
        resp->setStatusCode(drogon::k404NotFound);
        response = {{"result", "error"}, {"error", "Invalid hash: " + hash_str}};
        resp->setBody(response.dump());
    }
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    add_security_headers(resp);
    callback(resp);
}

std::pair<std::string, std::string> Server::parse_labeltag(const std::string& labeltag) const {
    size_t colon_pos = labeltag.find(':');
    if (colon_pos == std::string::npos || colon_pos == 0 || colon_pos == labeltag.length() - 1) {
        return {"", ""};
    }
    return {labeltag.substr(0, colon_pos), labeltag.substr(colon_pos + 1)};
}

void Server::handle_delete_object(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback) {
    dbEntry entry;
    auto resp = drogon::HttpResponse::newHttpResponse();

    {
        std::map<std::string, std::string> params;
        if (!validate_write_request(req, resp, {"hash"}, params)) {
            add_security_headers(resp);
            callback(resp);
            return;
        }

        // Validate hash format
        auto hashValidation = InputValidator::validateHash(params["hash"]);
        if (!hashValidation.valid) {
            resp->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Invalid hash: " + hashValidation.error}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            add_security_headers(resp);
            callback(resp);
            return;
        }

        if (!db_->get(params["hash"], entry)) {
            resp->setStatusCode(drogon::k404NotFound);
            nlohmann::json response = {{"result", "error"}, {"error", "Object not found for: " + params["hash"]}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            callback(resp);
            return;
        }
    } // we only use sanitised data from here on out.

    // Remove all tags that reference this hash
    if (!db_->remove_by_hash(entry.hash)) {
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to remove some or all associated tags"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        add_security_headers(resp);
        callback(resp);
        return;
    }

    // Remove the file
    std::filesystem::path file_path = config_.object_store_path / entry.hash;
    if (std::filesystem::exists(file_path) && std::filesystem::is_regular_file(file_path)) {
        try {
            std::filesystem::remove(file_path);
        } catch (const std::filesystem::filesystem_error& e) {
            LOG_ERROR("Error deleting object file: {}", e.what());
            resp->setStatusCode(drogon::k500InternalServerError);
            nlohmann::json response = {{"result", "error"}, {"error", "Failed to delete object file: " + std::string(e.what())}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            callback(resp);
            return;
        }
    }

    nlohmann::json response = {{"result", "success"}};
    resp->setBody(response.dump());
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    add_security_headers(resp);
    callback(resp);
}

void Server::handle_exists(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback, const std::string& key) {
    nlohmann::json response;
    dbEntry entry;
    if (db_->get(key, entry)) {
        response = {{"result", "success"}, {"exists", true}};
    } else {
        response = {{"result", "success"}, {"exists", false}};
    }
    auto resp = drogon::HttpResponse::newHttpResponse();
    resp->setBody(response.dump());
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    add_security_headers(resp);
    callback(resp);
}

void Server::handle_get_version(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback) {
    // Get the key from the URL path
    auto path = req->getPath();
    auto pos = path.rfind('/');
    std::string key = (pos != std::string::npos) ? path.substr(pos + 1) : "";
    std::string hash_str = key;

    // Check if the key is a label:tag format
    dbEntry entry;
    auto resp = drogon::HttpResponse::newHttpResponse();
    if (!db_->get(key, entry)) {
        nlohmann::json response = {{"result", "failed"}, {"error", "Failed to get version for: " + key}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        add_security_headers(resp);
        callback(resp);
        return;
    }

    nlohmann::json response;
    if (entry.metadata.contains("version")) {
        response = {{"result", "success"}, {"version", entry.metadata["version"]}};
    } else {
        response = {{"result", "failed"}, {"error", "No version found for: " + key}};
    }
    resp->setBody(response.dump());
    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
    add_security_headers(resp);
    callback(resp);
}

} // namespace simple_object_storage