#pragma once

#include <vector>
#include <cstdint>
#include <string>
#include <memory>
#include <map>
#include <deque>
#include <atomic>
#include <functional>
#include <chrono>

#include <mutex>
#include "config.h"
#include "rsa.h"
#include "aes.h"
#include "x509.h"
#include "sha.h"
#include "pkcs12.h"

namespace netplus {
    class buffer;
    class socket;

    // =====================================================================
    // TLS Session Cache — enables TLS 1.2 abbreviated handshakes.
    // Shared across connections; protected by its own mutex.
    // =====================================================================
    struct TlsSessionCacheEntry {
        std::vector<uint8_t> session_id;         // session ID for resumption
        std::vector<uint8_t> master_secret;      // 48 bytes
        uint16_t             cipher_suite;       // e.g. 0x002F (AES-128-CBC-SHA)
        std::chrono::steady_clock::time_point created;
    };

    class TlsSessionCache {
    public:
        // Look up a session by session_id; returns true if found and not expired.
        bool lookup(const std::vector<uint8_t>& session_id,
                    TlsSessionCacheEntry& out) const {
            std::lock_guard<std::mutex> lk(mtx_);
            auto it = map_.find(session_id);
            if (it == map_.end()) return false;
            auto age = std::chrono::steady_clock::now() - it->second.created;
            if (age > max_age_) { return false; }
            out = it->second;
            return true;
        }

        // Store a session keyed by session_id.
        void store(const std::vector<uint8_t>& session_id,
                   const std::vector<uint8_t>& master_secret,
                   uint16_t cipher_suite) {
            std::lock_guard<std::mutex> lk(mtx_);
            if (map_.size() >= max_entries_ && map_.find(session_id) == map_.end()) {
                evict_oldest();
            }
            map_[session_id] = { session_id, master_secret, cipher_suite,
                                 std::chrono::steady_clock::now() };
        }

        // Look up by host:port key (for client-side reconnection).
        bool lookup_by_host(const std::string& host_key,
                            TlsSessionCacheEntry& out) const {
            std::lock_guard<std::mutex> lk(mtx_);
            auto it = host_map_.find(host_key);
            if (it == host_map_.end()) return false;
            auto age = std::chrono::steady_clock::now() - it->second.created;
            if (age > max_age_) { return false; }
            out = it->second;
            return true;
        }

        // Store by host:port key (client-side).
        void store_by_host(const std::string& host_key,
                           const std::vector<uint8_t>& session_id,
                           const std::vector<uint8_t>& master_secret,
                           uint16_t cipher_suite) {
            std::lock_guard<std::mutex> lk(mtx_);
            host_map_[host_key] = { session_id, master_secret, cipher_suite,
                                    std::chrono::steady_clock::now() };
            // Also store by session_id for server-side lookup
            if (map_.size() >= max_entries_ && map_.find(session_id) == map_.end()) {
                evict_oldest();
            }
            map_[session_id] = host_map_[host_key];
        }

        void setMaxAge(std::chrono::seconds s) { max_age_ = s; }
        void setMaxEntries(size_t n) { max_entries_ = n; }

    private:
        void evict_oldest() {
            auto oldest = map_.begin();
            for (auto it = map_.begin(); it != map_.end(); ++it) {
                if (it->second.created < oldest->second.created)
                    oldest = it;
            }
            if (oldest != map_.end()) map_.erase(oldest);
        }

        mutable std::mutex mtx_;
        std::map<std::vector<uint8_t>, TlsSessionCacheEntry> map_;       // by session_id
        std::map<std::string, TlsSessionCacheEntry>          host_map_;  // by host:port
        size_t max_entries_ = 4096;
        std::chrono::seconds max_age_{7200};  // 2 hours default
    };

    class tls {
    public:
        // --- Cryptographic Helper Methods (Static) ---
        struct CertificateBundle {
            netplus::x509cert    cert;
            std::vector<uint8_t> privateKeyDer;
            netplus::rsa         rsa_key;  // Pre-loaded RSA private key
            uint8_t              ecPrivateKey[32] = {0};  // EC P-256 private scalar
            bool                 has_ec_key = false;
            std::vector<std::vector<uint8_t>> chain;  // CA chain certs in DER (per-SNI)

