#pragma once

#include <drogon/HttpResponse.h>
#include <drogon/HttpTypes.h>
#include <openssl/evp.h>
#include "pkcs11_common.h"
#include <string>
#include <sstream>
#include <iomanip>
#include <regex>
#include <optional>
#include <map>
#include <vector>
#include <random>
#include <mutex>
#include <chrono>
#include <unordered_map>

namespace provisioner {
    namespace utils {

        /**
         * RAII wrapper for incremental SHA256 hashing using the OpenSSL 3.0 EVP API.
         *
         * Usage:
         *   SHA256Hasher hasher;
         *   hasher.update(data1, len1);
         *   hasher.update(data2, len2);
         *   std::string hex = hasher.finalize();
         */
        class SHA256Hasher {
            EVP_MD_CTX *ctx_;
        public:
            SHA256Hasher() : ctx_(EVP_MD_CTX_new()) {
                EVP_DigestInit_ex(ctx_, EVP_sha256(), nullptr);
            }
            ~SHA256Hasher() { if (ctx_) EVP_MD_CTX_free(ctx_); }

            SHA256Hasher(const SHA256Hasher&) = delete;
            SHA256Hasher& operator=(const SHA256Hasher&) = delete;
            SHA256Hasher(SHA256Hasher&& o) noexcept : ctx_(o.ctx_) { o.ctx_ = nullptr; }
            SHA256Hasher& operator=(SHA256Hasher&& o) noexcept {
                if (ctx_) EVP_MD_CTX_free(ctx_);
                ctx_ = o.ctx_; o.ctx_ = nullptr;
                return *this;
            }

            void update(const void *data, size_t len) {
                EVP_DigestUpdate(ctx_, data, len);
            }

            std::string finalize() {
                unsigned char hash[EVP_MAX_MD_SIZE];
                unsigned int len = 0;
                EVP_DigestFinal_ex(ctx_, hash, &len);
                std::ostringstream oss;
                for (unsigned int i = 0; i < len; ++i)
                    oss << std::hex << std::setfill('0') << std::setw(2) << static_cast<int>(hash[i]);
                return oss.str();
            }
        };
        
        // ===== Configuration Paths =====
        // Package defaults are read first, then user config overrides
        constexpr const char* CONFIG_DEFAULTS_PATH = "/usr/share/rpi-sb-provisioner/defaults/config";
        constexpr const char* CONFIG_USER_PATH = "/etc/rpi-sb-provisioner/config";
        constexpr const char* CUSTOMER_KEY_STORAGE_DIR = "/etc/rpi-sb-provisioner/keys";
        
        // ===== Firmware Information =====
        
        /**
         * Information about a firmware file
         */
        struct FirmwareInfo {
            std::string version;        // e.g., "2024-01-15"
            std::string filename;       // e.g., "pieeprom-2024-01-15.bin"
            std::string filepath;       // Full path to file
            std::string releaseChannel; // e.g., "default", "latest", "beta"
            uintmax_t size;             // File size in bytes
        };
        
        // ===== Key Information =====
        
        /**
         * Information about a cryptographic key
         */
        struct KeyInfo {
            std::string algorithm;      // "RSA", "EC", "DSA", etc.
            int keySize;                // Key size in bits (e.g., 2048, 4096)
            bool isPrivateKey;          // Whether this is a private key
            std::string fingerprint;    // Secure-boot key hash (the value programmed into device OTP); empty for non-RSA-2048 keys
            bool isFitForPurpose;       // Whether key meets Pi secure boot requirements (RSA-2048)
            std::string statusMessage;  // Human-readable status message
            std::string statusLevel;    // "valid", "warning", or "error"
            std::string errorMessage;   // Error details if parsing failed
            bool success;               // Whether key parsing succeeded
            
            KeyInfo() : keySize(0), isPrivateKey(false), isFitForPurpose(false), success(false) {}
        };
        
        /**
         * Canonicalize a config file path, rejecting traversal attempts.
         * Returns the absolute, normalized path when valid.
         */
        std::optional<std::string> canonicalizeConfigPath(const std::string& path);

        /**
         * True when a canonical PEM key path may be validated: either under the
         * provisioner key storage directory or the configured CUSTOMER_KEY_FILE_PEM.
         */
        bool isAllowedCustomerKeyPath(const std::string& canonicalPath);

