#include "save_header.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <zlib.h>

enum { PREFIX_LIMIT = 1024 * 1024, CHUNK = 4096 };
static uint16_t u16(const unsigned char *p) { return (uint16_t)(p[0] | ((uint16_t)p[1] << 8)); }
static uint32_t u32(const unsigned char *p) { return (uint32_t)u16(p) | ((uint32_t)u16(p + 2) << 16); }
static int read_at(FILE *f, uint64_t position, void *out, size_t length) {
  if (position > INT64_MAX || fseeko(f, (off_t)position, SEEK_SET)) return 0;
  return fread(out, 1, length, f) == length;
}
static int suffix(const char *name, size_t length, const char *tail) {
  size_t count = strlen(tail);
  return length >= count && !memcmp(name + length - count, tail, count) &&
    (length == count || name[length - count - 1] == '/');
}

typedef struct { uint32_t local, compressed; uint16_t method, flags, name_length; char name[4097]; } Entry;
static int find_level(FILE *file, uint64_t file_size, Entry *selected) {
  if (file_size < 22) return FACTORIO_SAVE_INVALID_ZIP;
  size_t count = file_size < 65557 ? (size_t)file_size : 65557;
  unsigned char *tail = malloc(count);
  if (!tail) return FACTORIO_SAVE_IO;
  if (!read_at(file, file_size - count, tail, count)) { free(tail); return FACTORIO_SAVE_IO; }
  size_t end = count;
  for (size_t i = count - 22 + 1; i-- > 0;) {
    if (u32(tail + i) == 0x06054b50 && (size_t)u16(tail + i + 20) == count - i - 22) {
      end = i; break;
    }
  }
  if (end == count) { free(tail); return FACTORIO_SAVE_INVALID_ZIP; }
  unsigned char *e = tail + end;
  uint32_t central_size = u32(e + 12), central_offset = u32(e + 16);
  uint16_t entries = u16(e + 10);
  if (u16(e + 4) || u16(e + 6) || entries != u16(e + 8) || entries == 0xffff ||
      central_size == 0xffffffff || central_offset == 0xffffffff) {
    free(tail); return FACTORIO_SAVE_UNSUPPORTED;
  }
  uint64_t central_end = (uint64_t)central_offset + central_size;
  if (central_end > file_size - count + end) { free(tail); return FACTORIO_SAVE_INVALID_ZIP; }
  free(tail);
  uint64_t cursor = central_offset;
  int priority = 0;
  for (unsigned index = 0; index < entries; ++index) {
    unsigned char header[46];
    if (cursor > central_end || central_end - cursor < sizeof(header) ||
        !read_at(file, cursor, header, sizeof(header)) || u32(header) != 0x02014b50) return FACTORIO_SAVE_INVALID_ZIP;
    uint16_t name_length = u16(header + 28);
    uint64_t next = cursor + 46 + name_length + u16(header + 30) + u16(header + 32);
    if (next > central_end || name_length > 4096) return FACTORIO_SAVE_UNSUPPORTED;
    char name[4097];
    if (!read_at(file, cursor + 46, name, name_length)) return FACTORIO_SAVE_IO;
    if (memchr(name, 0, name_length)) return FACTORIO_SAVE_INVALID_ZIP;
    name[name_length] = 0;
    int match = suffix(name, name_length, "level.dat0") ? 2 : suffix(name, name_length, "level.dat");
    if (match && match >= priority) {
      if (match == priority) return FACTORIO_SAVE_INVALID_ZIP; /* ambiguous roots */
      if (u16(header + 34) || u32(header + 20) == 0xffffffff || u32(header + 42) == 0xffffffff)
        return FACTORIO_SAVE_UNSUPPORTED;
      *selected = (Entry){.local = u32(header + 42), .compressed = u32(header + 20),
        .method = u16(header + 10), .flags = u16(header + 8), .name_length = name_length};
      memcpy(selected->name, name, name_length + 1);
      priority = match;
    }
    cursor = next;
  }
  if (!priority) return FACTORIO_SAVE_NO_LEVEL;
  if ((selected->flags & 0x41) || (selected->method != 0 && selected->method != 8)) return FACTORIO_SAVE_UNSUPPORTED;
  return FACTORIO_SAVE_OK;
}

static int feed_inner(z_stream *inner, unsigned char *bytes, size_t length, unsigned char *header, size_t requested) {
  inner->next_in = bytes; inner->avail_in = (uInt)length;
  inner->next_out = header + inner->total_out; inner->avail_out = (uInt)(requested - inner->total_out);
  int status = inflate(inner, Z_NO_FLUSH);
  if (status != Z_OK && status != Z_STREAM_END && status != Z_BUF_ERROR) return FACTORIO_SAVE_COMPRESSION;
  if (inner->total_out == requested) return FACTORIO_SAVE_OK;
  return status == Z_STREAM_END ? FACTORIO_SAVE_SHORT_HEADER : -1;
}