            // Load cert+key from file(s) or inline PEM data.
            // Auto-detects:
            //   Inline PEM → certPath/keyPath contain "-----BEGIN" PEM data
            //   .p12/.pfx  → PKCS#12 (certPath only, password required)
            //   .pem/.crt/.cer/.der/.key → PEM or DER file (certPath + keyPath)
            bool loadFromFile(const std::string& certPath,
                              const std::string& keyPath = "",
                              const std::string& password = "");
        };

        // Handshake State Machine
        enum class HsState {
            // Server states
            READ_CLIENT_HELLO,
            TLS13_SEND_SERVER_HELLO,
            TLS13_SEND_ENCRYPTED_FLIGHT,
            TLS13_WAIT_CLIENT_FINISHED,
            SEND_SERVER_FLIGHT,
            SEND_ABBREVIATED_FLIGHT,   // TLS 1.2 session resumption (server)
            WAIT_CKE,
            WAIT_CCS,
            WAIT_FIN,
            SEND_CCS_FIN,
            ESTABLISHED,
            // Client states
            CLI_FLUSH_CH,
            CLI_WAIT_SERVER_HELLO,
            CLI_TLS13_WAIT_ENC_EXT,
            CLI_TLS13_WAIT_CERT,
            CLI_TLS13_WAIT_CERT_VERIFY,
            CLI_TLS13_WAIT_FINISHED,
            CLI_TLS13_SEND_FINISHED,
            CLI_TLS13_HANDSHAKE_DONE,
            CLI_WAIT_CERT,
            CLI_TLS12_WAIT_SKE,   // TLS 1.2 ECDHE: read optional ServerKeyExchange
            CLI_WAIT_SHD,
            CLI_SEND_CKE,
            CLI_SEND_CCS_FIN,
            CLI_WAIT_SCCS,
            CLI_WAIT_SFIN,
            DONE,
            FAIL,
            IDLE
        };

        // --- Buffer Management ---
        bool hasBufferedData() const {
            std::lock_guard<std::recursive_mutex> lk(mutex_);
            return !rx_tcp_buf.empty() || recv_off < rx_record_buf.size();
        }
        void pushReceivedData(const uint8_t* data, size_t len) {
            std::lock_guard<std::recursive_mutex> lk(mutex_);
            rx_tcp_buf.insert(rx_tcp_buf.end(), data, data + len);
        }
        bool hasPendingWrite() const {
            std::lock_guard<std::recursive_mutex> lk(mutex_);
            return (!send_queue.empty()) || (!send_record.empty()) || (send_off > 0);
        }

        // Getter methods
        bool getHandshakeDone() const {
            std::lock_guard<std::recursive_mutex> lk(mutex_);
            return handshakeDone;
        }
        const std::string& getRequestedHostname() const { return requested_hostname; }

        // =================================================================
        // PUBLIC INTERFACE - Handshake and Data Transfer
        // =================================================================
        
        // TLS handshake operations
        void handshake_after_accept();
        void handshake_after_connect();

        // TLS record layer - send and receive application data
        size_t recvData(buffer& data, int flags);
        size_t sendData(buffer& data, int flags);

        // --- Initialization and State Management ---
        void resetTLS();
        void flush_out();

    public:

        // =================================================================
        // PROTECTED/PUBLIC METHODS - Organized by category
        // =================================================================

        // Helper methods for queuing raw bytes to send buffer
        void queueRaw(const std::vector<uint8_t>& v);
        void queueRaw(const uint8_t* data, size_t len);
        void queueRaw(std::vector<uint8_t>&& v);     

        // HKDF (HMAC-based Key Derivation Function)
        static std::vector<uint8_t> hkdf_extract(
            const std::vector<uint8_t>& salt,
            const std::vector<uint8_t>& ikm,
            uint16_t cipherSuite
        );

        static std::vector<uint8_t> hkdf_expand_label(
            const std::vector<uint8_t>& secret,
            const std::string& label,
            const std::vector<uint8_t>& context,
            size_t L,
            uint16_t cipherSuite
        );

        // TLS 1.2 PRF (Pseudo-Random Function)
        static std::vector<uint8_t> prf(
            const std::vector<uint8_t>& secret,
            const std::string& label,
            const std::vector<uint8_t>& seed,
            size_t outLen
        );

        // HMAC Helpers
        static std::vector<uint8_t> hmac_sha256(
            const std::vector<uint8_t>& key,
            const std::vector<uint8_t>& data
        );