        /**
         * Parse a PEM key file and extract metadata
         * 
         * @param path Path to the PEM key file
         * @return KeyInfo struct with key metadata
         */
        KeyInfo parseKeyFile(const std::string& path);
        
        /**
         * Parse a PKCS#11 key and extract metadata
         * Requires the HSM to be connected and accessible
         * 
         * @param uri PKCS#11 URI (e.g., pkcs11:object=mykey;type=private)
         * @param pin Optional PIN to use for validation (if empty, uses stored PIN file)
         * @return KeyInfo struct with key metadata
         */
        KeyInfo parsePkcs11Key(const std::string& uri, const std::string& pin = "");

        /**
         * A key object discovered on a PKCS#11 token.
         */
        struct Pkcs11Object {
            std::string uri;    // Full pkcs11: URI for this object
            std::string label;  // Object label (CKA_LABEL) if available
            std::string token;  // Token label the object lives on, if available
            std::string type;   // "private", "public", or "" if unknown
        };

        /**
         * Result of enumerating the PKCS#11 tokens/objects visible to the
         * pkcs11-provider (via p11-kit). Used to help users pick a key without
         * hand-crafting a URI.
         */
        struct Pkcs11Discovery {
            bool providerAvailable;             // pkcs11-provider could be loaded
            std::string errorMessage;           // Populated when discovery failed
            std::vector<Pkcs11Object> objects;  // Discovered key objects

            Pkcs11Discovery() : providerAvailable(false) {}
        };

        /**
         * Check whether the OpenSSL pkcs11-provider is installed and loadable.
         * Does not touch any token, so it never needs a PIN.
         *
         * @return true if the provider module loads
         */
        bool isPkcs11ProviderAvailable();

        /**
         * Enumerate the key objects visible across all PKCS#11 tokens, by
         * listing the bare "pkcs11:" store through the provider. Best-effort:
         * some tokens require a login (PIN) before private objects are listed.
         *
         * @param pin Optional PIN (if empty, uses the stored PIN file)
         * @return discovery result with provider status and any objects found
         */
        Pkcs11Discovery discoverPkcs11(const std::string& pin = "");

        // ===== PKCS#11 PIN Management =====
        
        /**
         * Path to the PKCS#11 PIN file. Single source of truth lives in
         * pkcs11_common.h (shared with the drogon-free rpi-sb-keyhelper).
         */
        constexpr const char* PKCS11_PIN_FILE = pkcs11::PIN_FILE;
        
        /**
         * Check if a PKCS#11 PIN is configured
         * 
         * @return true if PIN file exists and is non-empty
         */
        bool isPkcs11PinConfigured();
        
        /**
         * Save a PKCS#11 PIN securely
         * The PIN is stored in a separate file with restrictive permissions (0400)
         * 
         * @param pin The PIN to save
         * @return true if successful, false otherwise
         */
        bool savePkcs11Pin(const std::string& pin);
        
        /**
         * Remove the PKCS#11 PIN file
         *
         * @return true if successful or file didn't exist, false on error
         */
        bool removePkcs11Pin();

        /**
         * Check whether the file at path is stored device-wrapped (encrypted at
         * rest) rather than as legacy plaintext. Cheap header check.
         *
         * @param path Path to the secret file
         * @return true if the file begins with the device-wrap magic
         */
        bool isFileDeviceWrapped(const std::string& path);

        /**
         * Wrap an existing plaintext secret file in place with the device-bound
         * key, if it is not already wrapped. Used to migrate pre-existing
         * plaintext secrets (PIN, PEM key) ON EXPLICIT USER CONSENT - never
         * call this automatically. The result is written at 0400.
         *
         * @param path Path to the secret file to wrap in place
         * @return true if the file is wrapped on return (including the
         *         already-wrapped no-op case), false on failure
         */
        bool wrapFileInPlace(const std::string& path);

        /**
         * True if any configured secret (the PKCS#11 PIN or the customer PEM
         * key) is currently stored as legacy plaintext rather than wrapped at
         * rest. Evaluated fresh (cheap header reads) so it reflects the live
         * state after a migration. Drives the site-wide migration banner.
         */
        bool hasUnwrappedSecretsAtRest();

