// Image_Decode: any image file's bytes to RGBA pixels, the right way up, and the PNG previews .mcap
// files hold. stb_image reads PNG, JPEG, GIF, BMP, PSD and TGA; everything else (HEIC, AVIF, WebP,
// TIFF, camera RAW) goes to the system's own codecs: on Windows WIC, which reads what the installed
// Store extensions (HEIF, WebP, Raw Image) can, and on macOS ImageIO (imageio.jai). Linux reads only
// what stb_image does. A photo's EXIF orientation is applied, so a phone's portrait
// photo comes out standing up.
//
// What the bytes are comes from their first bytes (image_kind), never from a file name, except for
// the camera RAW formats that are TIFF inside: image_kind_from_name tells those apart.
//
// Videos (video.jai) are told apart here too, by their first bytes or their extension; video_poster
// reads a frame from one with FFmpeg, on Windows and macOS, and gif_info counts an animated GIF's frames.
//
// Thread-safe: every call stands alone, and WIC is set up and released on the calling thread; ImageIO
// needs no setting up.

Image_Kind :: enum u8 {
  UNKNOWN;
  PNG; JPEG; GIF; BMP; PSD; TGA;   // stb_image
  WEBP; HEIF; AVIF; TIFF; RAW;     // WIC on Windows, ImageIO on macOS
  VIDEO;                           // MP4, QuickTime, Matroska and WebM, AVI, WMV, MPEG: see video_poster
}

// Pixels from image_decode: RGBA, 4 bytes a pixel, rows packed. Free with image_free.
Image :: struct {
  pixels: *u8;
  width, height: s32;
  from_stb: bool;    // pixels are stb_image's, freed with stbi_image_free; otherwise the default allocator's
}

image_free :: (image: *Image) {
  if image.pixels {
    if image.from_stb  stbi_image_free(image.pixels);
    else               free(image.pixels,, context.default_allocator);
  }
  image.* = .{};
}

