#define STRVIEW_IMPLEMENTATION
#include "strview.h"

#include "probe.h"
#include "platform.h"

#include <ctype.h>
#include <errno.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.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 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 bool json_string(strview input, char* output, size_t capacity) {
  sv_trim_left(&input);
  if (input.len == 0 || input.data[0] != '"' || capacity == 0) return false;
  size_t written = 0;
  for (size_t i = 1; i < input.len && input.data[i] != '"'; i++) {
    char character = input.data[i];
    if (character == '\\' && i + 1 < input.len) {
      character = input.data[++i];
      if (character == 'n') character = '\n';
      else if (character == 'r') character = '\r';
      else if (character == 't') character = '\t';
    }
    if (written + 1 < capacity) output[written++] = character;
  }
  output[written] = '\0';
  return true;
}

static size_t matching_delimiter(strview input, char open, char close) {
  unsigned depth = 0;
  bool quoted = false;
  bool escaped = false;
  for (size_t i = 0; i < input.len; i++) {
    char character = input.data[i];
    if (quoted) {
      if (escaped) escaped = false;
      else if (character == '\\') escaped = true;
      else if (character == '"') quoted = false;
      continue;
    }
    if (character == '"') quoted = true;
    else if (character == open) depth++;
    else if (character == close && depth != 0 && --depth == 0) return i;
  }
  return SIZE_MAX;
}

static SubtitleKind subtitle_kind_for_codec(const char* codec) {
  if (strcmp(codec, "ass") == 0 || strcmp(codec, "ssa") == 0) return SUBTITLE_ASS;
  if (strcmp(codec, "dvd_subtitle") == 0 || strcmp(codec, "dvb_subtitle") == 0 ||
      strcmp(codec, "hdmv_pgs_subtitle") == 0 || strcmp(codec, "xsub") == 0) {
    return SUBTITLE_BITMAP;
  }
  return SUBTITLE_TEXT;
}

static bool json_truthy(strview object, const char* key) {
  strview value;
  return json_value(object, sv_from_cstr(key), &value) &&
         (sv_starts_with(value, svlit("1")) || sv_starts_with(value, svlit("true")));
}

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\""));
  bool attachment = sv_starts_with(type_value, svlit("\"attachment\""));
  if (!video && !audio && !subtitles && !attachment) {
    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);
  char codec_name[32] = {0};
  strview codec_value;
  if (json_value(object, svlit("\"codec_name\""), &codec_value)) {
    sv_trim_left(&codec_value);
    if (codec_value.len != 0 && codec_value.data[0] == '"') {
      sv_chop_left(&codec_value, 1);
      size_t length = 0;
      while (length < codec_value.len && codec_value.data[length] != '"' &&
             length + 1 < sizeof(codec_name)) length++;
      memcpy(codec_name, codec_value.data, length);
      codec_name[length] = '\0';
    }
  }
  if (attachment) {
    if (result->attachment_count < PROBE_ATTACHMENT_LIMIT) {
      ProbeAttachment* item = &result->attachments[result->attachment_count++];
      item->stream_index = (size_t)raw_index;
      strview value;
      if (json_value(object, svlit("\"filename\""), &value))
        json_string(value, item->filename, sizeof(item->filename));
      if (json_value(object, svlit("\"mimetype\""), &value))
        json_string(value, item->mimetype, sizeof(item->mimetype));
    }
    return;
  }
  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,
    };
    memcpy(result->streams[result->stream_count - 1].codec_name, codec_name, sizeof(codec_name));
    ProbeStream* stream = &result->streams[result->stream_count - 1];
    stream->subtitle_kind = subtitle_kind_for_codec(codec_name);
    strview value;
    if (json_value(object, svlit("\"language\""), &value))
      json_string(value, stream->language, sizeof(stream->language));
    if (json_value(object, svlit("\"title\""), &value))
      json_string(value, stream->title, sizeof(stream->title));
    if (json_truthy(object, "\"default\"")) stream->dispositions |= DISPOSITION_DEFAULT;
    if (json_truthy(object, "\"forced\"")) stream->dispositions |= DISPOSITION_FORCED;
    if (json_truthy(object, "\"hearing_impaired\"")) stream->dispositions |= DISPOSITION_HEARING_IMPAIRED;
    if (json_truthy(object, "\"visual_impaired\"")) stream->dispositions |= DISPOSITION_VISUAL_IMPAIRED;
    if (json_truthy(object, "\"comment\"")) stream->dispositions |= DISPOSITION_COMMENT;
    if (json_truthy(object, "\"original\"")) stream->dispositions |= DISPOSITION_ORIGINAL;
    if (json_truthy(object, "\"dub\"")) stream->dispositions |= DISPOSITION_DUB;
    if (json_truthy(object, "\"lyrics\"")) stream->dispositions |= DISPOSITION_LYRICS;
    if (json_truthy(object, "\"karaoke\"")) stream->dispositions |= DISPOSITION_KARAOKE;
    if (json_truthy(object, "\"clean_effects\"")) stream->dispositions |= DISPOSITION_CLEAN_EFFECTS;
    if (json_truthy(object, "\"attached_pic\"")) stream->dispositions |= DISPOSITION_ATTACHED_PIC;
    if (json_truthy(object, "\"timed_thumbnails\"")) stream->dispositions |= DISPOSITION_TIMED_THUMBNAILS;
    if (json_truthy(object, "\"non_diegetic\"")) stream->dispositions |= DISPOSITION_NON_DIEGETIC;
    if (json_truthy(object, "\"captions\"")) stream->dispositions |= DISPOSITION_CAPTIONS;
    if (json_truthy(object, "\"descriptions\"")) stream->dispositions |= DISPOSITION_DESCRIPTIONS;
    if (json_truthy(object, "\"metadata\"")) stream->dispositions |= DISPOSITION_METADATA;
    if (json_truthy(object, "\"dependent\"")) stream->dispositions |= DISPOSITION_DEPENDENT;
    if (json_truthy(object, "\"still_image\"")) stream->dispositions |= DISPOSITION_STILL_IMAGE;
    if (json_truthy(object, "\"multilayer\"")) stream->dispositions |= DISPOSITION_MULTILAYER;
  }
  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 void parse_tags(strview object, ProbeResult* result) {
  strview tags;
  if (!json_value(object, svlit("\"tags\""), &tags)) return;
  bool found = false;
  size_t start = sv_find_char(tags, '{', &found);
  if (!found) return;
  sv_chop_left(&tags, start + 1);
  size_t end = matching_delimiter(sv_from_data(tags.data - 1, tags.len + 1), '{', '}');
  if (end == SIZE_MAX) return;
  tags.len = end - 1;
  while (tags.len != 0 && result->tag_count < PROBE_TAG_LIMIT) {
    sv_trim_left(&tags);
    if (tags.len == 0 || tags.data[0] == '}') break;
    char key[64] = {0};
    if (!json_string(tags, key, sizeof(key))) break;
    size_t quote = 1;
    bool escaped = false;
    while (quote < tags.len) {
      if (!escaped && tags.data[quote] == '"') break;
      escaped = !escaped && tags.data[quote] == '\\';
      if (tags.data[quote] != '\\') escaped = false;
      quote++;
    }
    sv_chop_left(&tags, quote + 1);
    bool colon_found = false;
    size_t colon = sv_find_char(tags, ':', &colon_found);
    if (!colon_found) break;
    sv_chop_left(&tags, colon + 1);
    sv_trim_left(&tags);
    char value[256] = {0};
    if (!json_string(tags, value, sizeof(value))) break;
    ProbeTag* tag = &result->tags[result->tag_count++];
    snprintf(tag->key, sizeof(tag->key), "%s", key);
    snprintf(tag->value, sizeof(tag->value), "%s", value);
    if (tags.len != 0 && tags.data[0] == '"') {
      size_t value_end = 1;
      bool value_escape = false;
      while (value_end < tags.len) {
        if (!value_escape && tags.data[value_end] == '"') break;
        value_escape = !value_escape && tags.data[value_end] == '\\';
        if (tags.data[value_end] != '\\') value_escape = false;
        value_end++;
      }
      sv_chop_left(&tags, value_end + 1);
    }
    sv_trim_left(&tags);
    if (tags.len != 0 && tags.data[0] == ',') sv_chop_left(&tags, 1);
  }
}

