#include "NetworkClient.h"
#include "GWWifi.h"             // WiFi management (thread-safe)

extern GwWifi gwWifi;   // Extern declaration of global WiFi instance

extern "C" {
    #include "puff.h"
}

static uint32_t crc32_update(uint32_t crc, const uint8_t* data, size_t len) {
    crc = ~crc;
    for (size_t i = 0; i < len; ++i) {
        crc ^= data[i];
        for (int bit = 0; bit < 8; ++bit) {
            uint32_t mask = -(int32_t)(crc & 1U);
            crc = (crc >> 1) ^ (0xEDB88320U & mask);
        }
    }
    return ~crc;
}

// Constructor
NetworkClient::NetworkClient(size_t reserveSize)
    : _doc(reserveSize),
      _valid(false),
      _jsonRaw(nullptr),
      _jsonRawLen(0),
      _imageWidth(0),
      _imageHeight(0),
      _numberPixels(0),
      _pictureBase64(nullptr),
      _pictureBase64Len(0)
{
}

NetworkClient::~NetworkClient() {
    if (_jsonRaw != nullptr) {
        free(_jsonRaw);
        _jsonRaw = nullptr;
        _jsonRawLen = 0;
    }
}

bool NetworkClient::findJsonIntField(const char* json, size_t len, const char* key, int& outValue) {
    if (json == nullptr || key == nullptr || len == 0) {
        return false;
    }

    char pattern[64];
    int plen = snprintf(pattern, sizeof(pattern), "\"%s\"", key);
    if (plen <= 0 || (size_t)plen >= sizeof(pattern)) {
        return false;
    }

    const char* keyPos = strstr(json, pattern);
    if (keyPos == nullptr) {
        return false;
    }

    const char* end = json + len;
    const char* colon = strchr(keyPos + plen, ':');
    if (colon == nullptr || colon >= end) {
        return false;
    }

    const char* p = colon + 1;
    while (p < end && (*p == ' ' || *p == '\t' || *p == '\r' || *p == '\n')) {
        ++p;
    }
    if (p >= end) {
        return false;
    }

    char* parseEnd = nullptr;
    long value = strtol(p, &parseEnd, 10);
    if (parseEnd == p) {
        return false;
    }
    outValue = (int)value;
    return true;
}

bool NetworkClient::extractJsonStringInPlace(char* json, size_t len, const char* key, char*& outValue, size_t& outLen) {
    outValue = nullptr;
    outLen = 0;

    if (json == nullptr || key == nullptr || len == 0) {
        return false;
    }

    char pattern[64];
    int plen = snprintf(pattern, sizeof(pattern), "\"%s\"", key);
    if (plen <= 0 || (size_t)plen >= sizeof(pattern)) {
        return false;
    }

    char* keyPos = strstr(json, pattern);
    if (keyPos == nullptr) {
        return false;
    }

    char* end = json + len;
    char* colon = strchr(keyPos + plen, ':');
    if (colon == nullptr || colon >= end) {
        return false;
    }

    char* p = colon + 1;
    while (p < end && (*p == ' ' || *p == '\t' || *p == '\r' || *p == '\n')) {
        ++p;
    }
    if (p >= end || *p != '"') {
        return false;
    }

    char* valueStart = p + 1;
    char* cur = valueStart;
    while (cur < end) {
        if (*cur == '\\') {
            ++cur;
            if (cur < end) {
                ++cur;
            }
            continue;
        }
        if (*cur == '"') {
            *cur = '\0';
            outValue = valueStart;
            outLen = (size_t)(cur - valueStart);
            return true;
        }
        ++cur;
    }

    return false;
}

// Skip GZIP Header an goto DEFLATE content
int NetworkClient::skipGzipHeader(const uint8_t* data, size_t len) {
    if (len < 10) return -1;

    if (data[0] != 0x1F || data[1] != 0x8B || data[2] != 8) {
        return -1;
    }

    size_t pos = 10;
    uint8_t flags = data[3];

    if (flags & 4) {
        if (pos + 2 > len) return -1;
        uint16_t xlen = data[pos] | (data[pos+1] << 8);
        pos += 2 + xlen;
    }

    if (flags & 8) {
        while (pos < len && data[pos] != 0) pos++;
        pos++;
    }

    if (flags & 16) {
        while (pos < len && data[pos] != 0) pos++;
        pos++;
    }

    if (flags & 2) pos += 2;

    if (pos >= len) return -1;

    return pos;
}

