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

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);
  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;

  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_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)
{
  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);
  if (!p_resp->p_arena)
  {
    free(p_resp);
    return NULL;
  }

  while (1)
  {
    int r = Seobeo_Handle_Read(p_handle);
    if (r < 0)
      return p_resp;
    if (r == -2)
      break;

    if (p_handle->read_buffer_len >= 4 && strstr((char*)p_handle->read_buffer, "\r\n\r\n") != NULL)
      break;

    if (r == 0)
      continue;
  }

  char *buf = (char*)p_handle->read_buffer;
  char *hdr_end = strstr(buf, "\r\n\r\n");
  if (!hdr_end)
    return p_resp;

  size_t hdr_len = hdr_end - buf + 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);

  void *p_cl_kv = Dowa_HashMap_Get_Ptr(p_resp->headers, "Content-Length");
  size_t body_len = 0;
  if (p_cl_kv)
  {
    const char *content_length_str = ((Seobeo_Request_Entry*)p_cl_kv)->value;
    body_len = atoi(content_length_str);
  }

  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);
      }

      if (total_read < body_len)
      {
        int r = Seobeo_Handle_Read(p_handle);
        if (r < 0 || r == -2)
          break;
        if (r == 0)
          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 * 8;
    size_t used = 0;
    char *body = download_path ? NULL : Dowa_Arena_Allocate(p_resp->p_arena, cap);

    while (1)
    {
      int n = Seobeo_Handle_Read(p_handle);
      if (n > 0)
      {
        if (download_path)
        {
          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)
        continue;
      else
        break;
    }

    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);

  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);

  if (p_req->headers_map)
    Dowa_HashMap_Free(p_req->headers_map);

  if (p_req->headers_array)
    Dowa_Array_Free(p_req->headers_array);

  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);

  if (p_resp->headers)
    Dowa_HashMap_Free(p_resp->headers);

  free(p_resp);
}