(***********************************************************************)
(*                                                                     *)
(*                               TAToo                                 *)
(*                                                                     *)
(*                     Kim Nguyen, LRI UMR8623                         *)
(*                   Université Paris-Sud & CNRS                       *)
(*                                                                     *)
(*  Copyright 2010-2012 Université Paris-Sud and Centre National de la *)
(*  Recherche Scientifique. All rights reserved.  This file is         *)
(*  distributed under the terms of the GNU Lesser General Public       *)
(*  License, with the special exception on linking described in file   *)
(*  ../LICENSE.                                                        *)
(*                                                                     *)
(***********************************************************************)

open Format

let time f arg msg =
  let t1 = Unix.gettimeofday () in
  let r = f arg in
  let t2 = Unix.gettimeofday () in
  let time = (t2 -. t1) *. 1000. in
  Logger.msg `STATS "%s: %fms" msg time;
  r


let compose_parallel run auto_list tree nodes () =
  match auto_list with
    [ auto ] -> [run auto tree nodes]
  | _ -> assert false

let compose_sequential run auto_list tree nodes () =
  [ List.fold_left (fun acc auto ->
    run auto tree acc) nodes auto_list ]


let restart_parallel run auto_list tree nodes () =
  match auto_list with
    [ auto ] -> List.map snd (run auto tree nodes)
  | _ -> assert false

let restart_sequential run auto_list tree nodes () =
  List.map (fun auto -> run auto tree nodes) auto_list

let main () =
  let () = Options.parse () in
  let doc =
    let fd, close_fd = match !Options.input_file with
      None | Some "-" | Some "/dev/stdin" -> stdin, ignore
    | Some input ->
        let fd = open_in input in fd, fun () -> close_in fd
    in
    let d = time Naive_tree.load_xml_file fd "parsing xml document" in
    close_fd (); d
  in
  let queries =
    time
      (fun l ->
        List.map (fun q ->
          Xpath.Parser.parse
            (Ulexing.from_utf8_string q)) l)
      !Options.queries
      "parsing XPath queries"
  in
  (* parallel, compose  ->     action
     true, true -> Ata.concat of all automata and single run
     true, false -> Ata.merge of all automata and single run
     false, true -> Eval first, then run on results then ...
     false, false -> Eval first on root, then second on root then ...
  *)
  let auto_list =
    time
      (fun l ->
        List.map (fun query -> Xpath.Compile.path query) l)
      queries
      "compiling XPath queries"
  in
  let auto_list =
    if !Options.parallel then
      match auto_list with
        fst :: rest ->
          let f =
            if !Options.compose then
              Ata.concat
            else
              Ata.merge
          in
          let big_auto = List.fold_left f fst rest in
          [big_auto]
      | _ -> assert false

    else
      auto_list
  in
  let output =
    match !Options.output_file with
    | None | Some "-" | Some "/dev/stdout" -> stdout
    | Some f -> open_out f
  in
  if !Options.stats then begin
    List.iter (fun query ->
      Logger.msg `STATS "Query: %a " Xpath.Ast.print_path query) queries;
    List.iter (fun auto ->
      Logger.msg `STATS "@[Automaton: @\n%a@]" Ata.print auto) auto_list;
  end;

  let module Naive = Run.Make(Naive_tree) in
  let result_list =
    let root = [ Naive_tree.root doc] in
    let f, msg =
      match !Options.parallel, !Options.compose with
        true, true ->
          compose_parallel Naive.eval auto_list doc root, "parallel/compose"
      | true, false ->
          restart_parallel Naive.full_eval auto_list doc root, "parallel/restart"
      | false, true ->
          compose_sequential Naive.eval auto_list doc root , "sequential/compose"
      | false, false ->
          restart_sequential Naive.eval auto_list doc root, "sequential/restart"
    in
    time f () ("evaluating query in " ^ msg ^ " mode")
  in
  time (fun () ->
    let count = ref 1 in
    List.iter (fun results ->
      output_string output "<xml_result num=\"";
      output_string output (string_of_int !count);
      output_string output "\" >\n";
      if !Options.count then begin
        output_string output (string_of_int (List.length results));
        output_char output '\n';
      end else
        List.iter (fun n ->
          Naive_tree.print_xml output doc n;
          output_char output '\n'
        ) results;
      output_string output "</xml_result>\n";
      incr count
    ) result_list;
    flush output;
    if output != stdout then close_out output

  ) () "serializing results"


let () =
  try
    main ()
  with
    Arg.Bad msg -> eprintf "Error: %s\n%!" msg; Options.usage (); exit 1
  | Sys_error msg -> eprintf "Error: %s\n%!" msg; exit 2
  | Tree.Parse_error msg ->
      eprintf "Error: %s, %s\n%!"
        (match !Options.input_file with
          Some s -> ("file " ^ s)
        | None -> "[stdin]") msg; exit 3
  | Xpath.Ulexer.Error (s, e, msg) -> eprintf "Error: character %i-%i: %s\n%!" s e msg; exit 4
  | e -> eprintf "FATAL ERROR: %s\n%!" (Printexc.to_string e); exit 128