#include "AuthApiClient.h"
#include <ArduinoJson.h>
#include "mbedtls/base64.h"
#include "SpiRamAllocator.cpp"

AuthApiClient::AuthApiClient(const String& apiBase, int requestId) : _apiBase(apiBase), _requestId(requestId) 
{

}

SpiRamAllocator alloc;

void AuthApiClient::setInsecure(bool enable) 
{ 
  _useInsecure = enable; 
}

void AuthApiClient::setRootCA(const char* pemRootCA) 
{
  _rootCA = pemRootCA; _useInsecure = false; _useBundle = false;
}
void AuthApiClient::useCertBundle(bool enable) 
{ 
  _useBundle = enable; 
  
  if (enable) 
  { _useInsecure = false; 
    _rootCA = nullptr; 
  } 
}

bool AuthApiClient::beginHttp(HTTPClient& http, WiFiClientSecure& client, const String& fullUrl) 
{
  client.setTimeout(15000);

#if defined(ARDUINO_ARCH_ESP32)
  if (_useBundle) 
  { 
    Serial.printf("Using certificate bundle\n");
    //client.setCACertBundle(nullptr); 
    return false;
  }
  else if (_rootCA) 
  { 
    Serial.printf("Using root certificate authority\n");
    client.setCACert(_rootCA); 
    }
  else if (_useInsecure) 
  { 
    Serial.printf("Using HTTP Insecure\n");
    client.setInsecure(); 
  }
  else 
  { 
    return false; 
  }
#else
  if (_rootCA == nullptr && !_useInsecure) 
  { 
    return false; 
  }
#endif
  return http.begin(client, fullUrl);
}

bool AuthApiClient::httpsPostJsonFullUrl(const String& fullUrl, const String& jsonBody, String& respBody) 
{
  WiFiClientSecure client; 
  HTTPClient http;

  Serial.printf("CREATE HTTP: %s\n", fullUrl.c_str());

  if (!beginHttp(http, client, fullUrl)) 
  {
    return false;
  } 

  Serial.printf("Adding Headers\n");

  http.addHeader("Content-Type", "application/json");

   const String& tok = accessToken();

  // Build Authorization header safely (avoid large temporaries)
  if (!tok.isEmpty()) 
  {
    Serial.printf("Adding Token\n");

    String auth;
    auth.reserve(8 + tok.length());  // "Bearer " + token
    auth = "Bearer ";
    auth += tok;

    // sanitize in case token includes stray CR/LF
    auth.replace("\r", "");
    auth.replace("\n", "");

    http.addHeader("Authorization", auth);
   
    //Serial.printf("TOKEN: %s\n", auth.c_str());
  }

  Serial.printf("JSON: %u %s \n", (unsigned)jsonBody.length(), jsonBody.c_str());

  Serial.printf("POST: Start\n");
  int code = http.POST(jsonBody);
  Serial.printf("POST: End\n");

  if (code <= 0) 
  { 
    Serial.printf("POST failed: %s (%d)\n", http.errorToString(code).c_str(), code);

    http.end(); 
    return false; 
  }

  respBody = http.getString(); 
  http.end();

  //Serial.printf("RESPONSE BODY: %s\n", respBody.c_str());
  return (code >= 200 && code < 300);
}