        static std::vector<uint8_t> hmac_sha384(
            const std::vector<uint8_t>& key,
            const std::vector<uint8_t>& data
        );
        
        static std::vector<uint8_t> hmac_sha1(
            const std::vector<uint8_t>& key,
            const std::vector<uint8_t>& data
        );

        // Signature Helpers
        static std::vector<uint8_t> sign_rsa_sha256_pkcs15(
            const std::vector<uint8_t>& data,
            const netplus::rsa& key
        );

        static std::vector<uint8_t> sign_rsa_pss_sha256(
            const std::vector<uint8_t>& data,
            netplus::rsa& key
        );

        static std::vector<uint8_t> sign_ecdsa_sha256(
            const std::vector<uint8_t>& data,
            const std::vector<uint8_t>& privateKey
        );

        // --- TLS 1.2 Record Layer ---
        std::vector<uint8_t> readTlsRecordAsync();

        std::vector<uint8_t> calculateHMAC(
            const std::vector<uint8_t>& content,
            uint8_t contentType,
            uint64_t seq,
            const std::vector<uint8_t>& macKey
        );

        std::vector<uint8_t> decryptRecordCBC(
            uint8_t recType,
            uint16_t ver,
            const std::vector<uint8_t>& frag
        );

        std::vector<uint8_t> decryptRecordGCM(
            uint8_t recType,
            uint16_t ver,
            const std::vector<uint8_t>& frag
        );

        std::vector<uint8_t> decryptTLS12Record(
            uint8_t recType,
            uint16_t ver,
            const std::vector<uint8_t>& frag
        );

        void sendTLS12Record(
            uint8_t recType,
            const std::vector<uint8_t>& content
        );
        void sendTLS12Record(
            uint8_t recType,
            const uint8_t* data, size_t len
        );

        void sendTLS12RecordGCM(
            uint8_t recType,
            const std::vector<uint8_t>& content
        );
        void sendTLS12RecordGCM(
            uint8_t recType,
            const uint8_t* data, size_t len
        );

        // --- TLS 1.2 Handshake Operations ---
        bool popHandshakeMsg(
            std::vector<uint8_t>& out,
            uint8_t& type
        );

        std::vector<uint8_t> fetchNextHandshakePlain();

        void sendHandshake(
            uint8_t handshakeType,
            const std::vector<uint8_t>& data
        );

        void sendServerHello_TLS12(
            bool secureReneg,
            uint16_t suite
        );

        // --- TLS 1.3 Record Layer ---
        std::vector<uint8_t> tls13_read_record_handshake();

        void tls13_send_record(
            uint8_t inner_type,
            const std::vector<uint8_t>& plaintext,
            bool handshake_keys
        );
        void tls13_send_record(
            uint8_t inner_type,
            const uint8_t* data, size_t len,
            bool handshake_keys
        );

        bool tls13_recv_record(
            uint8_t record_type,
            uint16_t legacy_ver,
            const uint8_t* data, size_t data_len,
            std::vector<uint8_t>& out_plain,
            uint8_t& out_inner_type,
            bool handshake_keys
        );

        void sendTLS13Record(
            uint8_t innerType,
            const std::vector<uint8_t>& plain
        );
        void sendTLS13Record(
            uint8_t innerType,
            const uint8_t* data, size_t len
        );

        // --- TLS 1.3 Handshake Operations ---
        std::vector<uint8_t> fetchNextHandshakeTLS13();

        std::vector<uint8_t> buildServerHello_TLS13(
            uint16_t cipher13,
            const std::vector<uint8_t>& client_session_id,
            uint16_t named_group,
            const std::vector<uint8_t>& pubkey
        );

        std::vector<uint8_t> tls13_build_certificate();
        std::vector<uint8_t> tls13_build_certificate_verify();
        std::vector<uint8_t> tls13_build_server_finished();

        void tls13_send_handshake(
            uint8_t hs_type,
            const std::vector<uint8_t>& body,
            bool handshake_keys
        );

        void queueClientHello();
        void tls13_send_finished(bool handshake_keys);

        // --- TLS 1.3 Transcript and Signing ---
        std::vector<uint8_t> tls13_transcript_hash() {
            // Use SHA-384 for TLS_AES_256_GCM_SHA384 (0x1302), SHA-256 otherwise
            if (chosenSuite == 0x1302 || tls13_cipher == 0x1302) {
                return sha384_hash(handshake_transcript);
            }
            return sha256_hash(handshake_transcript);
        }

