/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <cstring>
#include <fstream>
#include <iostream>
#include <optional>
#include <sstream>
#include <stdexcept>
#include <unordered_map>

#include "se_defines.h"

SE_GLOBALS()

class FormatError : public std::runtime_error {
 public:
  explicit FormatError(std::string f) : std::runtime_error(f) {}
};

class Reader {
 public:
  std::string data;
  size_t offset;
  explicit Reader(std::string str) : data(str), offset(0) {}

  uint32_t read_u32() {
    uint32_t res = 0;
    std::string raw_bytes = data.substr(offset, 4);
    if (raw_bytes.size() != 4) {
      throw FormatError("couldn't read u32");
    }
    SE_STEP(true);
    std::memcpy(&res, raw_bytes.c_str(), 4);
    offset += 4;
    return res;
  }

  std::string read_str(uint32_t len) {
    std::string res = data.substr(offset, len);
    if (res.size() != len) {
      throw FormatError("couldn't read string");
    }
    SE_STEP(true);
    offset += len;
    return res;
  }
};

class DataFrame {};

class DataHeader : public DataFrame {
 public:
  uint32_t size_x;
  uint32_t size_y;

  static const uint32_t MAGIC = 0x11EADE13;

  static DataHeader parse_data(Reader& rd) {
    DataHeader res;
    uint32_t size = rd.read_u32();
    if (size != 8) {
      throw FormatError("wrong_size");
    }
    SE_STEP(true);
    res.size_x = rd.read_u32();
    res.size_y = rd.read_u32();
    return res;
  }
};

class DataImage : public DataFrame {
 public:
  std::string data;
  static const uint32_t MAGIC = 0xDA7ADA7A;
  static DataImage parse_data(Reader& rd) {
    DataImage res;
    uint32_t size = rd.read_u32();
    if (size < 64) {
      throw FormatError("data too small");
    }
    SE_STEP(true);
    res.data = rd.read_str(size);
    return res;
  }
};

class DataMeta : public DataFrame {
 public:
  std::string section;
  std::unordered_map<std::string, std::string> metadata;

  static const uint32_t MAGIC = 0x11E7A000;

  DataMeta() : section("") {}

  DataMeta(const DataMeta& other)
      : section(other.section), metadata(other.metadata) {}

  DataMeta(
      std::string section,
      std::unordered_map<std::string, std::string> metadata)
      : section(section), metadata(metadata) {}

  static DataMeta parse_data(Reader& rd) {
    std::unordered_map<std::string, std::string> metadata;
    std::string section = rd.read_str(rd.read_u32());
    uint32_t num_elems = rd.read_u32();
    for (int i = 0; i < num_elems; i++) {
      std::string key = rd.read_str(rd.read_u32());
      std::string value = rd.read_str(rd.read_u32());
      metadata[key] = value;
    }
    return DataMeta(section, metadata);
  }
};

class DataFile {
 public:
  DataHeader header;
  DataImage image;
  std::unordered_map<std::string, DataMeta> meta;

  DataFile(
      DataHeader hdr,
      DataImage img,
      std::unordered_map<std::string, DataMeta> meta)
      : header(hdr), image(img), meta(meta) {}

  std::optional<std::string> resolve(std::string name) {
    size_t pos = name.find('.');
    if (pos == std::string::npos) {
      return std::nullopt;
    }
    SE_STEP(true)
    std::string sec = name.substr(0, pos);
    std::string key = name.substr(pos + 1);
    if (meta.find(sec) != meta.end()) {
      if (meta[sec].metadata.find(key) != meta[sec].metadata.end()) {
        return meta[sec].metadata[key];
      }
    }
    return std::nullopt;
  }
};

std::optional<DataFile> parse_file(Reader& rd) {
  if (rd.read_u32() != 0xC0DEC0DE) {
    throw FormatError("wrong_magic");
  }
  SE_STEP(true)
  uint32_t len = rd.read_u32();
  std::optional<DataHeader> header = std::nullopt;
  if (len < 2) {
    throw FormatError("wrong number of tags");
  }
  SE_STEP(true)

  std::unordered_map<std::string, DataMeta> meta;
  for (int i = 0; i < len; i++) {
    uint32_t tag = rd.read_u32();
    if (tag == DataHeader::MAGIC) {
      header = DataHeader::parse_data(rd);
      SE_STEP(true)
    } else if (tag == DataImage::MAGIC) {
      if (!header.has_value()) {
        throw FormatError("no_header");
      }
      SE_STEP(true)
      auto image = DataImage::parse_data(rd);
      if (image.data.size() != header->size_x * header->size_y) {
        throw FormatError("wrong_size");
      }
      SE_STEP(true)
      return DataFile(*header, image, meta);
    } else if (tag == DataMeta::MAGIC) {
      auto meta_frame = DataMeta::parse_data(rd);
      meta[meta_frame.section] = meta_frame;
    } else {
      throw FormatError("unknown_tag");
    }
  }
  return std::nullopt;
}

bool test_case(std::string str) {
  try {
    auto rd = Reader(str);
    auto file = parse_file(rd);
    if (file.has_value()) {
      SE_STEP(true)
      if (file->meta.size() >= 1) {
        SE_STEP(true)
        SE_TARGET_STATE((file->resolve("section.foo") == "bar"));
      }
    }
  } catch (FormatError& e) {
    std::cerr << "Error: " << e.what() << "\n";
    return false;
  }
  return true;
}

int main(int argc, char* argv[]) {
  std::istream* input;
  std::ifstream file;
  if (argc < 2) {
    input = &std::cin; // No file name provided, read from stdin
  } else {
    file.open(argv[1]);
    if (!file) {
      std::cerr << "Could not open the file.\n";
      return 1;
    }
    input = &file; // File opened successfully, read from file
  }
  std::string content(
      (std::istreambuf_iterator<char>(*input)),
      std::istreambuf_iterator<char>());
  test_case(content);
  return 0;
}
