(***********************************************************************)
(*                                                                     *)
(*                           Objective Caml                            *)
(*                                                                     *)
(*            Franois Pessaux, projet Cristal, INRIA Rocquencourt     *)
(*            Pierre Weis, projet Cristal, INRIA Rocquencourt          *)
(*            Jun Furuse, projet Cristal, INRIA Rocquencourt           *)
(*                                                                     *)
(*  Copyright 1999-2004                                                *)
(*  Institut National de Recherche en Informatique et en Automatique.  *)
(*  Distributed only by permission.                                    *)
(*                                                                     *)
(***********************************************************************)

(* $Id: jpeg.ml,v 1.16 2004/09/07 07:31:05 weis Exp $ *)

open Images;;
open Rgb24;;

external read : string -> int * int * string
    = "read_JPEG_file";;
external write : string -> string -> int -> int -> int -> unit
    = "write_JPEG_file";;

type in_handle;;

external open_in_header : string -> int * int * in_handle
    = "open_jpeg_file_for_read";;
external set_scale_denom : in_handle -> int -> unit
    = "jpeg_set_scale_denom";;
external open_in_start : in_handle -> int * int * in_handle
    = "open_jpeg_file_for_read_start";;
external read_scanline : in_handle -> string -> unit
    = "read_jpeg_scanline";;
external close_in : in_handle -> unit
    = "close_jpeg_file_for_read";;

type out_handle;;

external open_out : string -> int -> int -> int -> out_handle
    = "open_jpeg_file_for_write";;
external open_out_cmyk : string -> int -> int -> int -> out_handle
    = "open_jpeg_file_for_write_cmyk";;
external write_scanline : out_handle -> string -> unit
    = "write_jpeg_scanline";;
external close_out : out_handle -> unit
    = "close_jpeg_file_for_write";;

let open_in name =
  let _, _, ic = open_in_header name in
  open_in_start ic;;

let open_in_thumbnail name geom_spec =
  if geom_spec.Geometry.spec_aspect = Geometry.Dont_keep then
    raise (Invalid_argument "Jpeg.open_in_thumbnail: illegal geom_spec");
  let image_width, image_height, ic = open_in_header name in
  let scale =
    try
      let geom = Geometry.compute geom_spec image_width image_height in
(*
      prerr_endline (Printf.sprintf "Denom %d/%d" image_width geom.Geometry.geom_width);
*)
      image_width / geom.Geometry.geom_width
    with
    | _ -> 1 in
  let denom =
    if scale < 2 then 1 else
    if scale < 4 then 2 else
    if scale < 8 then 4 else 8 in
(*
  prerr_endline (Printf.sprintf "Denom %d" denom);
*)
  set_scale_denom ic denom;
  image_width, image_height, open_in_start ic;;

let load name load_opts =
  let prog = Images.load_progress load_opts in
  let w, h, ic = open_in name in
  let img = Rgb24.create w h in
  let scanline = String.create (w * 3) in
  for y = 0 to h - 1 do
    read_scanline ic scanline;
    Rgb24.set_scanline img y scanline;
    match prog with
    | Some p -> p (float (y + 1) /. float h)
    | None -> ()
  done;
  close_in ic;
  Rgb24 img;;

let load_thumbnail name load_opts geom_spec =
  let prog = Images.load_progress load_opts in
  let ow, oh, (w, h, ic) = open_in_thumbnail name geom_spec in
  let img = Rgb24.create w h in
  let scanline = String.create (w * 3) in
  for y = 0 to h - 1 do
    read_scanline ic scanline;
    Rgb24.set_scanline img y scanline;
    match prog with
    | Some p -> p (float (y + 1) /. float h)
    | None -> ()
  done;
  close_in ic;
  ow, oh, Rgb24 img;;

let save name opts image =
  let quality =
    match Images.save_quality opts with
    | Some q -> q
    | None -> 80
  in
  let prog = Images.save_progress opts in
  match image with
  | Rgb24 bmp ->
      let oc = open_out name bmp.width bmp.height quality in
      for y = 0 to bmp.height - 1 do
        write_scanline oc (Rgb24.get_scanline bmp y);
        match prog with
        | Some p -> p (float (y + 1) /. float bmp.height)
        | None -> ()
      done;
      close_out oc
  | _ -> raise Wrong_image_type;;

