#include "seobeo/seobeo.h"
#include <ctype.h>
#include <time.h>

static int64_t Seobeo_Client_Monotonic_Milliseconds(void)
{
  struct timespec now;
  clock_gettime(CLOCK_MONOTONIC, &now);
  return (int64_t)now.tv_sec * 1000 + now.tv_nsec / 1000000;
}

static boolean Seobeo_Client_Read_Timed_Out(int32 timeout_ms, int64_t last_progress_ms)
{
  return timeout_ms > 0 &&
         Seobeo_Client_Monotonic_Milliseconds() - last_progress_ms >= timeout_ms;
}

static size_t Seobeo_Client_Find_Header_Length(const uint8 *buffer, size_t length)
{
  if (!buffer || length < 4)
    return 0;
  for (size_t i = 0; i + 3 < length; i++)
  {
    if (buffer[i] == '\r' && buffer[i + 1] == '\n' &&
        buffer[i + 2] == '\r' && buffer[i + 3] == '\n')
      return i + 4;
  }
  return 0;
}

static const char *Seobeo_Client_Header_Value(
    Seobeo_Request_Entry *headers,
    const char *name)
{
  if (!headers || !name)
    return NULL;
  for (size_t i = 0; i < Dowa_Array_Length(headers); i++)
  {
    if (headers[i].key && strcasecmp(headers[i].key, name) == 0)
      return headers[i].value;
  }
  return NULL;
}

static boolean Seobeo_Client_Header_Has_Token(
    const char *value, const char *token)
{
  if (!value || !token)
    return FALSE;
  size_t token_length = strlen(token);
  const char *cursor = value;
  while (*cursor)
  {
    while (*cursor == ' ' || *cursor == '\t' || *cursor == ',')
      cursor++;
    const char *end = cursor;
    while (*end && *end != ',')
      end++;
    const char *trimmed_end = end;
    while (trimmed_end > cursor &&
           (trimmed_end[-1] == ' ' || trimmed_end[-1] == '\t'))
      trimmed_end--;
    if ((size_t)(trimmed_end - cursor) == token_length &&
        strncasecmp(cursor, token, token_length) == 0)
      return TRUE;
    cursor = end;
  }
  return FALSE;
}

static boolean Seobeo_Client_Decode_Chunked(
    const char *encoded, size_t encoded_length,
    char *decoded, size_t decoded_capacity, size_t *decoded_length)
{
  if (!encoded || !decoded || !decoded_length)
    return FALSE;
  size_t input_offset = 0;
  size_t output_offset = 0;
  while (input_offset < encoded_length)
  {
    size_t line_end = input_offset;
    while (line_end + 1 < encoded_length &&
           !(encoded[line_end] == '\r' && encoded[line_end + 1] == '\n'))
      line_end++;
    if (line_end + 1 >= encoded_length)
      return FALSE;

    size_t chunk_size = 0;
    boolean have_digit = FALSE;
    for (size_t i = input_offset; i < line_end && encoded[i] != ';'; i++)
    {
      uint8 c = (uint8)encoded[i];
      int32 digit;
      if (c >= '0' && c <= '9')
        digit = c - '0';
      else if (c >= 'a' && c <= 'f')
        digit = c - 'a' + 10;
      else if (c >= 'A' && c <= 'F')
        digit = c - 'A' + 10;
      else
        return FALSE;
      if (chunk_size > (SIZE_MAX - (size_t)digit) / 16)
        return FALSE;
      chunk_size = chunk_size * 16 + (size_t)digit;
      have_digit = TRUE;
    }
    if (!have_digit)
      return FALSE;
    input_offset = line_end + 2;
    if (chunk_size == 0)
    {
      *decoded_length = output_offset;
      return TRUE;
    }
    if (chunk_size > encoded_length - input_offset ||
        chunk_size > decoded_capacity - output_offset)
      return FALSE;
    memcpy(decoded + output_offset, encoded + input_offset, chunk_size);
    output_offset += chunk_size;
    input_offset += chunk_size;
    if (input_offset + 1 >= encoded_length ||
        encoded[input_offset] != '\r' || encoded[input_offset + 1] != '\n')
      return FALSE;
    input_offset += 2;
  }
  return FALSE;
}