// HTTP GET + GZIP Decompression (reading in chunks)
bool NetworkClient::httpGetGzip(const String& url, uint8_t*& outData, size_t& outLen) {

    const size_t capacity = READLIMIT;          // Read limit for data (can be adjusted in NetworkClient.h)
    uint8_t* buffer = (uint8_t*)malloc(capacity);

    // If not with WiFi connectetd then return without any activities
    if (!gwWifi.clientConnected()) {
        if (DEBUGING) {Serial.println("No WiFi connection");}
        return false;
    }
    
    // If frame buffer not correct allocated then return without any activities
    if (!buffer) {
        if (DEBUGING) {Serial.println("Malloc failed buffer");}
        return false;
    }

    HTTPClient http;

    // Timeouts to prevent hanging connections
    http.setConnectTimeout(CONNECTIONTIMEOUT);  // Connect timeout in ms (can be adjusted in NetworkClient.h)
    http.setTimeout(TCPREADTIMEOUT);            // Read timeout in ms (can be adjusted in NetworkClient.h)

    http.begin(url);

    // NEW: force server to close the connection after the response (prevents "stuck" keep-alive reads)
    http.addHeader("Connection", "close");

    // NEW: request gzip, but we will only decompress if the server actually answers with gzip
    http.addHeader("Accept-Encoding", "gzip");

    // NEW: register headers BEFORE GET() (more reliable with Arduino HTTPClient)
    if (DEBUGING) {
        // We need follow key words
        const char* keys[] = {
            "Content-Encoding",
            "Transfer-Encoding",
            "Content-Length"
        };
        // Read header
        http.collectHeaders(keys, 3);
    }

    int code = http.GET();
    if (code != HTTP_CODE_OK) {
        Serial.printf("HTTP Client ERROR: %d (%s)\n", code, http.errorToString(code).c_str());

        // Hard reset HTTP + socket
        WiFiClient* tmp = http.getStreamPtr();
        if (tmp) tmp->stop();   // Force close TCP socket
        
        http.end();
        free(buffer);
        return false;
    }
    else{       
        if (DEBUGING) {
            String ce = http.header("Content-Encoding");
            String te = http.header("Transfer-Encoding");
            String cl = http.header("Content-Length");

            // Print header informations
            Serial.printf("Content-Encoding=%s Transfer-Encoding=%s Content-Length=%s\n",
                        ce.c_str(),
                        te.c_str(),
                        cl.c_str());
        }                
    }

    WiFiClient* stream = http.getStreamPtr();

    size_t len = 0;
    uint32_t lastData = millis();
    const uint32_t READ_TIMEOUT = READDATATIMEOUT;   // Timeout for reading data (can be adjusted in NetworkClient.h)

    bool complete = false;
    bool aborting = false;   // NEW: remember if we must force-close socket

    // NEW: detect if server really sent gzip
    String ce = http.header("Content-Encoding");
    bool isGzip = ce.equalsIgnoreCase("gzip");

    // NEW: read expected body size if provided by server (prevents waiting forever for missing bytes)
    int total = http.getSize();  // returns Content-Length, or -1 if unknown/chunked

    // NEW: fail fast if server claims something larger than our buffer
    if (total > 0 && (size_t)total > capacity) {
        Serial.println("Response exceeds READLIMIT.");
        aborting = true;
    }

    // NEW: if not gzip, we will not try to decompress (prevents false "Decompress OK" / random success)
    // You can either handle plain JSON here or just fail-fast.
    if (!isGzip && !aborting) {
        if (DEBUGING) {
            Serial.println("Server response is NOT gzip (Content-Encoding != gzip).");
            Serial.println("Either disable Accept-Encoding: gzip or add plain-body handling here.");
        }

        // --- Plain-body handling (recommended): read full body into outData as-is ---
        // NEW: try to read Content-Length bytes if available (more robust)
        if (total > 0 && (size_t)total > capacity) {
            Serial.println("Plain response exceeds READLIMIT.");
            aborting = true;
        } else {
            // Read until we have all bytes (Content-Length) or until connection closes + buffer drains
            while ((http.connected() || (stream && stream->available())) && !aborting) {
                size_t avail = stream ? stream->available() : 0;
                if (avail == 0) {
                    if (millis() - lastData > READ_TIMEOUT) {
                        Serial.println("TIMEOUT waiting for data (plain)!");
                        aborting = true;
                        break;
                    }
                    delay(1);
                    continue;
                }

                if (len >= capacity) {
                    Serial.println("READLIMIT reached, aborting (plain).");
                    aborting = true;
                    break;
                }

                if (len + avail > capacity)
                    avail = capacity - len;

                int read = stream->readBytes(buffer + len, avail);
                if (read > 0) {
                    len += (size_t)read;
                    lastData = millis();
                }

                // NEW: stop reading as soon as we have the full response
                if (total > 0 && (int)len >= total) {
                    break; // we got full body
                }
            }
        }

        if (aborting) {
            // --- Added: Force-close connection only if aborted to avoid TCP RST storms ---
            if (stream) stream->stop();   // Force close TCP socket
            http.end();
            free(buffer);
            return false;
        }

        if (total > 0 && (int)len != total) {
            Serial.printf("Plain response incomplete: got=%d expected=%d\n", (int)len, total);
            if (stream) stream->stop();
            http.end();
            free(buffer);
            return false;
        }

        // Return plain body to caller
        outData = (uint8_t*)malloc(len + 1);
        if (!outData) {
            Serial.println("Malloc failed outData (plain).");
            // --- Added: Force-close connection only if aborted to avoid TCP RST storms ---
            if (stream) stream->stop();   // Force close TCP socket
            http.end();
            free(buffer);
            return false;
        }
        memcpy(outData, buffer, len);
        outData[len] = 0;
        outLen = len;

        http.end();
        free(buffer);
        return true;
    }

    // --- GZIP path (only if Content-Encoding is gzip) ---
    if (!aborting) {

        // NEW: read exactly Content-Length bytes when available (prevents partial-body timeout loops)
        while ((http.connected() || (stream && stream->available())) && !complete && !aborting) {

            size_t avail = stream ? stream->available() : 0;

            if (avail == 0) {
                // NEW: if Content-Length is known and we already read it all, stop immediately
                if (total > 0 && (int)len >= total) {
                    break;
                }

                if (millis() - lastData > READ_TIMEOUT) {
                    Serial.println("TIMEOUT waiting for data!");
                    aborting = true;   // NEW: mark abnormal exit
                    break;
                }
                delay(1);
                continue;
            }

            // NEW: safety check if buffer limit is reached
            if (len >= capacity) {
                Serial.println("READLIMIT reached, aborting.");
                aborting = true;
                break;
            }

            // NEW: if Content-Length is known, do not read beyond it
            if (total > 0) {
                size_t remaining = (size_t)total - len;
                if (avail > remaining) avail = remaining;
            }

            if (len + avail > capacity)
                avail = capacity - len;

            int read = stream->readBytes(buffer + len, avail);
            if (read <= 0) {
                // NEW: avoid tight loop if read returns zero
                delay(1);
                continue;
            }

            len += (size_t)read;
            lastData = millis();

            if (DEBUGING) {Serial.printf("Read chunk: %d (total: %d)\n", read, (int)len);}

            // NEW: if Content-Length is known and fully received, we can stop reading
            if (total > 0 && (int)len >= total) {
                break;
            }
        }

        // NEW: only attempt gzip parse/decompress after we have a complete body (when Content-Length is known)
        // This avoids wasting heap with repeated malloc/free and reduces fragmentation over long runtimes.
        if (!aborting) {
            if (total > 0 && (int)len != total) {
                Serial.printf("GZIP response incomplete: got=%d expected=%d\n", (int)len, total);
                aborting = true;
            }
        }

        if (!aborting) {
            if (len < 20) {
                aborting = true;
            } else {
                int headerOffset = skipGzipHeader(buffer, len);
                if (headerOffset < 0) {
                    aborting = true;
                } else {
                    size_t deflateLen = len - (size_t)headerOffset;
                    // GZIP trailer (CRC32 + ISIZE) is 8 bytes and not part of deflate stream.
                    if (deflateLen >= 8) {
                        deflateLen -= 8;
                    }

                    unsigned long srcLenForSize = (unsigned long)deflateLen;
                    unsigned long outNeeded = 0;
                    int sizeRes = puff(NIL, &outNeeded, buffer + headerOffset, &srcLenForSize);

                    if (sizeRes != 0) {
                        if (DEBUGING) {
                            Serial.printf("Decompress size probe failed: res=%d src=%lu\n", sizeRes, srcLenForSize);
                        }
                        aborting = true;
                    } else {
                        uint8_t* test = (uint8_t*)malloc((size_t)outNeeded + 1);
                        if (!test) {
                            Serial.println("Malloc failed test buffer, aborting.");
                            aborting = true;
                        } else {
                            unsigned long srcLen = (unsigned long)deflateLen;
                            unsigned long testLen = outNeeded;
                            int res = puff(test, &testLen, buffer + headerOffset, &srcLen);

                            if (res == 0) {
                                uint32_t trailerCrc =
                                    (uint32_t)buffer[len - 8] |
                                    ((uint32_t)buffer[len - 7] << 8) |
                                    ((uint32_t)buffer[len - 6] << 16) |
                                    ((uint32_t)buffer[len - 5] << 24);
                                uint32_t trailerIsize =
                                    (uint32_t)buffer[len - 4] |
                                    ((uint32_t)buffer[len - 3] << 8) |
                                    ((uint32_t)buffer[len - 2] << 16) |
                                    ((uint32_t)buffer[len - 1] << 24);
                                uint32_t calcCrc = crc32_update(0, test, (size_t)testLen);
                                uint32_t calcIsize = (uint32_t)testLen;

                                if (calcCrc != trailerCrc || calcIsize != trailerIsize) {
                                    Serial.printf(
                                        "GZIP CRC/ISIZE mismatch crc=%08lx/%08lx isize=%lu/%lu\n",
                                        (unsigned long)calcCrc,
                                        (unsigned long)trailerCrc,
                                        (unsigned long)calcIsize,
                                        (unsigned long)trailerIsize
                                    );
                                    free(test);
                                    aborting = true;
                                } else {
                                    test[testLen] = 0;
                                    if (DEBUGING) {Serial.printf("Decompress OK! Size: %lu bytes\n", testLen);}
                                    outData = test;
                                    outLen = (size_t)testLen;
                                    complete = true;
                                }
                            } else {
                                if (DEBUGING) {
                                    Serial.printf("Decompress failed: res=%d out=%lu src=%lu\n", res, testLen, srcLen);
                                }
                                free(test);
                                aborting = true;
                            }
                        }
                    }
                }
            }
        }
    }

    // --- Added: Force-close connection only if aborted to avoid TCP RST storms ---
    if (aborting && stream) stream->stop();   // NEW: stop() only on abnormal termination

    http.end();
    free(buffer);

    if (!complete) {
        Serial.println("Failed to complete decompress.");
        return false;
    }

    return true;
}