static int prefix(FILE *file, uint64_t file_size, const Entry *entry, unsigned char *header, size_t requested) {
  unsigned char local[30];
  if (!read_at(file, entry->local, local, sizeof(local)) || u32(local) != 0x04034b50) return FACTORIO_SAVE_INVALID_ZIP;
  if (u16(local + 8) != entry->method || u16(local + 6) != entry->flags ||
      u16(local + 26) != entry->name_length) return FACTORIO_SAVE_INVALID_ZIP;
  char name[4097];
  if (!read_at(file, (uint64_t)entry->local + 30, name, entry->name_length) ||
      memcmp(name, entry->name, entry->name_length)) return FACTORIO_SAVE_INVALID_ZIP;
  uint64_t data_offset = (uint64_t)entry->local + 30 + u16(local + 26) + u16(local + 28);
  if (data_offset > file_size || entry->compressed > file_size - data_offset) return FACTORIO_SAVE_INVALID_ZIP;
  z_stream outer = {0}, inner = {0};
  if (inflateInit(&inner) != Z_OK) return FACTORIO_SAVE_COMPRESSION;
  if (entry->method == 8 && inflateInit2(&outer, -MAX_WBITS) != Z_OK) {
    inflateEnd(&inner); return FACTORIO_SAVE_COMPRESSION;
  }
  int result = FACTORIO_SAVE_SHORT_HEADER;
  unsigned char input[CHUNK], decoded[CHUNK];
  uint64_t consumed = 0, produced = 0;
  while (consumed < entry->compressed && consumed < PREFIX_LIMIT) {
    size_t amount = entry->compressed - consumed < CHUNK ? (size_t)(entry->compressed - consumed) : CHUNK;
    if (!read_at(file, data_offset + consumed, input, amount)) { result = FACTORIO_SAVE_IO; break; }
    consumed += amount;
    if (entry->method == 0) {
      result = feed_inner(&inner, input, amount, header, requested);
      if (result >= 0) goto finish;
    } else {
      outer.next_in = input; outer.avail_in = (uInt)amount;
      do {
        outer.next_out = decoded; outer.avail_out = sizeof(decoded);
        int status = inflate(&outer, Z_NO_FLUSH);
        if (status != Z_OK && status != Z_STREAM_END && status != Z_BUF_ERROR) {
          result = FACTORIO_SAVE_COMPRESSION; goto finish;
        }
        size_t bytes = sizeof(decoded) - outer.avail_out;
        produced += bytes;
        if (produced > PREFIX_LIMIT) { result = FACTORIO_SAVE_PREFIX_LIMIT; goto finish; }
        result = feed_inner(&inner, decoded, bytes, header, requested);
        if (result >= 0) goto finish;
        if (status == Z_STREAM_END) { result = FACTORIO_SAVE_SHORT_HEADER; goto finish; }
        if (!bytes && !outer.avail_in) break;
        if (!bytes && status == Z_BUF_ERROR) { result = FACTORIO_SAVE_COMPRESSION; goto finish; }
      } while (outer.avail_in || !outer.avail_out);
    }
  }
  result = consumed >= PREFIX_LIMIT ? FACTORIO_SAVE_PREFIX_LIMIT : FACTORIO_SAVE_SHORT_HEADER;
finish:
  if (entry->method == 8) inflateEnd(&outer);
  inflateEnd(&inner);
  return result;
}

int factorio_save_prefix(const char *path, uint8_t *out, size_t length) {
  if (!path || !out) return FACTORIO_SAVE_IO;
  if (!length || length > PREFIX_LIMIT) return FACTORIO_SAVE_PREFIX_LIMIT;
  FILE *file = fopen(path, "rb");
  if (!file) return FACTORIO_SAVE_IO;
  if (fseeko(file, 0, SEEK_END)) { fclose(file); return FACTORIO_SAVE_IO; }
  off_t end = ftello(file);
  if (end < 0) { fclose(file); return FACTORIO_SAVE_IO; }
  Entry entry;
  int status = find_level(file, (uint64_t)end, &entry);
  unsigned char *temporary = NULL;
  if (!status) {
    temporary = malloc(length);
    if (!temporary) status = FACTORIO_SAVE_IO;
    else status = prefix(file, (uint64_t)end, &entry, temporary, length);
  }
  fclose(file);
  if (!status) memcpy(out, temporary, length);
  free(temporary);
  return status;
}

int factorio_save_header(const char *path, FactorioMapVersion *out) {
  if (!out) return FACTORIO_SAVE_IO;
  unsigned char header[9];
  int status = factorio_save_prefix(path, header, sizeof(header));
  if (!status && factorio_version_parse(header, 9, out)) status = FACTORIO_SAVE_SHORT_HEADER;
  return status;
}

const char *factorio_save_status_string(int status) {
  switch (status) {
    case FACTORIO_SAVE_OK: return "Version header read; full map is not loaded";
    case FACTORIO_SAVE_IO: return "Cannot read the selected file";
    case FACTORIO_SAVE_INVALID_ZIP: return "Invalid or ambiguous ZIP archive";
    case FACTORIO_SAVE_NO_LEVEL: return "No level.dat0 or level.dat entry";
    case FACTORIO_SAVE_UNSUPPORTED: return "Unsupported archive layout, compression or encryption";
    case FACTORIO_SAVE_COMPRESSION: return "Invalid compressed level header";
    case FACTORIO_SAVE_SHORT_HEADER: return "Incomplete map version header";
    case FACTORIO_SAVE_PREFIX_LIMIT: return "Header decompression limit reached";
    default: return "Unknown save-header error";
  }
}