static void Seobeo_Client_Parse_Url(const char *url, char **p_host, 
    char **p_port, char **p_path, boolean *p_use_tls, Dowa_Arena *p_arena)
{
  if (!url)
    return;

  const char *start = url;
  *p_use_tls = FALSE;

  if (strncmp(url, "https://", 8) == 0)
  {
    *p_use_tls = TRUE;
    start = url + 8;
  }
  else if (strncmp(url, "http://", 7) == 0)
  {
    *p_use_tls = FALSE;
    start = url + 7;
  }

  const char *slash = strchr(start, '/');
  const char *colon = strchr(start, ':');

  if (colon && (!slash || colon < slash))
  {
    size_t host_len = colon - start;
    *p_host = Dowa_Arena_Allocate(p_arena, host_len + 1);
    memcpy(*p_host, start, host_len);
    (*p_host)[host_len] = '\0';

    const char *port_start = colon + 1;
    size_t port_len = slash ? (slash - port_start) : strlen(port_start);
    *p_port = Dowa_Arena_Allocate(p_arena, port_len + 1);
    memcpy(*p_port, port_start, port_len);
    (*p_port)[port_len] = '\0';
  }
  else
  {
    size_t host_len = slash ? (slash - start) : strlen(start);
    *p_host = Dowa_Arena_Allocate(p_arena, host_len + 1);
    memcpy(*p_host, start, host_len);
    (*p_host)[host_len] = '\0';

    *p_port = Dowa_Arena_Allocate(p_arena, 8);
    strcpy(*p_port, *p_use_tls ? "443" : "80");
  }

  if (slash)
  {
    size_t path_len = strlen(slash);
    *p_path = Dowa_Arena_Allocate(p_arena, path_len + 1);
    strcpy(*p_path, slash);
  }
  else
  {
    *p_path = Dowa_Arena_Allocate(p_arena, 2);
    strcpy(*p_path, "/");
  }
}

Seobeo_Client_Request *Seobeo_Client_Request_Create(const char *url)
{
  Seobeo_Client_Request *p_req = malloc(sizeof(Seobeo_Client_Request));
  if (!p_req)
    return NULL;

  memset(p_req, 0, sizeof(Seobeo_Client_Request));

  p_req->p_arena = Dowa_Arena_Create(1024 * 1024 * 5);
  if (!p_req->p_arena)
  {
    free(p_req);
    return NULL;
  }

  size_t url_len = strlen(url);
  p_req->url = Dowa_Arena_Allocate(p_req->p_arena, url_len + 1);
  strcpy(p_req->url, url);

  Seobeo_Client_Parse_Url(url, &p_req->host, &p_req->port, &p_req->path, &p_req->use_tls, p_req->p_arena);

  p_req->method = Dowa_Arena_Allocate(p_req->p_arena, 4);
  strcpy(p_req->method, "GET");

  p_req->follow_redirects = FALSE;
  p_req->max_redirects = 10;
  p_req->timeout_ms = 0;

  return p_req;
}

void Seobeo_Client_Request_Set_Method(Seobeo_Client_Request *p_req, const char *method)
{
  if (!p_req || !method)
    return;

  size_t len = strlen(method);
  p_req->method = Dowa_Arena_Allocate(p_req->p_arena, len + 1);
  strcpy(p_req->method, method);
}

void Seobeo_Client_Request_Add_Header_Map(Seobeo_Client_Request *p_req, const char *key, const char *value)
{
  if (!p_req || !key || !value)
    return;

  char *key_copy = Dowa_Arena_Allocate(p_req->p_arena, strlen(key) + 1);
  strcpy(key_copy, key);

  char *value_copy = Dowa_Arena_Allocate(p_req->p_arena, strlen(value) + 1);
  strcpy(value_copy, value);

  Dowa_HashMap_Push_Arena(p_req->headers_map, key_copy, value_copy, p_req->p_arena);
}

void Seobeo_Client_Request_Add_Header_Array(Seobeo_Client_Request *p_req, const char *header)
{
  if (!p_req || !header)
    return;

  char *header_copy = Dowa_Arena_Allocate(p_req->p_arena, strlen(header) + 1);
  strcpy(header_copy, header);

  Dowa_Array_Push_Arena(p_req->headers_array, header_copy, p_req->p_arena);
}

