// .mcap: one capture in one file. A fixed magic, a file version stamp, the capture's Capture Text,
// its OCR text (version 2), its front chunks (version 3), a small preview, the images, then its back
// chunks (version 3). See docs/mcap.md for the layout.
//
// Chunks hold what models add (labels, embeddings, later masks), each named by a tag and a key (the
// model's name): small ones in front, where a library browser reads them without the images, large
// ones at the back. Chunks this build does not know are kept as they are when a file is rewritten.
//
// Reading gives views into the bytes it was given; nothing is copied. Images stay PNG.
//
// The main copy is mcap's (D:/jai_projects/active/mcap/modules/Capture_File); SnagitConverter has
// the same files.

MCAP_MAGIC     :: u8.[0x89, #char "M", #char "C", #char "A", #char "P", 0x0D, 0x0A, 0x1A];
MCAP_VERSION   :: 3;   // what is written; every version from 1 up is read
MCAP_EXTENSION :: "mcap";
MCAP_PREVIEW_MAX_SIDE :: 320;   // the preview's longest side, in pixels

Mcap_Rendered_Kind :: enum u8 {
  NONE             :: 0;
  SAME_AS_ORIGINAL :: 1;   // the rendered image is byte for byte the original, so it is stored once
  OWN              :: 2;
}

Mcap_Image :: struct {
  width:  u16;
  height: u16;
  png: [] u8;   // empty when there is none
}

// What a model added to a capture, in its own format. A file holds one chunk per tag and key.
Mcap_Chunk :: struct {
  tag: [4] u8;                   // what it is, in four ASCII letters, such as "labl" or "embd"
  key: string;                   // whose: the model's name
  data: [] u8;
}

Mcap :: struct {
  version: u32 = MCAP_VERSION;   // as read; always MCAP_VERSION when written
  text: string;                  // Capture Text
  ocr: string;                   // OCR text (ocr_text.jai), one section per model; "" when none, and always in version 1
  front_chunks: [] Mcap_Chunk;   // small: read with the front. None before version 3
  preview: Mcap_Image;
  original_name: string;         // the original image's name in the .snagx, such as {GUID}.png
  original: [] u8;               // PNG
  rendered_kind: Mcap_Rendered_Kind;
  rendered: [] u8;               // PNG; OWN only, see mcap_rendered
  back_chunks: [] Mcap_Chunk;    // large: after the images. None before version 3
}

// "labl" as a tag.
mcap_tag :: (text: string) -> [4] u8 {
  assert(text.count == 4, "a chunk tag is four letters");
  tag: [4] u8;
  for 0..3  tag[it] = text[it];
  return tag;
}

mcap_tag_text :: (tag: [4] u8) -> string {
  return tprint("%", to_string(tag.data, 4));
}

// The chunk with tag and key in chunks, or null.
mcap_find_chunk :: (chunks: [] Mcap_Chunk, tag: [4] u8, key: string) -> *Mcap_Chunk {
  for * chunks  if mcap_same_tag(it.tag, tag) && it.key == key  return it;
  return null;
}

mcap_same_tag :: (a: [4] u8, b: [4] u8) -> bool {
  for a  if it != b[it_index]  return false;
  return true;
}

// The rendered image (Snagit's thumbnail.png) whichever way it is stored; empty when there is none.
mcap_rendered :: (m: Mcap) -> [] u8 {
  if m.rendered_kind == {
    case .SAME_AS_ORIGINAL;  return m.original;
    case .OWN;               return m.rendered;
  }
  return .[];
}

// ---------------------------------------------------------------- writing

// The file's bytes. A part too large for its size field fails the whole file rather than writing one
// that cannot be read back.
mcap_serialize :: (m: Mcap) -> Serializer, ok: bool, error: string {
  size := MCAP_HEADER_SIZE + 4 + m.text.count + 4 + m.ocr.count + mcap_chunks_size(m.front_chunks) + 8 + m.preview.png.count + 2 + m.original_name.count + 4 + m.original.count + 1;
  if m.rendered_kind == .OWN  size += 4 + m.rendered.count;
  size += mcap_chunks_size(m.back_chunks);
  s := new_serializer(size);

  magic := MCAP_MAGIC;
  AddArray_NoSize(*s, magic);
  Add(*s, cast(u32) MCAP_VERSION);

  if !AddString(*s, m.text, u32)  return s, false, "the text is too large for its size field";
  if !AddString(*s, m.ocr, u32)   return s, false, "the OCR text is too large for its size field";
  ok, error := mcap_add_chunks(*s, m.front_chunks);
  if !ok  return s, false, error;
  Add(*s, m.preview.width);
  Add(*s, m.preview.height);
  if !AddSizedArray(*s, m.preview.png, u32)  return s, false, "the preview is too large for its size field";
  if !AddString(*s, m.original_name, u16)  return s, false, "the original image's name is too long";
  if !AddSizedArray(*s, m.original, u32)  return s, false, "the original image is too large for its size field";
  Add(*s, m.rendered_kind);
  if m.rendered_kind == .OWN && !AddSizedArray(*s, m.rendered, u32)  return s, false, "the rendered image is too large for its size field";
  ok, error = mcap_add_chunks(*s, m.back_chunks);
  if !ok  return s, false, error;

  assert(s.count == size);
  return s, true, "";
}

// Writes beside path and then swaps it in, so a failure leaves any earlier file as it was.
mcap_write :: (path: string, bytes: [] u8) -> ok: bool, error: string {
  temp_path := tprint("%.tmp", path);
  if !write_entire_file(temp_path, bytes.data, bytes.count)  return false, tprint("cannot write %", temp_path);
  if !atomic_file_replace_or_just_rename_if_destination_file_does_not_exist(temp_path, path) {
    file_delete(temp_path);
    return false, tprint("cannot move % to %", temp_path, path);
  }
  return true, "";
}

// ---------------------------------------------------------------- reading

// A whole .mcap. The parts are views into data. Exactly the file's bytes must be consumed.
mcap_parse :: (data: [] u8) -> Mcap, ok: bool, error: string {
  m: Mcap;
  d := to_deserializer(data);
  ok, error := mcap_read_front_part(*d, *m);
  if !ok  return m, false, error;

  m.original_name = ReadSizedString(*d, u16);
  m.original      = ReadSizedArray(*d, u8, u32);
  m.rendered_kind = Read(*d, Mcap_Rendered_Kind);
  if d.ok && m.rendered_kind == .OWN  m.rendered = ReadSizedArray(*d, u8, u32);
  if d.ok && m.version >= 3 {
    m.back_chunks, ok, error = mcap_read_chunks(*d);
    if !ok  return m, false, error;
  }
  if !d.ok  return m, false, "the file ends early (it is truncated or damaged)";
  if cast(u8) m.rendered_kind > cast(u8) Mcap_Rendered_Kind.OWN  return m, false, tprint("the rendered image is of unknown kind %", cast(u8) m.rendered_kind);
  if d.cursor != d.count  return m, false, tprint("% bytes follow the last part", d.count - d.cursor);
  return m, true, "";
}

// Only the front of the file, as a library browser needs it: the header, text, OCR text, front chunks
// and, with preview, the preview. One read of the start of the file, and a few more only when a part runs past
// it. The images are left unread and empty. The parts are views into a buffer from the context allocator.
mcap_read_front :: (path: string, preview := true) -> Mcap, ok: bool, error: string {
  m: Mcap;
  file, opened := file_open(path, log_errors = false);
  if !opened  return m, false, "cannot be opened";
  defer file_close(*file);
  length, length_ok := file_length(file);
  if !length_ok  return m, false, "cannot be read";

  wanted := min(length, MCAP_FRONT_READ_SIZE);
  for attempt: 1..MCAP_FRONT_READS {
    buffer := NewArray(wanted, u8, initialized = false);
    if !file_read(file, buffer.data, wanted) { array_free(buffer); return m, false, "cannot be read"; }

    d := to_deserializer(buffer);
    ok, error := mcap_read_front_part(*d, *m, preview);
    if ok  return m, true, "";
    if d.ok || wanted == length { array_free(buffer); return m, false, error; }   // not a matter of too few bytes

    // Short: the sizes read so far say how far the front reaches at least. Read that and as much
    // again as the first read, so the size fields after it (the OCR, the chunk count, each chunk's
    // size) come in the same read rather than one read each.
    needed := mcap_front_size(buffer, preview);
    array_free(buffer);
    if needed <= wanted || needed > length  return m, false, "the file ends early (it is truncated or damaged)";
    wanted = min(needed + MCAP_FRONT_READ_SIZE, length);
    if !file_set_position(file, 0)  return m, false, "cannot be read";
  }
  return m, false, tprint("its front was not all read in % reads", MCAP_FRONT_READS);
}

#scope_module

MCAP_HEADER_SIZE     :: 12;          // magic and version
MCAP_FRONT_READ_SIZE :: 64 * 1024;   // holds the header, text and preview of nearly every capture
MCAP_FRONT_READS     :: 16;          // reads of the front at most: each reaches past at least one more size field

// Header, text, OCR text, front chunks and, with preview, the preview.
mcap_read_front_part :: (d: *Deserializer, m: *Mcap, preview := true) -> ok: bool, error: string {
  magic := ReadArrayView(d, MCAP_MAGIC.count, u8);
  if !d.ok  return false, "is not an .mcap file (it is shorter than the magic)";
  for MCAP_MAGIC  if magic[it_index] != it  return false, "is not an .mcap file (its magic is wrong)";
  m.version = Read(d, u32);
  if !d.ok  return false, "the file ends early (it is truncated or damaged)";
  if m.version == 0  return false, "has file version 0, which does not exist";
  if m.version > MCAP_VERSION  return false, tprint("is file version %, but this build reads up to %; it was written by a later version", m.version, MCAP_VERSION);

  m.text           = ReadSizedString(d, u32);
  if m.version >= 2  m.ocr = ReadSizedString(d, u32);   // version 1 has no OCR block
  if d.ok && m.version >= 3 {                           // nor versions before 3 chunks
    ok: bool;
    error: string;
    m.front_chunks, ok, error = mcap_read_chunks(d);
    if !ok  return false, error;
  }
  if !preview {
    if !d.ok  return false, "the file ends early (it is truncated or damaged)";
    return true, "";
  }
  m.preview.width  = Read(d, u16);
  m.preview.height = Read(d, u16);
  m.preview.png    = ReadSizedArray(d, u8, u32);
  if !d.ok  return false, "the file ends early (it is truncated or damaged)";
  return true, "";
}

// How many bytes the front takes, as far as the size fields in prefix tell: the whole front when
// they are all in it, otherwise how far to read to reach the next size field.
mcap_front_size :: (prefix: [] u8, preview := true) -> s64 {
  d := to_deserializer(prefix);
  d.cursor = MCAP_MAGIC.count;
  version := Read(*d, u32);
  if !d.ok  return MCAP_HEADER_SIZE;
  position: s64 = MCAP_HEADER_SIZE;
  skip_sized :: (prefix: [] u8, position: *s64) -> reached: bool {   // a u32 size and what it counts
    if position.* + 4 > prefix.count { position.* += 4; return false; }
    d := to_deserializer(prefix);
    d.cursor = position.*;
    position.* += 4 + cast(s64) Read(*d, u32);
    return true;
  }
  if !skip_sized(prefix, *position)  return position;                       // the text
  if version >= 2 && !skip_sized(prefix, *position)  return position;     // the OCR text
  if version >= 3 {                                                        // the front chunks
    if position + 4 > prefix.count  return position + 4;
    d.cursor = position;
    count := Read(*d, u32);
    position += 4;
    for 1..count {
      position += 4;                                                       // the tag
      if !skip_sized(prefix, *position)  return position;                 // its size and payload
    }
  }
  if !preview  return position;
  position += 4;                                                           // the preview's width and height
  skip_sized(prefix, *position);                                           // the preview
  return position;
}

// Chunks: a u32 count, then each as its tag, a u32 size, and that many bytes: a u16 key size, the
// key, and the data.
mcap_chunks_size :: (chunks: [] Mcap_Chunk) -> s64 {
  size: s64 = 4;
  for chunks  size += 4 + 4 + 2 + it.key.count + it.data.count;
  return size;
}

mcap_add_chunks :: (s: *Serializer, chunks: [] Mcap_Chunk) -> ok: bool, error: string {
  Add(s, cast(u32) chunks.count);
  for chunks {
    payload := 2 + it.key.count + it.data.count;
    if it.key.count > 0xFFFF || payload > 0xFFFF_FFFF  return false, tprint("the % chunk of % is too large for its size field", mcap_tag_text(it.tag), it.key);
    tag := it.tag;
    AddArray_NoSize(s, tag);
    Add(s, cast(u32) payload);
    AddString(s, it.key, u16);
    AddArray_NoSize(s, it.data);
  }
  return true, "";
}

// Views into d's bytes; the list itself from the context allocator.
mcap_read_chunks :: (d: *Deserializer) -> [] Mcap_Chunk, ok: bool, error: string {
  count := Read(d, u32);
  if !d.ok  return .[], false, "the file ends early (it is truncated or damaged)";
  // Running out of bytes clears d.ok, as Read does: a front read then reads further, a whole parse fails.
  if cast(s64) count * 10 > d.count - d.cursor { d.ok = false; return .[], false, tprint("it says it has % chunks, more than its bytes can hold", count); }
  chunks := NewArray(count, Mcap_Chunk);
  for * chunk: chunks {
    tag := ReadArrayView(d, 4, u8);
    payload := Read(d, u32);
    if !d.ok || payload < 2  return chunks, false, "a chunk runs past the end of the file (it is truncated or damaged)";
    if cast(s64) payload > d.count - d.cursor { d.ok = false; return chunks, false, "a chunk runs past the end of the file (it is truncated or damaged)"; }
    for 0..3  chunk.tag[it] = tag[it];
    chunk.key = ReadSizedString(d, u16);
    if !d.ok || 2 + chunk.key.count > cast(s64) payload  return chunks, false, "a chunk's key runs past the chunk";
    chunk.data = ReadArrayView(d, payload - 2 - chunk.key.count, u8);
  }
  if !d.ok  return chunks, false, "the file ends early (it is truncated or damaged)";
  return chunks, true, "";
}

#scope_export

// Stores what a scan adds to the .mcap at path in one rewrite: with ocr, that OCR section in place of
// any by the same model (the others kept); and each of front, in place of any front chunk with the
// same tag and key (every other chunk kept, known to this build or not). Upgrades the file to the
// current version. Rewritten beside itself and swapped in (mcap_write), so a failure leaves it as it was.
mcap_store :: (path: string, ocr: *Mcap_Ocr_Section = null, front: [] Mcap_Chunk = .[]) -> ok: bool, error: string {
  data, read_ok := read_entire_file(path, log_errors = false);
  if !read_ok  return false, "cannot be read";
  defer free(data);
  m, parsed, parse_error := mcap_parse(cast([] u8) data);
  if !parsed  return false, parse_error;

  push_allocator(temp);
  if ocr {
    sections, ok, error := mcap_ocr_parse(m.ocr);
    if !ok  return false, error;
    replaced := false;
    for * sections  if it.model == ocr.model { it.* = ocr.*; replaced = true; break; }
    if !replaced  array_add(*sections, ocr.*);
    m.ocr = mcap_ocr_text(sections);
  }
  m.front_chunks = mcap_chunks_with(m.front_chunks, front);
  m.version = MCAP_VERSION;

  serializer, serialized, serialize_error := mcap_serialize(m);
  if !serialized  return false, serialize_error;
  written, write_error := mcap_write(path, serializer);
  return written, write_error;
}

// Stores chunk alone: in front, or with back after the images. See mcap_store.
mcap_store_chunk :: (path: string, chunk: Mcap_Chunk, back := false) -> ok: bool, error: string {
  if !back {
    ok, error := mcap_store(path, front = .[chunk]);
    return ok, error;
  }
  data, read_ok := read_entire_file(path, log_errors = false);
  if !read_ok  return false, "cannot be read";
  defer free(data);
  m, parsed, parse_error := mcap_parse(cast([] u8) data);
  if !parsed  return false, parse_error;
  push_allocator(temp);
  m.back_chunks = mcap_chunks_with(m.back_chunks, .[chunk]);
  m.version = MCAP_VERSION;
  serializer, serialized, serialize_error := mcap_serialize(m);
  if !serialized  return false, serialize_error;
  written, write_error := mcap_write(path, serializer);
  return written, write_error;
}

// chunks with each of added in place of the one with its tag and key, or after them. Temp.
mcap_chunks_with :: (chunks: [] Mcap_Chunk, added: [] Mcap_Chunk) -> [] Mcap_Chunk {
  list: [..] Mcap_Chunk;
  list.allocator = temp;
  array_add(*list, ..chunks);
  for added {
    existing := mcap_find_chunk(list, it.tag, it.key);
    if existing  existing.* = it;
    else         array_add(*list, it);
  }
  return list;
}

#load "ocr_text.jai";
#load "label_chunks.jai";

#import "Basic";
#import "File";
#import "File_Utilities";
#import "Serialization";