        /**
         * Scan the firmware directory for available firmware versions
         * 
         * @param deviceFamily The device family ("4", "5", or "2W")
         * @return A vector of FirmwareInfo sorted by version (newest first)
         */
        std::vector<FirmwareInfo> scanFirmwareDirectory(const std::string& deviceFamily);
        
        /**
         * Get the chip number for a device family
         * 
         * @param deviceFamily The device family ("4", "5", or "2W")
         * @return The chip number ("2711", "2712") or empty string if unknown
         */
        inline std::string getChipNumberForFamily(const std::string& deviceFamily) {
            if (deviceFamily == "4") return "2711";
            if (deviceFamily == "5") return "2712";
            return "";
        }
        
        // ===== CSRF Protection =====
        
        /**
         * CSRF token manager - handles generation and validation of tokens
         * Tokens are time-limited and single-use for maximum security
         */
        class CsrfTokenManager {
        public:
            static CsrfTokenManager& getInstance();
            
            /**
             * Generate a new CSRF token for a session
             * 
             * @param sessionId A unique identifier for the session (can be IP + User-Agent hash)
             * @return The generated token
             */
            std::string generateToken(const std::string& sessionId);
            
            /**
             * Validate a CSRF token
             * 
             * @param sessionId The session identifier
             * @param token The token to validate
             * @return true if valid, false otherwise
             */
            bool validateToken(const std::string& sessionId, const std::string& token);
            
            /**
             * Clean up expired tokens (call periodically)
             */
            void cleanupExpiredTokens();
            
        private:
            CsrfTokenManager() = default;
            
            struct TokenInfo {
                std::string token;
                std::chrono::steady_clock::time_point createdAt;
                bool used = false;
            };
            
            std::mutex mutex_;
            std::unordered_map<std::string, std::vector<TokenInfo>> sessionTokens_;
            
            static constexpr int TOKEN_VALIDITY_SECONDS = 3600;  // 1 hour
            static constexpr int MAX_TOKENS_PER_SESSION = 10;
            static constexpr int TOKEN_LENGTH = 32;
        };
        
        /**
         * Generate a session ID from request (IP + User-Agent hash)
         * 
         * @param req The HTTP request
         * @return A session identifier string
         */
        std::string getSessionIdFromRequest(const drogon::HttpRequestPtr& req);
        
        /**
         * Validate CSRF token from request header or body
         * 
         * @param req The HTTP request
         * @return true if valid, false otherwise
         */
        bool validateCsrfToken(const drogon::HttpRequestPtr& req);
        /**
         * Sanitize path components to prevent directory traversal attacks
         * 
         * @param input The path component string to sanitize
         * @return A sanitized string containing only safe characters
         */
        inline std::string sanitize_path_component(const std::string& input) {
            // Only allow alphanumeric characters and some safe symbols
            // Remove or replace any characters that could be used for path traversal
            std::regex unsafe_chars("[^a-zA-Z0-9_\\-]");
            return std::regex_replace(input, unsafe_chars, "_");
        }

        /**
         * Create a consistent error response that handles both HTML and JSON formats
         * 
         * @param req The original HTTP request
         * @param errorMessage The main error message to display
         * @param statusCode The HTTP status code to return
         * @param errorTitle Optional title for the error page
         * @param errorCode Optional error code identifier
         * @param errorDetails Optional technical details for debugging
         * @return HttpResponsePtr The properly formatted response
         */
        inline drogon::HttpResponsePtr createErrorResponse(
            const drogon::HttpRequestPtr& req,
            const std::string& errorMessage,
            drogon::HttpStatusCode statusCode = drogon::k500InternalServerError,
            const std::string& errorTitle = "",
            const std::string& errorCode = "",
            const std::string& errorDetails = "")
        {
            // Check if client accepts HTML
            auto acceptHeader = req->getHeader("Accept");
            if (!acceptHeader.empty() && acceptHeader.find("text/html") != std::string::npos) {
                // Return HTML error page
                drogon::HttpViewData viewData;
                viewData.insert("currentPage", std::string("")); // No active page in navbar
                viewData.insert("error_message", errorMessage);
                
                if (!errorTitle.empty()) {
                    viewData.insert("error_title", errorTitle);
                }
                
                if (!errorCode.empty()) {
                    viewData.insert("error_code", errorCode);
                }
                
                if (!errorDetails.empty()) {
                    viewData.insert("error_details", errorDetails);
                }
                
                auto resp = drogon::HttpResponse::newHttpViewResponse("error.csp", viewData);
                resp->setStatusCode(statusCode);
                return resp;
            } else {
                // Return JSON error
                Json::Value error;
                error["error"] = errorMessage;
                
                if (!errorCode.empty()) {
                    error["error_code"] = errorCode;
                }
                
                if (!errorDetails.empty()) {
                    error["details"] = errorDetails;
                }
                
                auto resp = drogon::HttpResponse::newHttpJsonResponse(error);
                resp->setStatusCode(statusCode);
                return resp;
            }
        }