static void parse_chapter(strview object, ProbeResult* result) {
  if (result->chapter_count >= PROBE_CHAPTER_LIMIT) return;
  strview value;
  double start = 0, end = 0;
  if (!json_value(object, svlit("\"start_time\""), &value) || !json_number(value, &start) ||
      !json_value(object, svlit("\"end_time\""), &value) || !json_number(value, &end)) return;
  ProbeChapter* chapter = &result->chapters[result->chapter_count++];
  chapter->start = start;
  chapter->end = end;
  if (json_value(object, svlit("\"title\""), &value))
    json_string(value, chapter->title, sizeof(chapter->title));
}

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);
    size_t end_offset = matching_delimiter(sv_from_data(streams.data - 1, streams.len + 1), '{', '}');
    if (end_offset == SIZE_MAX) {
      return false;
    }
    parse_stream(sv_from_data(streams.data, end_offset - 1), result);
    sv_chop_left(&streams, end_offset);
  }

  strview format = sv_from_data(json.data + format_offset, json.len - format_offset);
  bool format_open_found = false;
  size_t format_open = sv_find_char(format, '{', &format_open_found);
  size_t format_end = format_open_found
                        ? matching_delimiter(sv_from_data(format.data + format_open,
                                                         format.len - format_open), '{', '}')
                        : SIZE_MAX;
  if (format_end == SIZE_MAX) {
    return false;
  }
  format = sv_from_data(format.data + format_open + 1, format_end - 1);
  parse_tags(format, result);
  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;
      }
    }
  }
  bool chapters_found = false;
  size_t chapters_offset = sv_find(json, svlit("\"chapters\""), &chapters_found);
  if (chapters_found) {
    strview chapters = sv_from_data(json.data + chapters_offset, json.len - chapters_offset);
    bool open_found = false;
    size_t open = sv_find_char(chapters, '[', &open_found);
    if (open_found) {
      sv_chop_left(&chapters, open + 1);
      while (chapters.len != 0) {
        bool object_found = false;
        size_t object = sv_find_char(chapters, '{', &object_found);
        if (!object_found) break;
        sv_chop_left(&chapters, object);
        size_t close = matching_delimiter(chapters, '{', '}');
        if (close == SIZE_MAX) break;
        parse_chapter(sv_from_data(chapters.data + 1, close - 1), result);
        sv_chop_left(&chapters, close + 1);
      }
    }
  }
  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;
  char* const argv[] = {
    "ffprobe", "-v", "error", "-show_entries",
    "format=duration:format_tags:stream=index,codec_name,codec_type,duration:"
    "stream_tags=language,title,filename,mimetype:stream_disposition:"
    "chapter=start_time,end_time:chapter_tags=title", "-show_chapters", "-of", "json",
    (char*)path, NULL,
  };
  char* output = NULL;
  int exit_code = 0;
  if (!platform_capture(argv, &output, &exit_code, error, "ffprobe") || exit_code != 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;
}