let save_as_cmyk name opts trans image =
  let quality =
    match Images.save_quality opts with
    | Some q -> q
    | None -> 80 in
  let prog = Images.save_progress opts in
  let get_cmyk_scanline width scanline =
    let buf = String.create (width * 4) in
    for x = 0 to width - 1 do
      let r = int_of_char scanline.[x * 3 + 0] in
      let g = int_of_char scanline.[x * 3 + 1] in
      let b = int_of_char scanline.[x * 3 + 2] in
      let c, m, y, k = trans {r = r; g = g; b = b} in
      buf.[x * 4 + 0] <- char_of_int (255 - c);
      buf.[x * 4 + 1] <- char_of_int (255 - m);
      buf.[x * 4 + 2] <- char_of_int (255 - y);
      buf.[x * 4 + 3] <- char_of_int (255 - k)
    done;
    buf in
  match image with
  | Rgb24 bmp ->
      let oc = open_out_cmyk name bmp.width bmp.height quality in
      for y = 0 to bmp.height - 1 do
        let buf = get_cmyk_scanline bmp.width (Rgb24.get_scanline bmp y) in
        write_scanline oc buf;
        match prog with
        | Some p -> p (float (y + 1) /. float bmp.height)
        | None -> ()
      done;
      close_out oc
  | _ -> raise Wrong_image_type;;

let save_cmyk_sample name opts =
  let quality =
    match Images.save_quality opts with
    | Some q -> q
    | None -> 80 in
  let prog = Images.save_progress opts in
  let sample_point x y =
    let c = x / 16 * 17
    and m = (x mod 16) * 17
    and y = y / 16 * 17
    and k = (y mod 16) * 17 in
    c, m, y, k in
  let sample_scan y =
    let s = String.create (256 * 4) in
    for x = 0 to 255 do
      let c, m, y, k = sample_point x y in
      s.[x * 4 + 0] <- char_of_int c;
      s.[x * 4 + 1] <- char_of_int m;
      s.[x * 4 + 2] <- char_of_int y;
      s.[x * 4 + 3] <- char_of_int k;
    done;
    s in
  let oc = open_out_cmyk name 256 256 quality in
  for y = 0 to 256 - 1 do
    let buf = sample_scan y in
    write_scanline oc buf
  done;
  close_out oc;;

let rec find_jpeg_size ic =
  (* jump to the next 0xff *)
  let rec loop () =
    let rec jump_to_0xff () =
      let ch = int_of_char (input_char ic) in
      if ch <> 0xff then jump_to_0xff () in
    let rec jump_to_non_0xff () =
      let ch = int_of_char (input_char ic) in
      if ch = 0xff then jump_to_non_0xff ()
      else ch in
    jump_to_0xff ();
    let ch = jump_to_non_0xff () in
    let str = String.create 4 in
    match ch with
    | 0xda -> raise Not_found
    | _ when ch >= 0xc0 && ch <= 0xc3 ->
      really_input ic str 0 3;
      really_input ic str 0 4;
      int_of_char str.[2] lsl 8 + int_of_char str.[3], (* width *)
      int_of_char str.[0] lsl 8 + int_of_char str.[1]  (* height *)
    | _ ->
      (* skip this block *)
      let blocklen =
        really_input ic str 0 2;
        int_of_char str.[0] lsl 8 + int_of_char str.[1] in
      let s = String.create (blocklen - 2) in
      really_input ic s 0 (blocklen - 2);
      loop () in
  try loop () with
  | _ -> raise Not_found (* any error returns Not_found *);;

let check_header filename =
  let len = 2 in
  let ic = open_in_bin filename in
  try
    let str = String.create len in
    really_input ic str 0 len;
    if
      (* some jpeg's start with 7f58, the 7th bit is missing *)
      int_of_char str.[0] lor 0x80 = 0xff &&
      int_of_char str.[1] lor 0x80 = 0xd8
      (* int_of_char str.[0] = 0xff && int_of_char str.[1] = 0xd8 *)
      (* && String.sub str 6 4 = "JFIF" --- strict *) then begin
      let w, h =
        try find_jpeg_size ic with
        | Not_found -> -1, -1 in
      Pervasives.close_in ic;
      { header_width = w;
        header_height = h;
        header_infos = []; }
    end else
      raise Wrong_file_type
  with
  | _ ->
      Pervasives.close_in ic;
      raise Wrong_file_type;;

add_methods Jpeg
  { check_header = check_header;
    load = Some load;
    save = Some save;
    load_sequence = None;
    save_sequence = None};;