// Decompress JSON
bool NetworkClient::fetchAndDecompressJson(const String& url) {

    _valid = false;
    _doc.clear();
    _imageWidth = 0;
    _imageHeight = 0;
    _numberPixels = 0;
    _pictureBase64 = nullptr;
    _pictureBase64Len = 0;

    if (_jsonRaw != nullptr) {
        free(_jsonRaw);
        _jsonRaw = nullptr;
        _jsonRawLen = 0;
    }

    uint8_t* raw = nullptr;
    size_t rawLen = 0;

    if (!httpGetGzip(url, raw, rawLen)) {
        Serial.println("GZIP download/decompress failed.");
        return false;
    }

    char* json = reinterpret_cast<char*>(raw);
    bool ok = true;
    ok = findJsonIntField(json, rawLen, "number_pixels", _numberPixels) && ok;
    ok = findJsonIntField(json, rawLen, "width", _imageWidth) && ok;
    ok = findJsonIntField(json, rawLen, "height", _imageHeight) && ok;
    ok = extractJsonStringInPlace(json, rawLen, "picture_base64", _pictureBase64, _pictureBase64Len) && ok;

    if (!ok) {
        Serial.println("JSON field extraction failed.");
        free(raw);
        return false;
    }

    if (_imageWidth <= 0 || _imageHeight <= 0 || _pictureBase64Len == 0) {
        Serial.printf("JSON invalid geometry/data w=%d h=%d b64=%u\n",
                      _imageWidth,
                      _imageHeight,
                      (unsigned int)_pictureBase64Len);
        free(raw);
        return false;
    }

    _jsonRaw = raw;
    _jsonRawLen = rawLen;

    if (DEBUGING) {
        Serial.printf("JSON fields OK: num=%d w=%d h=%d b64=%u\n",
                      _numberPixels,
                      _imageWidth,
                      _imageHeight,
                      (unsigned int)_pictureBase64Len);
    }
    _valid = true;
    return true;
}

JsonDocument& NetworkClient::json() {
    return _doc;
}

int NetworkClient::imageWidth() const {
    return _imageWidth;
}

int NetworkClient::imageHeight() const {
    return _imageHeight;
}

int NetworkClient::numberPixels() const {
    return _numberPixels;
}

const char* NetworkClient::pictureBase64() const {
    return _pictureBase64;
}

size_t NetworkClient::pictureBase64Len() const {
    return _pictureBase64Len;
}

bool NetworkClient::isValid() const {
    return _valid;
}