        /**
         * Read a single configuration value from the config file
         * 
         * @param key The configuration key to look for
         * @param logAccessToAudit Whether to log file access to the audit log
         * @return An optional string containing the value if found
         */
        std::optional<std::string> getConfigValue(const std::string& key, bool logAccessToAudit = true);
        
        /**
         * Read all configuration values from the config file
         * 
         * @param logAccessToAudit Whether to log file access to the audit log
         * @return A map of all key-value pairs in the config file
         */
        std::map<std::string, std::string> getAllConfigValues(bool logAccessToAudit = true);

        /**
         * Atomically merge values into the user configuration file.
         *
         * Reads the full merged configuration (defaults overridden by the user
         * config), applies the given overrides, then persists the whole set to
         * CONFIG_USER_PATH with each value shell-quoted. Passing an empty value
         * for a key clears it (e.g. dropping a stale GOLD_MASTER_OS_FILE when its
         * image is deleted).
         *
         * Concurrency: the entire read-modify-write is serialised against every
         * other config writer by a process-wide lock, so concurrent writers
         * cannot lose each other's updates. The write itself is atomic (values
         * are written to a temp file which is then renamed over the target), so
         * any reader - in this process or an external one, such as the signing
         * shell scripts - always observes either the complete old file or the
         * complete new one, never a partial write. Readers therefore need no
         * lock. All config writers MUST go through this function to preserve that
         * guarantee.
         *
         * @param updates Key/value overrides to merge into the config
         * @return The full merged config on success (for callers that need the
         *         post-write state), or std::nullopt if the write failed
         */
        std::optional<std::map<std::string, std::string>> setConfigValues(
            const std::map<std::string, std::string>& updates);

        /**
         * Convenience wrapper around setConfigValues for a single key.
         *
         * @param key The configuration key to set
         * @param value The value to write (empty string clears the key)
         * @return true if the config file was written successfully
         */
        bool setConfigValue(const std::string& key, const std::string& value);

        /**
         * Quote a value for safe inclusion in the shell-sourced config file.
         *
         * The config file is written by this service but sourced by the signing
         * shell scripts (e.g. rpi-sb-pkcs11-sign.sh), so any value containing a
         * shell metacharacter (space, ';', '$', ...) must be quoted or the
         * shell mangles it. PKCS#11 URIs always contain ';' (type=private), so
         * an unquoted value is silently truncated at the first ';'. Returns the
         * value unchanged when it needs no quoting, otherwise single-quoted with
         * embedded single quotes escaped as '\'' (the shlex.quote convention).
         *
         * @param value The raw configuration value
         * @return A representation safe to write as the right-hand side of KEY=
         */
        std::string shellQuoteConfigValue(const std::string& value);

        /**
         * Inverse of shellQuoteConfigValue for the read path.
         *
         * Strips a single layer of single or double quotes (unescaping the
         * embedded sequences those quoting styles use) so a round-tripped value
         * reads back identically. Values with no surrounding quotes - the bare
         * lines in the packaged defaults file and legacy user configs - are
         * returned unchanged, keeping the reader backward-compatible.
         *
         * @param value The raw right-hand side of a KEY= line
         * @return The unquoted value
         */
        std::string unquoteConfigValue(const std::string& value);

        /**
         * Create an appropriate error response for configuration file read errors
         * 
         * @param req The original HTTP request
         * @param configKey The configuration key that was being searched for (optional)
         * @return HttpResponsePtr The error response
         */
        drogon::HttpResponsePtr createConfigErrorResponse(
            const drogon::HttpRequestPtr& req,
            const std::string& configKey = "");
    } // namespace utils
} // namespace provisioner 