#include "put_handler.hpp"
#include "logger.hpp"
#include "hash.hpp"
#include "compress.hpp"
#include "string_utils.hpp"
#include "utils.hpp"
#include "validation.hpp"
#include <drogon/MultiPart.h>

#include <random>
#include <chrono>
#include <fstream>
#include <sstream>

/*

IMPORTANT: handling of tags.

When an object is uploaded:
- the object is stored in the object store
- the database is primarily 



*/

namespace simple_object_storage {

PutHandler::PutHandler(Server& server) : server_(server) {}

void PutHandler::handle_upload_object(const drogon::HttpRequestPtr& req, std::function<void(const drogon::HttpResponsePtr &)>&& callback) {
    // Check all request parameters first before processing any data
    auto resp = drogon::HttpResponse::newHttpResponse();
    
    std::map<std::string, std::string> params;
    if (!server_.validate_write_request(req, resp, {}, params)) {  // No required params now since token is in header
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Check content length first
    auto contentLengthHeader = req->getHeader("content-length");
    if (!contentLengthHeader.empty()) {
        try {
            size_t contentLength = std::stoull(contentLengthHeader);
            const size_t MAX_UPLOAD_SIZE = 6ULL * 1024 * 1024 * 1024; // 6GB
            if (contentLength > MAX_UPLOAD_SIZE) {
                resp->setStatusCode(drogon::k413RequestEntityTooLarge);
                nlohmann::json response = {{"result", "error"}, {"error", "File too large"}};
                resp->setBody(response.dump());
                resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
                server_.add_security_headers(resp);
                callback(resp);
                return;
            }
        } catch (const std::exception&) {
            // Invalid content-length header, continue with parsing
        }
    }
    
    // Parse the multipart form data
    drogon::MultiPartParser fileParser;
    if (fileParser.parse(req) != 0) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to parse multipart data"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }
    
    auto &files = fileParser.getFiles();
    const drogon::HttpFile* fileData = nullptr;
    for (auto& file : files) {
        if (file.getItemName() == "file") {
            fileData = &file;
            break;
        }
    }
    
    if (!fileData) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "No file provided in upload"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Parse metadata if provided
    nlohmann::json metadata;
    
    // First check if metadata was sent as a form field
    auto &parameters = fileParser.getParameters();
    auto metadataParam = parameters.find("metadata");
    if (metadataParam != parameters.end()) {
        try {
            metadata = nlohmann::json::parse(metadataParam->second);
        } catch (const nlohmann::json::parse_error& e) {
            resp->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Invalid JSON metadata: " + std::string(e.what())}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            callback(resp);
            return;
        }
    } else {
        // If not found as parameter, check if it was sent as a file
        for (auto& file : files) {
            if (file.getItemName() == "metadata") {
                try {
                    auto fileContent = file.fileContent();
                    std::string content(fileContent.data(), fileContent.size());
                    metadata = nlohmann::json::parse(content);
                } catch (const nlohmann::json::parse_error& e) {
                    resp->setStatusCode(drogon::k400BadRequest);
                    nlohmann::json response = {{"result", "error"}, {"error", "Invalid JSON metadata: " + std::string(e.what())}};
                    resp->setBody(response.dump());
                    resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
                    callback(resp);
                    return;
                }
                break;
            }
        }
    }

    // Validate metadata using comprehensive validation
    auto metadataValidation = InputValidator::validateMetadata(metadata);
    if (!metadataValidation.valid) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Invalid metadata: " + metadataValidation.error}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Validate required metadata fields
    if (!metadata.contains("labeltags") || !metadata["labeltags"].is_array() || metadata["labeltags"].empty()) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", "Missing or invalid required metadata field: labeltags (must be non-empty array of label:tag pairs)"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Extract and validate labeltags
    std::vector<std::string> labeltags;
    for (const auto& lt : metadata["labeltags"]) {
        if (!lt.is_string()) {
            resp->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Invalid label:tag pair format - must be a string"}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            server_.add_security_headers(resp);
            callback(resp);
            return;
        }
        labeltags.push_back(lt.get<std::string>());
    }
    
    // Validate all labeltags using input validator
    auto labeltagsValidation = InputValidator::validateLabelTags(labeltags);
    if (!labeltagsValidation.valid) {
        resp->setStatusCode(drogon::k400BadRequest);
        nlohmann::json response = {{"result", "error"}, {"error", labeltagsValidation.error}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Validate and add filename to metadata if not provided
    std::string filename = fileData->getFileName();
    if (!filename.empty()) {
        auto filenameValidation = InputValidator::validateFilename(filename);
        if (!filenameValidation.valid) {
            resp->setStatusCode(drogon::k400BadRequest);
            nlohmann::json response = {{"result", "error"}, {"error", "Invalid filename: " + filenameValidation.error}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            server_.add_security_headers(resp);
            callback(resp);
            return;
        }
    }
    
    // Validate filename in metadata if present (updated)
    if (metadata.contains("filename")) {
        if (metadata["filename"].is_string()) {
            auto filenameInMetadata = metadata["filename"].get<std::string>();
            auto filenameValidation = InputValidator::validateFilename(filenameInMetadata);
            if (!filenameValidation.valid) {
                resp->setStatusCode(drogon::k400BadRequest);
                nlohmann::json response = {{"result", "error"}, {"error", "Invalid filename in metadata: " + filenameValidation.error}};
                resp->setBody(response.dump());
                resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
                server_.add_security_headers(resp);
                callback(resp);
                return;
            }
        }
    } else if (!filename.empty()) {
        metadata["filename"] = filename;
    }

    // Now that all parameters are validated, process the upload
    
    // Generate a random number for the temporary filename
    std::mt19937_64 rng(std::chrono::high_resolution_clock::now().time_since_epoch().count());
    std::uniform_int_distribution<uint64_t> dist;
    uint64_t random_num = dist(rng);
    std::string temp_filename = "temp_" + std::to_string(random_num);

    // Create temporary file
    std::filesystem::path temp_path = server_.config_.object_store_path / temp_filename;
    std::ofstream temp_file(temp_path, std::ios::binary);
    if (!temp_file.is_open()) {
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to create temporary file"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    temp_file.close();
    
    // Use Drogon's built-in file saving for large files
    try {
        fileData->saveAs(temp_path.string());
    } catch (const std::exception& e) {
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to write to temporary file: " + std::string(e.what())}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        std::filesystem::remove(temp_path);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Ensure the temporary file is removed even if errors occur
    ScopeFileDeleter temp_file_deleter(temp_path);

    // Calculate SHA-256 hash
    std::string hash = hash_file(temp_path.string());
    if (hash.empty()) {
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to calculate hash"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

    // Add file metadata
    add_file_metadata(temp_path.string(), metadata);

    // Check if filename ends with ".tgz" using the utility function
    if (utils::ends_with(metadata["filename"], ".tgz")) {
        metadata["tgz_content_hash"] = get_hash_from_tgz(temp_path.string());
    }

    // Move file to final location (using SHA-256 hash as filename)
    std::filesystem::path final_path = server_.config_.object_store_path / hash;

    if (!std::filesystem::exists(final_path)) {
        try {
            std::filesystem::rename(temp_path, final_path);
            temp_file_deleter.release();
        } catch (const std::filesystem::filesystem_error& e) {
            LOG_ERROR("Error renaming temp file: {}", e.what());
            resp->setStatusCode(drogon::k500InternalServerError);
            nlohmann::json response = {{"result", "error"}, {"error", "Failed to store object file"}};
            resp->setBody(response.dump());
            resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
            callback(resp);
            return;
        }
    }

    // Update database index
    dbEntry entry;
    entry.hash = hash;
    entry.labeltags = metadata["labeltags"].get<std::vector<std::string>>();
    entry.metadata = metadata;

    if (!server_.db_->update_or_insert(entry)) {
        resp->setStatusCode(drogon::k500InternalServerError);
        nlohmann::json response = {{"result", "error"}, {"error", "Failed to update database index"}};
        resp->setBody(response.dump());
        resp->setContentTypeCode(drogon::CT_APPLICATION_JSON);
        // Attempt to clean up the moved file if index fails
        try { if (std::filesystem::exists(final_path)) std::filesystem::remove(final_path); } catch(...) {};
        server_.add_security_headers(resp);
        callback(resp);
        return;
    }

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

void PutHandler::add_file_metadata(const std::string& file_path, nlohmann::json& metadata) const {
    // get the file size
    metadata["file_size"] = std::filesystem::file_size(file_path);

    // get the file modification time
    auto ftime = std::filesystem::last_write_time(file_path);
    auto sctp = std::chrono::time_point_cast<std::chrono::system_clock::duration>(
        ftime - std::filesystem::file_time_type::clock::now()
        + std::chrono::system_clock::now()
    );
    metadata["file_modification_time"] = std::chrono::system_clock::to_time_t(sctp);
}


} // namespace simple_object_storage 