void Seobeo_Client_Request_Set_Body(Seobeo_Client_Request *p_req, const char *body, size_t length)
{
  if (!p_req || !body)
    return;

  if (length == 0)
    length = strlen(body);

  p_req->body = Dowa_Arena_Allocate(p_req->p_arena, length + 1);
  memcpy(p_req->body, body, length);
  p_req->body[length] = '\0';
  p_req->body_length = length;
}

void Seobeo_Client_Request_Set_Follow_Redirects(Seobeo_Client_Request *p_req, boolean follow, int32 max_redirects)
{
  if (!p_req)
    return;

  p_req->follow_redirects = follow;
  p_req->max_redirects = max_redirects > 0 ? max_redirects : 10;
}

void Seobeo_Client_Request_Set_Timeout_Milliseconds(Seobeo_Client_Request *p_req, int32 timeout_ms)
{
  if (!p_req)
    return;
  p_req->timeout_ms = timeout_ms > 0 ? timeout_ms : 0;
}

void Seobeo_Client_Request_Set_Download_Path(Seobeo_Client_Request *p_req, const char *path)
{
  if (!p_req || !path)
    return;

  size_t len = strlen(path);
  p_req->download_path = Dowa_Arena_Allocate(p_req->p_arena, len + 1);
  strcpy(p_req->download_path, path);
}

static int Seobeo_Client_Build_Request_Header(Seobeo_Client_Request *p_req, char *buffer, size_t buffer_size)
{
  int offset = 0;

  offset += snprintf(buffer + offset, buffer_size - offset,
                     "%s %s HTTP/1.1\r\n"
                     "Host: %s\r\n",
                     p_req->method, p_req->path, p_req->host);

  boolean has_content_length = FALSE;
  boolean has_connection = FALSE;

  if (p_req->headers_map)
  {
    size_t count = Dowa_Array_Length(p_req->headers_map);
    for (size_t i = 0; i < count; i++)
    {
      const char *key = p_req->headers_map[i].key;
      const char *value = p_req->headers_map[i].value;

      if (strcasecmp(key, "content-length") == 0)
        has_content_length = TRUE;
      if (strcasecmp(key, "connection") == 0)
        has_connection = TRUE;

      offset += snprintf(buffer + offset, buffer_size - offset,
                         "%s: %s\r\n", key, value);
    }
  }

  if (p_req->headers_array)
  {
    size_t count = Dowa_Array_Length(p_req->headers_array);
    for (size_t i = 0; i < count; i++)
    {
      const char *header = p_req->headers_array[i];

      if (strncasecmp(header, "content-length:", 15) == 0)
        has_content_length = TRUE;
      if (strncasecmp(header, "connection:", 11) == 0)
        has_connection = TRUE;

      offset += snprintf(buffer + offset, buffer_size - offset,
                         "%s\r\n", header);
    }
  }

  if (p_req->body && !has_content_length)
  {
    offset += snprintf(buffer + offset, buffer_size - offset,
                       "Content-Length: %zu\r\n", p_req->body_length);
  }

  if (!has_connection)
  {
    offset += snprintf(buffer + offset, buffer_size - offset,
                       "Connection: close\r\n");
  }

  offset += snprintf(buffer + offset, buffer_size - offset, "\r\n");

  return offset;
}