// What the bytes hold, from their first bytes. A TIFF may be a camera's RAW file; see image_kind_from_name.
image_kind :: (bytes: [] u8) -> Image_Kind {
  starts :: (bytes: [] u8, at: s64, magic: string) -> bool {
    if bytes.count < at + magic.count  return false;
    return memcmp(bytes.data + at, magic.data, magic.count) == 0;
  }
  if starts(bytes, 0, "\x89PNG\r\n\x1A\n")  return .PNG;
  if starts(bytes, 0, "\xFF\xD8\xFF")       return .JPEG;
  if starts(bytes, 0, "GIF87a") || starts(bytes, 0, "GIF89a")  return .GIF;
  if starts(bytes, 0, "BM") && bytes.count > 26  return .BMP;
  if starts(bytes, 0, "8BPS")               return .PSD;
  if starts(bytes, 0, "RIFF") && starts(bytes, 8, "WEBP")  return .WEBP;
  if starts(bytes, 0, "II*\0") || starts(bytes, 0, "MM\0*")  return .TIFF;
  if starts(bytes, 0, "IIRO") || starts(bytes, 0, "IIU\0") || starts(bytes, 0, "FUJIFILMCCD-RAW")  return .RAW;   // Olympus, Panasonic, Fuji
  if starts(bytes, 4, "ftyp") {
    // ISO base media: HEIF and AVIF images, Canon's CR3, or a video, by the major brand and the
    // compatible ones. AVIF first: an AVIF often names the general mif1 brand too.
    size := be32(bytes, 0);
    end  := min(cast(s64) size, bytes.count);
    if end < 16  return .UNKNOWN;
    found: [#run enum_highest_value(Image_Kind) + 1] bool;
    found[ftyp_brand_kind(slice_bytes(bytes, 8, 4))] = true;
    at := 16;
    while at + 4 <= end {
      found[ftyp_brand_kind(slice_bytes(bytes, at, 4))] = true;
      at += 4;
    }
    if found[Image_Kind.AVIF]  return .AVIF;
    if found[Image_Kind.HEIF]  return .HEIF;
    if found[Image_Kind.RAW]   return .RAW;
    return .VIDEO;   // an ftyp with no image brand: an MP4, MOV, 3GP and the like
  }
  if starts(bytes, 4, "moov") || starts(bytes, 4, "mdat") || starts(bytes, 4, "wide") || starts(bytes, 4, "free")  return .VIDEO;   // old QuickTime
  if starts(bytes, 0, "\x1A\x45\xDF\xA3")  return .VIDEO;                         // Matroska and WebM (EBML)
  if starts(bytes, 0, "RIFF") && starts(bytes, 8, "AVI ")  return .VIDEO;
  if starts(bytes, 0, "\x30\x26\xB2\x75\x8E\x66\xCF\x11")  return .VIDEO;     // ASF: WMV
  if starts(bytes, 0, "\x00\x00\x01\xBA") || starts(bytes, 0, "\x00\x00\x01\xB3")  return .VIDEO;   // MPEG program stream, or bare MPEG video
  if starts(bytes, 0, "FLV\x01")  return .VIDEO;
  if bytes.count >= 189 && bytes[0] == 0x47 && bytes[188] == 0x47  return .VIDEO;   // MPEG transport stream: a sync byte every 188
  if bytes.count >= 197 && bytes[4] == 0x47 && bytes[196] == 0x47  return .VIDEO;   // ... with a 4-byte time code before each (.m2ts, .mts)
  return .UNKNOWN;
}

// The container a video is in, for people to read: "MP4", "QuickTime", "WebM" ... from its first
// bytes and, where those are shared (MP4 and QuickTime, Matroska and WebM), its extension.
video_format_name :: (head: [] u8, name: string) -> string {
  extension := to_lower_copy(path_extension(name),, temp);
  starts :: (bytes: [] u8, at: s64, magic: string) -> bool {
    return bytes.count >= at + magic.count && memcmp(bytes.data + at, magic.data, magic.count) == 0;
  }
  if starts(head, 4, "ftyp") {
    if starts(head, 8, "qt  ")  return "QuickTime";
    if starts(head, 8, "3g")    return "3GP";
    return "MP4";
  }
  if starts(head, 4, "moov") || starts(head, 4, "mdat") || starts(head, 4, "wide") || starts(head, 4, "free")  return "QuickTime";
  if starts(head, 0, "\x1A\x45\xDF\xA3")  return ifx extension == "webm" then "WebM" else "Matroska";
  if starts(head, 0, "RIFF")  return "AVI";
  if starts(head, 0, "\x30\x26\xB2\x75")  return "WMV";
  if starts(head, 0, "FLV")  return "FLV";
  if starts(head, 0, "\x00\x00\x01")  return "MPEG";
  if head.count && (head[0] == 0x47 || (head.count > 4 && head[4] == 0x47))  return "MPEG-TS";
  return "video";
}

// The kind for bytes that image_kind cannot place, or places as TIFF, from the file's name: camera RAW
// formats built on TIFF, and TGA, which has no magic of its own. kind is image_kind's answer.
image_kind_from_name :: (kind: Image_Kind, name: string) -> Image_Kind {
  extension := to_lower_copy(path_extension(name),, temp);
  if kind == .TIFF {
    for RAW_TIFF_EXTENSIONS  if extension == it  return .RAW;
    return .TIFF;
  }
  if kind == .UNKNOWN && extension == "tga"  return .TGA;
  return kind;
}

RAW_TIFF_EXTENSIONS :: string.["dng", "cr2", "nef", "nrw", "arw", "srf", "sr2", "pef", "erf", "kdc", "dcr", "mos", "3fr", "iiq", "srw", "x3f", "mrw"];

image_kind_name :: (kind: Image_Kind) -> string {
  if kind == {
    case .PNG;   return "PNG";
    case .JPEG;  return "JPEG";
    case .GIF;   return "GIF";
    case .BMP;   return "BMP";
    case .PSD;   return "PSD";
    case .TGA;   return "TGA";
    case .WEBP;  return "WebP";
    case .HEIF;  return "HEIF";
    case .AVIF;  return "AVIF";
    case .TIFF;  return "TIFF";
    case .RAW;   return "RAW";
    case .VIDEO; return "video";
  }
  return "unknown";
}

// The image in bytes as RGBA, turned and flipped as its EXIF orientation (or, for HEIF and AVIF, its
// container) says, unless orient is false. kind may be passed when known; image_kind is used otherwise.
// On failure, why.
image_decode :: (bytes: [] u8, kind := Image_Kind.UNKNOWN, orient := true) -> Image, ok: bool, error: string {
  image: Image;
  if !bytes.count  return image, false, "it is empty";
  if kind == .UNKNOWN  kind = image_kind(bytes);
  if kind == .VIDEO  return image, false, "it is a video";

  stb_kinds :: Image_Kind.[.PNG, .JPEG, .GIF, .BMP, .PSD, .TGA, .UNKNOWN];
  decoded := false;
  error: string;
  if array_find(stb_kinds, kind) && bytes.count < S32_MAX {
    w, h, channels: s32;
    pixels := stbi_load_from_memory(bytes.data, cast(s32) bytes.count, *w, *h, *channels, 4);
    if pixels && w > 0 && h > 0 {
      image = .{pixels = pixels, width = w, height = h, from_stb = true};
      decoded = true;
    } else {
      if pixels  stbi_image_free(pixels);
      error = to_string(stbi_failure_reason());
    }
  }
  container_orientation := 0;   // ImageIO's, for HEIF and AVIF on macOS; 0 when not known
  if !decoded {
    #if OS == .WINDOWS {
      wic_ok: bool;
      wic_error: string;
      image, wic_ok, wic_error = wic_decode(bytes);
      if !wic_ok {
        if error  return image, false, tprint("% (stb_image: %)", wic_error, error);
        return image, false, wic_error;
      }
    } else #if OS == .MACOS {
      imageio_ok: bool;
      imageio_error: string;
      image, imageio_ok, imageio_error, container_orientation = imageio_decode(bytes);
      if !imageio_ok {
        if error  return image, false, tprint("% (stb_image: %)", imageio_error, error);
        return image, false, imageio_error;
      }
    } else {
      return image, false, ifx error then tprint("stb_image: %", error) else tprint("% images are only read on Windows and macOS", image_kind_name(kind));
    }
  }

  if orient {
    orientation := 1;
    if (kind == .HEIF || kind == .AVIF) && container_orientation {
      // ImageIO gives the image as stored, and the container's irot and imir as an orientation.
      orientation = container_orientation;
    } else if kind == .HEIF || kind == .AVIF {
      // The container's irot and imir are the orientation; the EXIF tag only repeats it. WIC applies
      // them (HEIF requires a reader to); a quarter turn it left out shows as the sides not swapped.
      heif := heif_info(bytes);
      if heif.rotation == 90 || heif.rotation == 270 {
        if heif.width && image.width == heif.width && image.height == heif.height {
          orientation = ifx heif.rotation == 90 then 8 else 6;   // irot turns anticlockwise
        }
      }
    } else {
      exif := exif_read(bytes, kind);
      orientation = exif.orientation;
    }
    if orientation > 1 && orientation <= 8  image_orient(*image, orientation);
  }
  return image, true, "";
}

