(* Ocsigen
 * ocsigen_stream.ml Copyright (C) 2005 Vincent Balat
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, with linking exception;
 * either version 2.1 of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
*)

open Ocsigen_lib

exception Interrupted of exn
exception Cancelled
exception Already_read
exception Finalized

type 'a stream = 'a step Lwt.t Lazy.t

and 'a step =
  | Finished of 'a stream option (* If there is another stream following
                                    (useful for substreams) *)
  | Cont of 'a * 'a stream       (* Current buffer, what follows *)

type outcome = [`Success | `Failure]

type 'a t =
  { mutable stream : 'a stream;
    mutable in_use : bool;
    mutable finalizer : outcome -> unit Lwt.t }

let net_buffer_size = ref 8192

let set_net_buffer_size i = net_buffer_size := i

let empty follow =
  match follow with
    None    -> Lwt.return (Finished None)
  | Some st -> Lwt.return (Finished (Some (Lazy.from_fun st)))

let cont stri f =
  Lwt.return (Cont (stri, Lazy.from_fun f))

let make ?finalize:(g = fun _ -> Lwt.return ()) f =
  { stream = Lazy.from_fun f; in_use = false; finalizer = g }

let next = Lazy.force

let rec get_aux st =
  lazy
    (Lwt.try_bind
       (fun () -> Lazy.force st.stream)
       (fun e ->
          Lwt.return
            (match e with
               Cont (s, rem) -> st.stream <- rem; Cont (s, get_aux st)
             | _             -> e))
       (fun e ->
          st.stream <- lazy (Lwt.fail e);
          Lwt.fail (Interrupted e)))

let get st =
  if st.in_use then raise Already_read;
  st.in_use <- true;
  get_aux st

(** read the stream until the end, without decoding *)
let rec consume_aux st =
  next st >>= fun e ->
  match e with
  | Cont (_, f)        -> consume_aux f
  | Finished (Some ss) -> consume_aux ss
  | Finished None      -> Lwt.return ()

let cancel st =
  let st' = st.stream in
  st.stream <- lazy (Lwt.fail Cancelled);
  consume_aux st'

let consume st =
  consume_aux st.stream

let finalize st status =
  let f = st.finalizer in
  st.finalizer <- (fun _ -> Lwt.return ());
  f status >>= fun () ->
  st.stream <- lazy (Lwt.fail Finalized);
  Lwt.return ()

let add_finalizer st g =
  let f = st.finalizer in
  st.finalizer <- fun status -> f status >>= fun () -> g status

(****)

(** String streams *)

exception Stream_too_small
exception Stream_error of string
exception String_too_large

let string_of_stream m s =
  let buff = Buffer.create (m/4) in
  let rec aux i s =
    next s >>= function
    | Finished _ -> Lwt.return buff
    | Cont (s, f) ->
      let i = i + String.length s in
      if i > m
      then Lwt.fail String_too_large
      else (Buffer.add_string buff s; aux i f)
  in
  aux 0 s >|= Buffer.contents

let enlarge_stream = function
  | Finished _a -> Lwt.fail Stream_too_small
  | Cont (s, f) ->
    let long = String.length s in
    let max = !net_buffer_size in
    if long >= max
    then Lwt.fail Input_is_too_large
    else
      next f >>= fun e ->
      match e with
      | Finished _ -> Lwt.fail Stream_too_small
      | Cont (r, ff) ->
        let long2 = String.length r in
        let long3=long+long2 in
        let new_s = s^r in
        if long3 <= max
        then Lwt.return (Cont (new_s, ff))
        else let long4 = long3 - max in
          cont (String.sub new_s 0 max)
            (fun () ->
               Lwt.return (Cont (String.sub new_s max long4, ff)))