static Seobeo_Client_Response *Seobeo_Client_Parse_Response(
    Seobeo_Handle *p_handle,
    const char *download_path,
    int32 timeout_ms)
{
  Seobeo_Client_Response *p_resp = malloc(sizeof(Seobeo_Client_Response));
  if (!p_resp)
    return NULL;

  memset(p_resp, 0, sizeof(Seobeo_Client_Response));

  p_resp->p_arena = Dowa_Arena_Create(1024 * 1024 * 10); // 10 MB 
  if (!p_resp->p_arena)
  {
    free(p_resp);
    return NULL;
  }

  int64_t last_progress_ms = Seobeo_Client_Monotonic_Milliseconds();
  while (TRUE)
  {
    int r = Seobeo_Handle_Read(p_handle);
    if (r == -2)
      break;
    if (r < 0)
    {
      Seobeo_Client_Response_Destroy(p_resp);
      return NULL;
    }
    if (r > 0)
      last_progress_ms = Seobeo_Client_Monotonic_Milliseconds();

    if (Seobeo_Client_Find_Header_Length(
            p_handle->read_buffer, p_handle->read_buffer_len) > 0)
      break;

    if (r == 0)
    {
      if (Seobeo_Client_Read_Timed_Out(timeout_ms, last_progress_ms))
      {
        Seobeo_Log(SEOBEO_ERROR, "HTTP response header timed out\n");
        Seobeo_Client_Response_Destroy(p_resp);
        return NULL;
      }
      usleep(1000);
      continue;
    }
  }

  size_t hdr_len = Seobeo_Client_Find_Header_Length(
      p_handle->read_buffer, p_handle->read_buffer_len);
  if (hdr_len == 0)
  {
    Seobeo_Client_Response_Destroy(p_resp);
    return NULL;
  }

  char *buf = Dowa_Arena_Allocate(p_resp->p_arena, hdr_len + 1);
  memcpy(buf, p_handle->read_buffer, hdr_len);
  buf[hdr_len] = '\0';
  char *hdr_end = buf + hdr_len - 4;

  char version[16];
  int status_code;
  char status_text[256];
  int scan_result = sscanf(buf, "%15s %d %255[^\r\n]", version, &status_code, status_text);

  if (scan_result >= 2)
  {
    p_resp->status_code = status_code;
    if (scan_result >= 3)
    {
      size_t len = strlen(status_text);
      p_resp->status_text = Dowa_Arena_Allocate(p_resp->p_arena, len + 1);
      strcpy(p_resp->status_text, status_text);
    }
  }

  char *line = buf + strlen(version) + 1;
  while (*line && !isdigit(*line))
    line++;
  while (*line && isdigit(*line))
    line++;
  while (*line && *line == ' ')
    line++;
  while (*line && *line != '\r')
    line++;
  line += 2;

  while (line < hdr_end)
  {
    char *next = strstr(line, "\r\n");
    if (!next)
      break;

    char *colon = memchr(line, ':', next - line);
    if (colon)
    {
      size_t key_len = colon - line;
      size_t value_len = next - colon - 1;

      char *val_start = colon + 1;
      if (*val_start == ' ')
      {
        val_start++;
        value_len--;
      }

      char *key = Dowa_Arena_Allocate(p_resp->p_arena, key_len + 1);
      memcpy(key, line, key_len);
      key[key_len] = '\0';

      char *val = Dowa_Arena_Allocate(p_resp->p_arena, value_len + 1);
      memcpy(val, val_start, value_len);
      val[value_len] = '\0';

      Dowa_HashMap_Push_Arena(p_resp->headers, key, val, p_resp->p_arena);

      if (strcasecmp(key, "Location") == 0)
      {
        p_resp->redirect_url = Dowa_Arena_Allocate(p_resp->p_arena, value_len + 1);
        strcpy(p_resp->redirect_url, val);
      }
    }

    line = next + 2;
  }

  Seobeo_Handle_Consume(p_handle, (uint32)hdr_len);

  size_t body_len = 0;
  const char *content_length = Seobeo_Client_Header_Value(
      p_resp->headers, "Content-Length");
  const char *transfer_encoding = Seobeo_Client_Header_Value(
      p_resp->headers, "Transfer-Encoding");
  boolean is_chunked = Seobeo_Client_Header_Has_Token(
      transfer_encoding, "chunked");
  if (content_length && !is_chunked)
  {
    body_len = (size_t)strtoull(content_length, NULL, 10);
  }

  FILE *p_file = NULL;
  if (download_path)
  {
    p_file = fopen(download_path, "wb");
    if (!p_file)
    {
      Seobeo_Log(SEOBEO_ERROR, "Failed to open file for writing: %s\n", download_path);
      return p_resp;
    }
  }

  if (body_len > 0)
  {
    char *body = download_path ? NULL : Dowa_Arena_Allocate(p_resp->p_arena, body_len + 1);
    size_t total_read = 0;

    while (total_read < body_len)
    {
      size_t available = p_handle->read_buffer_len;
      size_t to_copy = (body_len - total_read) < available ? (body_len - total_read) : available;

      if (to_copy > 0)
      {
        if (download_path)
          fwrite(p_handle->read_buffer, 1, to_copy, p_file);
        else
          memcpy(body + total_read, p_handle->read_buffer, to_copy);

        total_read += to_copy;
        Seobeo_Handle_Consume(p_handle, (uint32)to_copy);
        last_progress_ms = Seobeo_Client_Monotonic_Milliseconds();
      }

      if (total_read < body_len)
      {
        int r = Seobeo_Handle_Read(p_handle);
        if (r == -2 || r < 0)
        {
          if (p_file) fclose(p_file);
          Seobeo_Client_Response_Destroy(p_resp);
          return NULL;
        }
        if (r > 0)
          last_progress_ms = Seobeo_Client_Monotonic_Milliseconds();
        if (r == 0)
        {
          if (Seobeo_Client_Read_Timed_Out(timeout_ms, last_progress_ms))
          {
            Seobeo_Log(SEOBEO_ERROR, "HTTP response body timed out\n");
            if (p_file) fclose(p_file);
            Seobeo_Client_Response_Destroy(p_resp);
            return NULL;
          }
          usleep(1000);
          continue;
        }
      }
    }

    if (!download_path)
    {
      body[body_len] = '\0';
      p_resp->body = body;
      p_resp->body_length = body_len;
    }
    else
    {
      p_resp->body_length = total_read;
    }
  }
  else
  {
    size_t cap = 1024 * 1024 * 5;
    size_t used = 0;
    char *body = (!download_path || is_chunked)
        ? Dowa_Arena_Allocate(p_resp->p_arena, cap) : NULL;

    while (1)
    {
      int n = p_handle->read_buffer_len > 0
          ? (int)p_handle->read_buffer_len
          : Seobeo_Handle_Read(p_handle);
      if (n > 0)
      {
        last_progress_ms = Seobeo_Client_Monotonic_Milliseconds();
        if (download_path && !is_chunked)
        {
          fwrite(p_handle->read_buffer, 1, p_handle->read_buffer_len, p_file);
          used += p_handle->read_buffer_len;
        }
        else
        {
          if (used + p_handle->read_buffer_len >= cap)
          {
            Seobeo_Log(SEOBEO_WARNING, "Response body too large, truncating...\n");
            break;
          }
          memcpy(body + used, p_handle->read_buffer, p_handle->read_buffer_len);
          used += p_handle->read_buffer_len;
        }
        Seobeo_Handle_Consume(p_handle, (uint32)p_handle->read_buffer_len);
      }
      else if (n == -2)
        break;
      else if (n == 0)
      {
        if (Seobeo_Client_Read_Timed_Out(timeout_ms, last_progress_ms))
        {
          Seobeo_Log(SEOBEO_ERROR, "HTTP response body timed out\n");
          if (p_file) fclose(p_file);
          Seobeo_Client_Response_Destroy(p_resp);
          return NULL;
        }
        usleep(1000);
        continue;
      }
      else
      {
        if (p_file) fclose(p_file);
        Seobeo_Client_Response_Destroy(p_resp);
        return NULL;
      }
    }

    if (is_chunked)
    {
      char *decoded = Dowa_Arena_Allocate(p_resp->p_arena, used + 1);
      size_t decoded_length = 0;
      if (!decoded ||
          !Seobeo_Client_Decode_Chunked(
              body, used, decoded, used, &decoded_length))
      {
        if (p_file) fclose(p_file);
        Seobeo_Client_Response_Destroy(p_resp);
        return NULL;
      }
      decoded[decoded_length] = '\0';
      if (download_path)
      {
        fwrite(decoded, 1, decoded_length, p_file);
        p_resp->body_length = decoded_length;
      }
      else
      {
        p_resp->body = decoded;
        p_resp->body_length = decoded_length;
      }
    }
    else if (!download_path)
    {
      p_resp->body = body;
      p_resp->body_length = used;
      if (used < cap)
        body[used] = '\0';
    }
    else
    {
      p_resp->body_length = used;
    }
  }

  if (p_file)
    fclose(p_file);

  return p_resp;
}