bool AuthApiClient::httpsPostJsonReturnBytesFullUrl(const String& fullUrl, const String& jsonBody, uint8_t*& outBufPS, size_t& outLen) 
{
  WiFiClientSecure client; 
  HTTPClient http;

  Serial.printf("CREATE HTTP: %s\n", fullUrl.c_str());

  if (!beginHttp(http, client, fullUrl)) 
  {
    return false;
  } 

Serial.printf("Free heap: %u\n", (unsigned)ESP.getFreeHeap());
  Serial.printf("Adding Headers");

  http.addHeader("Content-Type", "application/json");

  // Build Authorization header safely (avoid large temporaries)
  if (!_accessToken.isEmpty()) 
  {
    static String auth;                 // reuse to reduce heap churn
    auth.remove(0);

    Serial.printf("Adding Token\n");

    auth.reserve(8 + _accessToken.length());  // "Bearer " + token
    auth = "Bearer ";
    auth += _accessToken;

    // sanitize in case token includes stray CR/LF
    auth.replace("\r", "");
    auth.replace("\n", "");

    http.addHeader("Authorization", auth);
   
    //Serial.printf("TOKEN: %s\n", auth.c_str());
  }

Serial.printf("Free heap: %u\n", (unsigned)ESP.getFreeHeap());
  //Serial.printf("JSON: %u %s \n", (unsigned)jsonBody.length(), jsonBody.c_str());

  Serial.printf("POST: Start\n");
  int code = http.POST(jsonBody);
  Serial.printf("Free heap: %u\n", (unsigned)ESP.getFreeHeap());
  Serial.printf("POST: End\n");

  // Content-Length if provided, else -1 for chunked
  int contentLen = http.getSize();
  NetworkClient* s = http.getStreamPtr();

  // Allocate initial PSRAM buffer
  size_t cap = 0;
  if (contentLen > 0) 
  {
    cap = (size_t)contentLen;
  } else 
  {
    // Unknown length: start with 64 KB and grow as needed
    cap = 64 * 1024;
  }

  uint8_t* buf = (uint8_t*) ps_malloc(cap);

  if (!buf) 
  {
    Serial.println("PSRAM alloc failed");
    http.end();
    return false;
  }

  size_t total = 0;
  const uint32_t deadline = millis() + 15000;

  while (http.connected() && (contentLen != 0 || contentLen == -1)) 
  {
    if (millis() > deadline) 
    {
      Serial.println("Receive timeout");
      free(buf);
      http.end();
      return false;
    }

    size_t avail = s->available();
    if (!avail) { delay(1); continue; }

    // Determine how many bytes we can/should read this iteration
    size_t want = (contentLen > 0) ? (size_t)min((int)avail, contentLen) : avail;

    // Ensure capacity (grow if needed for chunked or short Content-Length)
    if (total + want > cap) {
      size_t newCap = max(cap * 2, total + want);   // grow exponentially
      uint8_t* grown = (uint8_t*) ps_realloc(buf, newCap);

      if (!grown) 
      {
        Serial.println("PSRAM realloc failed");
        free(buf);
        http.end();
        return false;
      }

      buf = grown;
      cap = newCap;
    }

    int r = s->read(buf + total, want);

    if (r > 0) 
    {
      total += (size_t)r;

      if (contentLen > 0) 
      {
        contentLen -= r;
      }
    }
  }

  http.end();

  outBufPS = buf;
  outLen   = total;

  //Serial.printf("RESPONSE BODY: %s\n", respBody.c_str());
  return (code >= 200 && code < 300);
}

bool AuthApiClient::httpsGetFullUrl(const String& fullUrl, String& respBody) {
  WiFiClientSecure client; HTTPClient http;
  if (!beginHttp(http, client, fullUrl)) return false;
  if (_accessToken.length() > 0) http.addHeader("Authorization", "Bearer " + _accessToken);
  int code = http.GET();
  if (code <= 0) { http.end(); return false; }
  respBody = http.getString(); 
  http.end();
  return (code >= 200 && code < 300);
}

bool AuthApiClient::base64UrlDecode(const String& in, String& out) 
{
  String s = in; s.replace('-', '+'); s.replace('_', '/'); while (s.length() % 4 != 0) s += '=';
  auto idx=[](char c)->int{if(c>='A'&&c<='Z')return c-'A';if(c>='a'&&c<='z')return c-'a'+26;if(c>='0'&&c<='9')return c-'0'+52;if(c=='+')return 62;if(c=='/')return 63;return-1;};
  out="";int val=0,valb=-8;for(size_t i=0;i<s.length();i++){char c=s[i];if(c=='=')break;int d=idx(c);if(d<0)continue;val=(val<<6)+d;valb+=6;if(valb>=0){out+=char((val>>valb)&0xFF);valb-=8;}}return true;
}

time_t AuthApiClient::jwtExp(const String& jwt) {
  int dot1=jwt.indexOf('.'); int dot2=jwt.indexOf('.',dot1+1); if(dot1<0||dot2<0)return 0; String payloadB64=jwt.substring(dot1+1,dot2);
  String payloadJson; if(!base64UrlDecode(payloadB64,payloadJson)) return 0;
  StaticJsonDocument<1024> doc; if(deserializeJson(doc,payloadJson)) return 0;
  if(!doc.containsKey("exp")) return 0; return (time_t)doc["exp"].as<long>();
}

bool AuthApiClient::login(const String& userName, const String& password) 
{
  String url=_apiBase+"/logon"; 
  //SpiRamAllocator alloc;
  JsonDocument req(&alloc); 
  
  req["id"]=_requestId; 
  JsonObject p=req.createNestedObject("payload"); 
  p["userName"]=userName; 
  p["password"]=password;

  String body,resp; 
  
  serializeJson(req,body); 

  Serial.printf("JSON: %s\n", body);
  
  if(!httpsPostJsonFullUrl(url,body,resp)) 
  {
    return false;
  }

  JsonDocument doc(&alloc); 
  
  if(deserializeJson(doc,resp)) 
  {
    return false; 
  }
  
  if(!(doc["success"]|false)) 
  {
    return false;
  }


   setAccessToken(doc["payload"]["token"].as<String>()); 
  _refreshToken=doc["payload"]["refreshToken"].as<String>(); 
  _tokenExpEpoch=jwtExp(_accessToken);
  _authed=(_accessToken.length()>0&&_tokenExpEpoch>0); 
  
  return _authed;
}