let rec stream_want s len =
  (* returns a stream with at least len bytes read if possible *)
  match s with
  | Finished _ -> Lwt.return s
  | Cont (stri, _f)  ->
    if String.length stri >= len then
      Lwt.return s
    else
      Lwt.catch
        (fun () -> enlarge_stream s >>= fun r -> Lwt.return (`OK r))
        (function
          | Stream_too_small -> Lwt.return `Too_small
          | e -> Lwt.fail e)
      >>= function
      | `OK r -> stream_want r len
      | `Too_small -> Lwt.return s


let current_buffer = function
  | Finished _  -> raise Stream_too_small
  | Cont (s, _) -> s

let rec skip s k = match s with
  | Finished _ -> raise Stream_too_small
  | Cont (s, f) ->
    let len = String.length s in
    let len64 = Int64.of_int len in
    if Int64.compare k len64 <= 0
    then
      let k = Int64.to_int k in
      Lwt.return (Cont (String.sub s k (len - k), f))
    else (enlarge_stream (Cont ("", f)) >>=
          (fun s -> skip s (Int64.sub k len64)))

let substream delim s =
  let ldelim = String.length delim in
  if ldelim = 0 then Lwt.fail (Stream_error "Empty delimiter")
  else
    let rdelim = Pcre.(regexp (quote delim)) in
    let rec aux =
      function
      | Finished _ -> Lwt.fail Stream_too_small
      | Cont (s, f) as stre ->
        let len = String.length s in
        if len < ldelim
        then enlarge_stream stre >>= aux
        else try
            let p,_ = Ocsigen_lib.Netstring_pcre.search_forward rdelim s 0 in
            cont (String.sub s 0 p)
              (fun () ->
                 empty
                   (Some (fun () -> Lwt.return (Cont (String.sub s p (len - p),
                                                      f)))))
          with Not_found ->
            let pos = (len + 1 - ldelim) in
            cont (String.sub s 0 pos)
              (fun () ->
                 next f >>= function
                 | Finished _ -> Lwt.fail Stream_too_small
                 | Cont (s', f') ->
                   aux
                     (Cont (String.sub s pos (len - pos) ^ s',
                            f'))
              )
    in aux s

(*****************************************************************************)

(*VVV Is it the good place for this? *)

let of_file filename =
  let fd = Lwt_unix.of_unix_file_descr
      (Unix.openfile filename [Unix.O_RDONLY;Unix.O_NONBLOCK] 0o666)
  in
  let ch = Lwt_io.(of_fd ~mode:input) fd in
  let buf = Bytes.create 1024 in
  let rec aux () =
    Lwt_io.read_into ch buf 0 1024 >>= fun n ->
    if n = 0 then empty None else
      (* Streams should be immutable, thus we always make a copy
         of the buffer *)
      cont (Bytes.sub_string buf 0 n) aux
  in make ~finalize:(fun _ -> Lwt_unix.close fd) aux

let of_string s =
  make (fun () -> cont s (fun () -> empty None))

(** Convert a {!Lwt_stream.t} to an {!Ocsigen_stream.t}. *)
let of_lwt_stream stream =
  let rec aux () =
    Lwt_stream.get stream >>= function
    | Some e -> cont e aux
    | None -> empty None
  in make aux

(** Convert an {!Ocsigen_stream.t} into a {!Lwt_stream.t}.
    @param is_empty function to skip empty chunk.
*)
let to_lwt_stream ?(is_empty = (fun _ -> false)) o_stream =
  let stream = ref (get o_stream) in
  let rec wrap () =
    next !stream >>= function
    | Finished None -> o_stream.finalizer `Success >>= fun () -> Lwt.return None
    | Finished (Some next) -> stream := next; wrap ()
    | Cont (value, next) ->
      stream := next;
      if is_empty value
      then wrap ()
      else Lwt.return (Some value)
  in Lwt_stream.from wrap

module StringStream = struct

  type out = string t
  type m = (string stream -> string step Lwt.t) Lazy.t

  let empty : m = lazy (fun c -> Lazy.force c)
  let concat (m: m) (f: m) : m =
    lazy (fun c -> Lazy.force m (lazy (Lazy.force f c)))
  let put (s : string) : m = lazy (fun c -> Lwt.return (Cont (s, c)))

  let make_stream (m: m) : string stream =
    lazy (Lazy.force m (lazy (Lwt.return (Finished None))))

  let make (m: m) : out = make (fun () -> Lazy.force (make_stream m))

end