// Turns and flips the pixels as EXIF orientation 1-8 says, into new pixels from the default allocator.
image_orient :: (image: *Image, orientation: s64) {
  if orientation <= 1 || orientation > 8 || !image.pixels  return;
  w, h := image.width, image.height;
  swaps := orientation >= 5;
  out_w := ifx swaps then h else w;
  out_h := ifx swaps then w else h;
  out := cast(*u32) alloc(cast(s64) w * h * 4,, context.default_allocator);
  source := cast(*u32) image.pixels;
  for dy: 0..out_h - 1 {
    row := out + cast(s64) dy * out_w;
    for dx: 0..out_w - 1 {
      sx, sy: s32;
      if orientation == {
        case 2;  sx, sy = w - 1 - dx, dy;            // mirrored left to right
        case 3;  sx, sy = w - 1 - dx, h - 1 - dy;    // half a turn
        case 4;  sx, sy = dx, h - 1 - dy;            // mirrored top to bottom
        case 5;  sx, sy = dy, dx;                    // transposed
        case 6;  sx, sy = dy, h - 1 - dx;            // stored a quarter turn anticlockwise: turned clockwise
        case 7;  sx, sy = w - 1 - dy, h - 1 - dx;    // transversed
        case 8;  sx, sy = w - 1 - dy, dx;            // stored a quarter turn clockwise: turned anticlockwise
      }
      row[dx] = source[cast(s64) sy * w + sx];
    }
  }
  image_free(image);
  image.* = .{pixels = cast(*u8) out, width = out_w, height = out_h};
}