        std::vector<uint8_t> ecdsa_sha256_sign(const std::vector<uint8_t>& in);

        bool tls13_pop_hs_from_buf(std::vector<uint8_t>& out_msg);

        void tls13_transcript_add_raw(const std::vector<uint8_t>& hsmsg);

        void tls13_insert_message_hash_if_needed();

        // --- TLS 1.3 Key Derivation ---
        void tls13_derive_handshake_keys(const std::vector<uint8_t>& ecdhe_shared);
        void tls13_derive_application_keys();

        // =================================================================
        // PUBLIC MEMBER VARIABLES - Organized by category
        // =================================================================

        HsState hs_state = HsState::IDLE;

        // --- Handshake State Flags ---
        bool is_client = false;
        bool handshakeStarted = false;
        bool handshakeDone = false;
        bool secure_reneg = true;
        bool extended_master_secret = false;
        bool ccs_received = false;
        bool cli_cke_sent = false;
        bool cli_ccs_fin_sent = false;
        bool client_offered_tls13 = false;
        bool encflight_queued = false;
        bool compat_ccs_sent = false;
        bool tls13_server_hello_queued = false;
        bool tls13_encflight_queued = false;
        bool serverFinishedQueued = false;
        bool serverFlightQueued = false;

        // --- TLS 1.2 Session Resumption ---
        // Shared session cache (set by server or client code)
        TlsSessionCache* session_cache = nullptr;

        // Server-side: true when abbreviated handshake is in progress
        bool is_resumption = false;

        // Client-side: set these before connect to attempt resumption
        bool attempt_resumption = false;
        std::vector<uint8_t> resumption_session_id;
        std::vector<uint8_t> resumption_master_secret;
        uint16_t resumption_cipher = 0;

        // --- TLS Version and Cipher Suite ---
        bool is_tls13 = false;
        uint16_t neg_version = 0x0303;        // 0x0303 = TLS 1.2, 0x0304 = TLS 1.3
        uint16_t tls13_cipher = 0x1301;       // TLS_AES_128_GCM_SHA256 (default)
        uint16_t chosenSuite = 0x002F;        // TLS 1.2 cipher suite
        uint16_t selected_group = 0;          // Named group for key share

        // --- SNI (Server Name Indication) Support ---
        std::shared_ptr<const std::map<std::string, CertificateBundle>> cert_map;
        std::string requested_hostname;
        CertificateBundle* selected_cert_bundle = nullptr;

        // --- ALPN (Application-Layer Protocol Negotiation) Support ---
        std::string client_alpn_protocols;
        std::string selected_alpn;

        // Callback for ALPN protocol selection.
        // Receives the list of client-offered protocols, returns the selected one.
        // If not set, the first client-offered protocol is accepted.
        std::function<std::string(const std::vector<std::string>&)> alpn_select_cb;

        // --- Transcript and Handshake Data ---
        std::vector<uint8_t> handshake_transcript;
        std::vector<uint8_t> clientHelloRawBytes;
        std::vector<uint8_t> tls12_transcript_before_client_finished;
        std::vector<uint8_t> tls13_client_keyshare;
        std::vector<uint8_t> client_session_id;

        // --- TLS 1.2 Keys and Secrets ---
        std::vector<uint8_t> masterSecret;
        std::vector<uint8_t> clientRandom;
        std::vector<uint8_t> serverRandom;
        std::vector<uint8_t> mac_key;
        std::vector<uint8_t> client_mac_key;

        // --- TLS 1.2 GCM implicit IVs (4 bytes each) ---
        uint8_t tls12_client_write_iv[4] = {0};
        uint8_t tls12_server_write_iv[4] = {0};
        bool tls12_is_gcm = false;

        // --- TLS 1.2 ECDHE state ---
        uint8_t tls12_ecdhe_priv[32] = {0};  // Server ephemeral P-256 private key

        // --- TLS 1.2 Record Layer Encryption/Decryption ---
        std::unique_ptr<netplus::aes> aes = nullptr;
        std::unique_ptr<netplus::aes> aes_recv = nullptr;
        uint64_t send_seq = 0;
        uint64_t recv_seq = 0;

        // --- Send Buffer Management ---
        std::deque<std::vector<uint8_t>> send_queue;
        std::vector<uint8_t> send_record;
        size_t send_off = 0;
        std::atomic<bool> pendingIocpWrite{false};

