#include "save_layout.h"
#include <stdlib.h>

enum { LIMIT = 1024 * 1024, NODES = 4096, DEPTH = 64 };
typedef struct {
  const uint8_t *bytes;
  size_t length, cursor, needed;
  unsigned nodes;
  int status;
} Reader;
static void fail(Reader *r, int status) { if (!r->status) r->status = status; }
static const uint8_t *take(Reader *r, size_t size) {
  if (r->status) return NULL;
  if (size > LIMIT || r->cursor > LIMIT - size) { fail(r, FACTORIO_SAVE_PREFIX_LIMIT); return NULL; }
  size_t end = r->cursor + size;
  if (end > r->length) {
    r->needed = end; fail(r, FACTORIO_SAVE_SHORT_HEADER); return NULL;
  }
  const uint8_t *result = r->bytes + r->cursor;
  r->cursor = end;
  return result;
}
static uint32_t number(Reader *r, unsigned size) {
  const uint8_t *p = take(r,size);
  uint32_t result = 0;
  if (p) for (unsigned i=0;i<size;++i) result |= (uint32_t)p[i] << (8*i);
  return result;
}
static uint32_t compact(Reader *r, unsigned extended) {
  uint32_t value = number(r,1);
  return value == 255 ? number(r,extended) : value;
}
static void string(Reader *r) { size_t count = compact(r,4); take(r,count); }
static void boolean(Reader *r) { if (number(r,1)>1) fail(r,FACTORIO_SAVE_UNSUPPORTED); }
static void version(Reader *r) { for(unsigned i=0;i<3;++i) compact(r,2); }
static void localised(Reader *r, unsigned depth) {
  if (r->status) return;
  if (depth>=DEPTH || ++r->nodes>NODES) { fail(r,FACTORIO_SAVE_PREFIX_LIMIT); return; }
  if (number(r,1)) {
    string(r);
    unsigned children = number(r,1);
    for(unsigned i=0;i<children && !r->status;++i) localised(r,depth+1);
  }
}
static void ending(Reader *r) {
  localised(r,0); localised(r,0);
  uint32_t count = compact(r,4);
  if (count>NODES) { fail(r,FACTORIO_SAVE_PREFIX_LIMIT); return; }
  for(uint32_t i=0;i<count && !r->status;++i) localised(r,0);
  localised(r,0); string(r);
}

int factorio_save_layout_bytes(const uint8_t *prefix, size_t length,
                              FactorioSaveLayout *out, size_t *needed) {
  if (!prefix || !out || !needed) return FACTORIO_SAVE_IO;
  *needed=0;
  if (length<9) { *needed=9; return FACTORIO_SAVE_SHORT_HEADER; }
  FactorioMapVersion v;
  if (factorio_version_parse(prefix,9,&v)) return FACTORIO_SAVE_SHORT_HEADER;
  if (v.major!=2 || v.minor!=0 || v.patch!=77 || v.build || v.flag)
    return FACTORIO_SAVE_UNSUPPORTED;
  Reader r = {.bytes=prefix,.length=length,.cursor=9};
  // ScenarioLocation: three compact-length strings.
  string(&r); string(&r); string(&r);
  if(number(&r,1)>2) fail(&r,FACTORIO_SAVE_UNSUPPORTED); // Difficulty
  boolean(&r); boolean(&r); string(&r);
  for(unsigned i=0;i<4;++i) boolean(&r);
  version(&r); number(&r,4); boolean(&r); number(&r,4);
  uint32_t mods=compact(&r,4);
  if(mods>NODES) fail(&r,FACTORIO_SAVE_PREFIX_LIMIT);
  for(uint32_t i=0;i<mods && !r.status;++i) { string(&r); version(&r); number(&r,4); }
  number(&r,4);
  // PropertyTree type=5 (dictionary), any-type boolean, uint32 child count.
  // Reject other/nonempty settings instead of guessing recursive wire layouts.
  if(number(&r,1)!=5) fail(&r,FACTORIO_SAVE_UNSUPPORTED);
  boolean(&r);
  if(number(&r,4)) fail(&r,FACTORIO_SAVE_UNSUPPORTED);
  ending(&r); ending(&r); boolean(&r);
  if(r.status) { *needed=r.needed; return r.status; }
  size_t position=r.cursor;
  if (!take(&r,24)) { *needed=r.needed; return r.status; }
  FactorioSaveLayout result = {.header_offset=position,.next_offset=r.cursor,.mod_count=mods};
  if(factorio_map_header_parse(prefix+position,24,v,&result.header)) return FACTORIO_SAVE_SHORT_HEADER;
  *out=result;
  return FACTORIO_SAVE_OK;
}

int factorio_save_map_header(const char *path, FactorioSaveLayout *out) {
  if(!path || !out) return FACTORIO_SAVE_IO;
  size_t requested=9;
  // Retry exactly the bytes the parser needs. Avoid assuming a fixed offset or
  // requiring a large prefix when a valid short stream contains the whole header.
  for(unsigned attempt=0;attempt<128;++attempt) {
    uint8_t *prefix=malloc(requested);
    if(!prefix) return FACTORIO_SAVE_IO;
    int status=factorio_save_prefix(path,prefix,requested);
    size_t needed=0;
    if(!status) status=factorio_save_layout_bytes(prefix,requested,out,&needed);
    free(prefix);
    if(status!=FACTORIO_SAVE_SHORT_HEADER || needed<=requested) return status;
    requested=needed;
  }
  return FACTORIO_SAVE_PREFIX_LIMIT;
}

const char *factorio_save_layout_status_string(int status) {
  switch(status) {
    case FACTORIO_SAVE_OK: return "Map header read; full world is not loaded";
    case FACTORIO_SAVE_UNSUPPORTED: return "Unsupported save metadata or archive; currently supports 2.0.77 with empty startup settings";
    case FACTORIO_SAVE_SHORT_HEADER: return "Incomplete scenario metadata or map header";
    case FACTORIO_SAVE_PREFIX_LIMIT: return "Scenario prefix, recursion or read limit reached";
    default: return factorio_save_status_string(status);
  }
}