#define STRVIEW_IMPLEMENTATION
#include "strview.h"

#include "probe.h"

#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

struct ProbeCacheEntry {
  struct ProbeCacheEntry* next;
  const char* path;
  ProbeResult result;
};

static bool probe_error(Error* error, const char* message, const char* path) {
  error->line = 0;
  snprintf(error->message, sizeof(error->message), message, path);
  return false;
}

static char* arena_copy(Arena* arena, const char* text) {
  size_t size = strlen(text) + 1;
  char* copy = arena_alloc(arena, size);
  if (copy != NULL) {
    memcpy(copy, text, size);
  }
  return copy;
}

void probe_cache_init(ProbeCache* cache, ProbeRunner runner, void* context) {
  *cache = (ProbeCache){
    .arena = arena_create(),
    .runner = runner != NULL ? runner : probe_run_ffprobe,
    .runner_context = context,
  };
}

bool probe_cache_get(ProbeCache* cache, const char* path, ProbeResult* result, Error* error) {
  for (ProbeCacheEntry* entry = cache->entries; entry != NULL; entry = entry->next) {
    if (sv_equal(sv_from_cstr(entry->path), sv_from_cstr(path))) {
      *result = entry->result;
      return true;
    }
  }

  ProbeResult probed = {0};
  cache->invocation_count++;
  if (!cache->runner(path, &probed, cache->runner_context, error)) {
    return false;
  }
  ProbeCacheEntry* entry = arena_alloc(&cache->arena, sizeof(*entry));
  char* cached_path = arena_copy(&cache->arena, path);
  if (entry == NULL || cached_path == NULL) {
    return probe_error(error, "out of memory while caching '%s'", path);
  }
  *entry = (ProbeCacheEntry){.next = cache->entries, .path = cached_path, .result = probed};
  cache->entries = entry;
  *result = probed;
  return true;
}

void probe_cache_destroy(ProbeCache* cache) {
  arena_destroy(&cache->arena);
  *cache = (ProbeCache){0};
}

static bool append_bytes(char** buffer, size_t* length, size_t* capacity, const char* data, size_t count) {
  if (count > SIZE_MAX - *length - 1) {
    return false;
  }
  size_t needed = *length + count + 1;
  if (needed > *capacity) {
    size_t new_capacity = *capacity == 0 ? 4096 : *capacity;
    while (new_capacity < needed) {
      if (new_capacity > SIZE_MAX / 2) {
        new_capacity = needed;
        break;
      }
      new_capacity *= 2;
    }
    char* grown = realloc(*buffer, new_capacity);
    if (grown == NULL) {
      return false;
    }
    *buffer = grown;
    *capacity = new_capacity;
  }
  memcpy(*buffer + *length, data, count);
  *length += count;
  (*buffer)[*length] = '\0';
  return true;
}

static bool read_process(int fd, char** output) {
  char* buffer = NULL;
  size_t length = 0;
  size_t capacity = 0;
  char chunk[4096];
  for (;;) {
    ssize_t count = read(fd, chunk, sizeof(chunk));
    if (count == 0) {
      break;
    }
    if (count < 0) {
      if (errno == EINTR) {
        continue;
      }
      free(buffer);
      return false;
    }
    if (!append_bytes(&buffer, &length, &capacity, chunk, (size_t)count)) {
      free(buffer);
      return false;
    }
  }
  if (buffer == NULL) {
    buffer = calloc(1, 1);
  }
  *output = buffer;
  return buffer != NULL;
}

static bool json_value(strview object, strview key, strview* value) {
  bool found = false;
  size_t offset = sv_find(object, key, &found);
  if (!found) {
    return false;
  }
  sv_chop_left(&object, offset + key.len);
  while (object.len != 0 && (sv_space_predicate(object.data[0]) || object.data[0] == ':')) {
    sv_chop_left(&object, 1);
  }
  *value = object;
  return object.len != 0;
}

static bool json_number(strview input, double* value) {
  sv_trim_left(&input);
  if (input.len != 0 && input.data[0] == '"') {
    sv_chop_left(&input, 1);
  }
  size_t length = 0;
  while (length < input.len &&
         (sv_numeric_predicate(input.data[length]) || input.data[length] == '.' ||
          input.data[length] == '-' || input.data[length] == '+' ||
          input.data[length] == 'e' || input.data[length] == 'E')) {
    length++;
  }
  if (length == 0 || length >= 128) {
    return false;
  }
  char buffer[128];
  memcpy(buffer, input.data, length);
  buffer[length] = '\0';
  char* after = NULL;
  errno = 0;
  *value = strtod(buffer, &after);
  return errno == 0 && after == buffer + length && *value >= 0;
}