        // --- Receive Buffer Management ---
        std::vector<uint8_t> rx_record_buf;
        std::vector<uint8_t> rx_handshake_buf;
        std::vector<uint8_t> rx_tcp_buf;
        size_t rx_handshake_off = 0;
        size_t recv_off = 0;

        // --- X25519 Key Share (for TLS 1.3) ---
        std::vector<uint8_t> client_keyshare_x25519;
        std::vector<uint8_t> server_keyshare_x25519;
        std::vector<uint8_t> server_priv_x25519;
        std::vector<uint8_t> server_pub_x25519;
        std::vector<uint8_t> client_priv_x25519;
        std::vector<uint8_t> x25519_shared;

        // --- ECDHE Key Share (P-256, for TLS 1.3) ---
        std::vector<uint8_t> client_keyshare_ecdhe;
        std::vector<uint8_t> server_keyshare_ecdhe;
        uint8_t client_priv_ecdhe[32] = {0};
        uint8_t server_priv_ecdhe[32] = {0};
        std::vector<uint8_t> ecdhe_shared;

        // --- ECC P-256 Private Key ---
        uint8_t ec_priv[32] = {0};
        bool has_ec_key = false;
        bool tls13_got_key_share = false;

        // --- TLS 1.3 Handshake Traffic Keys ---
        uint8_t tls13_hs_key_c2s[32] = {0};
        uint8_t tls13_hs_iv_c2s[12] = {0};
        uint8_t tls13_hs_key_s2c[32] = {0};
        uint8_t tls13_hs_iv_s2c[12] = {0};
        std::unique_ptr<netplus::aes> aes13_hs_send = nullptr;
        std::unique_ptr<netplus::aes> aes13_hs_recv = nullptr;
        uint64_t tls13_seq_hs_send = 0;
        uint64_t tls13_seq_hs_recv = 0;
        uint64_t tls13_hs_send_seq = 0;
        uint64_t tls13_hs_recv_seq = 0;

        // --- TLS 1.3 Application Traffic Keys ---
        uint8_t tls13_app_key_c2s[32] = {0};
        uint8_t tls13_app_iv_c2s[12] = {0};
        uint8_t tls13_app_key_s2c[32] = {0};
        uint8_t tls13_app_iv_s2c[12] = {0};
        std::unique_ptr<netplus::aes> aes13_app_send = nullptr;
        std::unique_ptr<netplus::aes> aes13_app_recv = nullptr;
        uint64_t tls13_seq_app_send = 0;
        uint64_t tls13_seq_app_recv = 0;
        uint64_t tls13_app_send_seq = 0;
        uint64_t tls13_app_recv_seq = 0;

        // --- TLS 1.3 Traffic Secrets ---
        std::vector<uint8_t> tls13_early_secret;
        std::vector<uint8_t> tls13_hs_secret;
        std::vector<uint8_t> tls13_master_secret;
        std::vector<uint8_t> tls13_c_hs_secret;
        std::vector<uint8_t> tls13_s_hs_secret;
        std::vector<uint8_t> tls13_c_ap_secret;
        std::vector<uint8_t> tls13_s_ap_secret;
        std::vector<uint8_t> tls13_s_app_secret;
        std::vector<uint8_t> tls13_c_app_secret;

        // --- TLS 1.3 IV and IV Material ---
        std::vector<uint8_t> cli_read_iv;
        std::vector<uint8_t> cli_write_iv;

        // --- TLS 1.3 Extension Negotiation ---
        bool hasTLS13SV = false;
        bool hasTLS13KS = false;

        // --- Peer RSA key (client stores server's public key here) ---
        netplus::rsa peer_rsa;

        // --- TLS compatibility flags ---
        // Set to true to send a TLS 1.2-only ClientHello without TLS 1.3 extensions.
        // Use for legacy embedded devices that choke on key_share / supported_versions.
        bool client_tls12_only = false;

        // --- Client TLS 1.2 ECDHE state ---
        bool server_cert_is_ecdsa = false;
        std::vector<uint8_t> tls12_srv_ecdhe_pub;  // server ECDHE public key from ServerKeyExchange

        // --- Hostname (deprecated/legacy) ---
        mutable std::recursive_mutex mutex_;
        std::string hostname;
        
        // --- Socket reference for sending data ---
        socket* _socket = nullptr;
        void setSocket(socket* s) { _socket = s; }
    };
}