// The preview an .mcap holds: the pixels scaled so the longest side is at most max_side, as a PNG
// (temp, or stb's when reused is set: see below). The width and height are the preview's.
image_preview_png :: (image: Image, max_side: s64) -> png: [] u8, width: s64, height: s64, ok: bool, error: string {
  w, h := image.width, image.height;
  pixels := image.pixels;
  longest := max(w, h);
  if longest > max_side {
    scale := cast(float64) max_side / longest;
    out_w := max(cast(s32) 1, cast(s32) (w * scale + 0.5));
    out_h := max(cast(s32) 1, cast(s32) (h * scale + 0.5));
    scaled := NewArray(cast(s64) out_w * out_h * 4, u8, initialized = false,, temp);
    if !stbir_resize_uint8_srgb(pixels, w, h, 0, scaled.data, out_w, out_h, 0, 4, 3, 0)  return .[], 0, 0, false, "its image cannot be scaled";
    pixels = scaled.data;
    w, h = out_w, out_h;
  }
  length: s32;
  encoded := stbi_write_png_to_mem(pixels, w * 4, w, h, 4, *length);
  if !encoded  return .[], 0, 0, false, "the preview cannot be encoded";
  // stb_image_write allocates with the C runtime's malloc, and stb_image frees with its free: both
  // libraries are linked against the same runtime.
  defer stbi_image_free(encoded);
  return cast([] u8) copy_string(string.{length, encoded},, temp), w, h, true, "";
}

// The width and height a PNG says it has, from its header, without decoding it.
png_size :: (png: [] u8) -> width: s64, height: s64, ok: bool {
  if png.count < 24  return 0, 0, false;
  for u8.[0x89, #char "P", #char "N", #char "G", 0x0D, 0x0A, 0x1A, 0x0A]  if png[it_index] != it  return 0, 0, false;
  return be32(png, 16), be32(png, 20), true;
}

#load "exif.jai";
#load "video.jai";
#if OS == .WINDOWS  #load "wic.jai";
#if OS == .MACOS    #load "imageio.jai";
#if OS == .WINDOWS || OS == .MACOS  #load "video_ffmpeg.jai";

#scope_module

be16 :: (b: [] u8, at: s64) -> u16 {
  if at < 0 || at + 2 > b.count  return 0;
  return ((cast(u16) b[at]) << 8) | (cast(u16) b[at + 1]);
}

be32 :: (b: [] u8, at: s64) -> u32 {
  if at < 0 || at + 4 > b.count  return 0;
  return ((cast(u32) b[at]) << 24) | ((cast(u32) b[at + 1]) << 16) | ((cast(u32) b[at + 2]) << 8) | (cast(u32) b[at + 3]);
}

be64 :: (b: [] u8, at: s64) -> u64 {
  return ((cast(u64) be32(b, at)) << 32) | (cast(u64) be32(b, at + 4));
}

le32 :: (b: [] u8, at: s64) -> u32 {
  if at < 0 || at + 4 > b.count  return 0;
  return (cast(u32) b[at]) | ((cast(u32) b[at + 1]) << 8) | ((cast(u32) b[at + 2]) << 16) | ((cast(u32) b[at + 3]) << 24);
}

// count bytes at at, cut short at the end; empty when at is past it.
slice_bytes :: (b: [] u8, at: s64, count: s64) -> [] u8 {
  if at < 0 || at >= b.count || count <= 0  return .[];
  result: [] u8;
  result.data  = b.data + at;
  result.count = min(count, b.count - at);
  return result;
}

ftyp_brand_kind :: (brand: [] u8) -> Image_Kind {
  if brand.count != 4  return .UNKNOWN;
  text := cast(string) brand;
  if text == "avif" || text == "avis"  return .AVIF;
  if text == "heic" || text == "heix" || text == "hevc" || text == "hevx" || text == "heim" || text == "heis" || text == "mif1" || text == "msf1"  return .HEIF;
  if text == "crx "  return .RAW;   // Canon CR3
  return .UNKNOWN;
}

#import "Basic";
#import "String";
#import "Math";
#import "stb_image";
#import "stb_image_write";
#import "stb_image_resize";
#if OS == .WINDOWS  #import "Windows";