bool AuthApiClient::loadDisk(const String& unitId, const String& diskId) 
{
  String url=_apiBase+"/operations/loadimage"; 
  //SpiRamAllocator alloc;
  JsonDocument req(&alloc); 
  
  req["id"]=_requestId; 
  JsonObject p=req.createNestedObject("payload"); 
  p["unitId"]=unitId; 
  p["diskId"]=diskId;

  String body,resp; 
  
  serializeJson(req,body); 

  req.clear();
  
  if(!httpsPostJsonFullUrl(url,body,resp)) 
  {
    return false;
  }

  JsonDocument doc;
  
  if(deserializeJson(doc,resp)) 
  {
    return false; 
  }
  
  if(!(doc["success"]|false)) 
  {
    return false;
  }
 
  return _authed;
}

bool AuthApiClient::loadDiskSector(const String& unitId, const String& diskId, uint32_t lba, uint32_t offset, void* buffer, size_t bufsize) 
{
  if (!buffer || bufsize == 0) 
  {
    Serial.println("API: invalid buffer");
    return false;
  }

  String url=_apiBase+"/operations/loadsector"; 
  //SpiRamAllocator alloc;
  JsonDocument req(&alloc); 
  
  req["id"]=_requestId; 
  JsonObject p=req.createNestedObject("payload"); 
  p["unitId"]=unitId; 
  p["diskId"]=diskId;
  //p["lba"]=lba;
  //p["offset"]=offset;
  //p["length"]=bufsize;

  String body,resp; 
  
  Serial.printf("API: Serialise\n");
  serializeJson(req,body); 

  Serial.printf("API: Start\n");

  Serial.printf("Free heap: %u\n", (unsigned)ESP.getFreeHeap());

uint8_t* psBuf = nullptr;
size_t psLen = 0;

  if(!httpsPostJsonReturnBytesFullUrl(url,body, psBuf, psLen)) 
  {
    Serial.printf("API: Failed\n");
    return false;
  }

  memcpy(buffer, psBuf, min(psLen, bufsize));

  free(psBuf);
  psBuf = nullptr;

  Serial.printf("API: End\n");

/*
  JsonDocument doc; 
  
  if(deserializeJson(doc,resp)) 
  {
    return false; 
  }
  
  if(!(doc["success"]|false)) 
  {
    return false;
  }
 
  // --- Get payload.data ---
  const char* encoded = doc["payload"]["data"];

  if (!encoded) 
  {
    return false;
  }

  size_t out_len = 0;
  int res = mbedtls_base64_decode(
      (unsigned char*)buffer, bufsize, &out_len,
      (const unsigned char*)encoded, strlen(encoded)
  );

  if (res != 0) 
  {
    // decoding failed
    return false;
  }*/

  return true;
}

bool AuthApiClient::refresh() 
{
  if(_refreshToken.length()==0)return false; String url=_apiBase+"/logon/refresh";
  StaticJsonDocument<320> req; req["id"]=_requestId; JsonObject p=req.createNestedObject("payload"); p["refreshToken"]=_refreshToken;
  String body,resp; serializeJson(req,body); if(!httpsPostJsonFullUrl(url,body,resp)) return false;
  StaticJsonDocument<2048> doc; if(deserializeJson(doc,resp)) return false; if(!(doc["success"]|false)) return false;
  _accessToken=doc["payload"]["token"].as<String>(); _refreshToken=doc["payload"]["refreshToken"].as<String>(); _tokenExpEpoch=jwtExp(_accessToken);
  _authed=(_accessToken.length()>0&&_tokenExpEpoch>0); return _authed;
}

bool AuthApiClient::ensureTokenFresh(long refreshEarlySeconds) 
{
  if(!_authed)
  {
    Serial.printf("Token Refresh: Not Authorised");
    return false;
  } 
   
  time_t nowT=time(nullptr); 
  if(nowT==0)
  {
    Serial.printf("Token Refresh: Refresh Overdue");
    return refresh(); 
  }
   
  if(nowT>=(_tokenExpEpoch-refreshEarlySeconds))
  {
    Serial.printf("Token Refresh: Not Early");
    return refresh();     
  }
  
  return true;
}

bool AuthApiClient::postJson(const String& path, const String& jsonBody, String& respBody, bool autoRefresh) 
{
  if(autoRefresh&&_authed)
  {
    ensureTokenFresh(); 
  }
  
  String url=_apiBase+path; 
  return httpsPostJsonFullUrl(url,jsonBody,respBody);
}

void AuthApiClient::setAccessToken(const String& token) 
{
  _accessToken.reserve(token.length());
  _accessToken = token;
  _accessToken.replace("\r", "");
  _accessToken.replace("\n", "");
}

bool AuthApiClient::get(const String& path, String& respBody, bool autoRefresh) 
{
  if(autoRefresh&&_authed) 
  {
    ensureTokenFresh(); 
  }

  String url=_apiBase+path; 
  return httpsGetFullUrl(url,respBody);
}