static Seobeo_Client_Response *Seobeo_Client_Execute_Single(Seobeo_Client_Request *p_req)
{
  Seobeo_Handle *p_handle = Seobeo_Stream_Handle_Client_Create(p_req->host, p_req->port, p_req->use_tls);
  if (!p_handle || p_handle->socket < 0)
  {
    if (p_handle)
      Seobeo_Handle_Destroy(p_handle);
    return NULL;
  }

  char request_buffer[8192];
  int request_len = Seobeo_Client_Build_Request_Header(p_req, request_buffer, sizeof(request_buffer));

  Seobeo_Handle_Queue(p_handle, (uint8*)request_buffer, (uint32)request_len);

  if (p_req->body && p_req->body_length > 0)
    Seobeo_Handle_Queue(p_handle, (uint8*)p_req->body, (uint32)p_req->body_length);

  if (Seobeo_Handle_Flush(p_handle) < 0)
  {
    Seobeo_Handle_Destroy(p_handle);
    return NULL;
  }

  Seobeo_Client_Response *p_resp =
      Seobeo_Client_Parse_Response(p_handle, p_req->download_path, p_req->timeout_ms);

  Seobeo_Handle_Destroy(p_handle);

  return p_resp;
}

Seobeo_Client_Response *Seobeo_Client_Request_Execute(Seobeo_Client_Request *p_req)
{
  if (!p_req)
    return NULL;

  Seobeo_Client_Response *p_resp = Seobeo_Client_Execute_Single(p_req);

  if (!p_resp)
    return NULL;

  int redirect_count = 0;
  while (p_req->follow_redirects &&
         p_resp->redirect_url &&
         (p_resp->status_code == 301 || p_resp->status_code == 302 ||
          p_resp->status_code == 303 || p_resp->status_code == 307 ||
          p_resp->status_code == 308) &&
         redirect_count < p_req->max_redirects)
  {
    char *redirect_url = malloc(strlen(p_resp->redirect_url) + 1);
    strcpy(redirect_url, p_resp->redirect_url);

    Seobeo_Client_Response_Destroy(p_resp);

    // Relative redirect
    if (redirect_url[0] == '/')
    {
      size_t path_len = strlen(redirect_url);
      p_req->path = Dowa_Arena_Allocate(p_req->p_arena, path_len + 1);
      strcpy(p_req->path, redirect_url);

      size_t url_len = strlen(p_req->host) + strlen(p_req->port) + path_len + 16;
      p_req->url = Dowa_Arena_Allocate(p_req->p_arena, url_len);
      snprintf(p_req->url, url_len, "%s://%s:%s%s",
               p_req->use_tls ? "https" : "http",
               p_req->host, p_req->port, p_req->path);
    }
    else
    {
      size_t url_len = strlen(redirect_url);
      p_req->url = Dowa_Arena_Allocate(p_req->p_arena, url_len + 1);
      strcpy(p_req->url, redirect_url);

      Seobeo_Client_Parse_Url(redirect_url, &p_req->host, &p_req->port, &p_req->path, &p_req->use_tls, p_req->p_arena);
    }

    free(redirect_url);

    p_resp = Seobeo_Client_Execute_Single(p_req);
    if (!p_resp)
      return NULL;

    redirect_count++;
  }

  return p_resp;
}

void Seobeo_Client_Request_Destroy(Seobeo_Client_Request *p_req)
{
  if (!p_req)
    return;

  if (p_req->p_arena)
    Dowa_Arena_Free(p_req->p_arena);

  // Note: headers_map and headers_array are allocated using arena functions
  // (Dowa_HashMap_Push_Arena, Dowa_Array_Push_Arena), so they are freed
  // when the arena is freed above. Do not call Dowa_HashMap_Free or
  // Dowa_Array_Free on them.

  free(p_req);
}

void Seobeo_Client_Response_Destroy(Seobeo_Client_Response *p_resp)
{
  if (!p_resp)
    return;

  if (p_resp->p_arena)
    Dowa_Arena_Free(p_resp->p_arena);

  // Note: headers are allocated using Dowa_HashMap_Push_Arena, so they are
  // freed when the arena is freed above. Do not call Dowa_HashMap_Free.

  free(p_resp);
}