static void parse_stream(strview object, ProbeResult* result) {
  strview type_value;
  if (!json_value(object, svlit("\"codec_type\""), &type_value)) {
    return;
  }
  bool video = sv_starts_with(type_value, svlit("\"video\""));
  bool audio = sv_starts_with(type_value, svlit("\"audio\""));
  bool subtitles = sv_starts_with(type_value, svlit("\"subtitle\""));
  if (!video && !audio && !subtitles) {
    return;
  }

  double raw_index = 0;
  strview index_value;
  if (!json_value(object, svlit("\"index\""), &index_value) ||
      !json_number(index_value, &raw_index) || raw_index > (double)SIZE_MAX) return;

  double duration = 0;
  strview duration_value;
  bool known = json_value(object, svlit("\"duration\""), &duration_value) &&
               json_number(duration_value, &duration);
  ValueType type = video ? TYPE_VIDEO : audio ? TYPE_AUDIO : TYPE_SUBTITLES;
  if (result->stream_count < PROBE_STREAM_LIMIT) {
    result->streams[result->stream_count++] = (ProbeStream){
      .type = type,
      .stream_index = (size_t)raw_index,
      .duration_known = known,
      .duration = duration,
    };
  }
  if (video) {
    result->has_video = true;
    if (known && (!result->video_duration_known || duration > result->video_duration)) {
      result->video_duration_known = true;
      result->video_duration = duration;
    }
  } else if (audio) {
    result->has_audio = true;
    if (known && (!result->audio_duration_known || duration > result->audio_duration)) {
      result->audio_duration_known = true;
      result->audio_duration = duration;
    }
  } else {
    result->has_subtitles = true;
    if (known && (!result->subtitle_duration_known || duration > result->subtitle_duration)) {
      result->subtitle_duration_known = true;
      result->subtitle_duration = duration;
    }
  }
}

static bool parse_probe_json(const char* text, ProbeResult* result) {
  strview json = sv_from_cstr(text);
  bool streams_found = false;
  bool format_found = false;
  size_t streams_offset = sv_find(json, svlit("\"streams\""), &streams_found);
  size_t format_offset = sv_find(json, svlit("\"format\""), &format_found);
  if (!streams_found || !format_found || streams_offset >= format_offset) {
    return false;
  }

  strview streams = sv_from_data(json.data + streams_offset, format_offset - streams_offset);
  bool array_found = false;
  size_t array_offset = sv_find_char(streams, '[', &array_found);
  if (!array_found) {
    return false;
  }
  sv_chop_left(&streams, array_offset + 1);
  while (streams.len != 0) {
    bool object_found = false;
    size_t object_offset = sv_find_char(streams, '{', &object_found);
    if (!object_found) {
      break;
    }
    sv_chop_left(&streams, object_offset + 1);
    bool end_found = false;
    size_t end_offset = sv_find_char(streams, '}', &end_found);
    if (!end_found) {
      return false;
    }
    parse_stream(sv_from_data(streams.data, end_offset), result);
    sv_chop_left(&streams, end_offset + 1);
  }

  strview format = sv_from_data(json.data + format_offset, json.len - format_offset);
  bool format_end_found = false;
  size_t format_end = sv_find_char(format, '}', &format_end_found);
  if (!format_end_found) {
    return false;
  }
  format.len = format_end;
  strview duration_value;
  double container_duration = 0;
  if (json_value(format, svlit("\"duration\""), &duration_value) &&
      json_number(duration_value, &container_duration)) {
    if (result->has_video && !result->video_duration_known) {
      result->video_duration_known = true;
      result->video_duration = container_duration;
    }
    if (result->has_audio && !result->audio_duration_known) {
      result->audio_duration_known = true;
      result->audio_duration = container_duration;
    }
    if (result->has_subtitles && !result->subtitle_duration_known) {
      result->subtitle_duration_known = true;
      result->subtitle_duration = container_duration;
    }
    for (size_t i = 0; i < result->stream_count; i++) {
      if (!result->streams[i].duration_known) {
        result->streams[i].duration_known = true;
        result->streams[i].duration = container_duration;
      }
    }
  }
  return result->has_video || result->has_audio || result->has_subtitles;
}

bool probe_run_ffprobe(const char* path, ProbeResult* result, void* context, Error* error) {
  (void)context;
  int pipe_fds[2];
  if (pipe(pipe_fds) != 0) {
    return probe_error(error, "cannot create pipe for '%s'", path);
  }

  pid_t child = fork();
  if (child < 0) {
    close(pipe_fds[0]);
    close(pipe_fds[1]);
    return probe_error(error, "cannot start ffprobe for '%s'", path);
  }
  if (child == 0) {
    close(pipe_fds[0]);
    dup2(pipe_fds[1], STDOUT_FILENO);
    dup2(pipe_fds[1], STDERR_FILENO);
    close(pipe_fds[1]);
    char* const argv[] = {
      "ffprobe", "-v", "error", "-show_entries",
      "format=duration:stream=index,codec_type,duration", "-of", "json", (char*)path, NULL,
    };
    execvp(argv[0], argv);
    _exit(127);
  }

  close(pipe_fds[1]);
  char* output = NULL;
  bool read_ok = read_process(pipe_fds[0], &output);
  close(pipe_fds[0]);
  int status = 0;
  while (waitpid(child, &status, 0) < 0 && errno == EINTR) {
  }
  if (!read_ok || !WIFEXITED(status) || WEXITSTATUS(status) != 0) {
    free(output);
    return probe_error(error, "ffprobe failed for '%s'", path);
  }

  *result = (ProbeResult){0};
  bool parsed = parse_probe_json(output, result);
  free(output);
  if (!parsed) {
    return probe_error(error, "could not parse ffprobe output for '%s'", path);
  }
  return true;
}