#!/usr/bin/env cash
(*
# vim:ft=ocaml 
*)

(* Prepare a LaTeX run for two-way communication with OCaml
   By Geoffrey Alan Washburn <washburn@acm.org> based upon
   code by Scott Pakin <scott+pt@pakin.org>                 *)

(* 
 This is file `ocamltex', Version 0.5
 
 Copyright (C) 2006 by Geoffrey Alan Washburn <washburn@acm.org>
 Copyright (C) 2004 by Scott Pakin <scott+pt@pakin.org>

 This file may be distributed and/or modified under the conditions
 of the LaTeX Project Public License, either version 1.2 of this
 license or (at your option) any later version.  The latest
 version of this license is in:

    http://www.latex-project.org/lppl.txt

 and version 1.2 or later is part of all distributions of LaTeX
 version 1999/12/01 or later. *)

(* Initialize the top-level we are going to interact with *) 
let () = Toploop.initialize_toplevel_env ();;

(* Load the regular expression library *)
#load "str.cma";;
open Cash;;
open Str;;

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

(* Set the default "program name" *)
let progname = ref "ocamltex"

(* Set the default "jobname" *)
let jobname = ref "texput" 

(* Display the "usage" information when run? *)
let showUsage = ref false 

(* Set the default file names *)
let toocaml = ref ((!jobname) ^ ".toml")
let fromocaml = ref ((!jobname) ^ ".frml")
let toflag = ref ((!jobname) ^ ".tfml")
let fromflag = ref ((!jobname) ^ ".ffml")
let doneflag = ref ((!jobname) ^ ".dfml") 
let logfile = ref ((!jobname) ^ ".lgml")

(* Create a string reference used to communicate results from "eval" *)
let evalData = ref ""

(* By default use "latex" *)
let latexprog = try ref (getenv "OCAMLTEX") with _ -> ref "latex" 

(* Separator character for directory names *)
let dirsep = "/"

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

(* Print an error message and exit with an error code *)
let error (msg : string) : 'a = 
  begin
    print_string ((!progname) ^ " (error) : " ^ msg ^ "\n"); 
    exit 1
  end  

(* basename *)
let basename (fn : string) =
  let parts = split (regexp dirsep) fn in
  let rec last (l : string list) : string = 
    (match l with 
     | []     -> (error "Argument to basename doesn't have a basecase?!")
     | [base] -> base
     | h::t   -> last t) in
   last (parts)

(* If a file exists, delete it *)
let delete_file_exists f = if (is_file_existing_fn f) then delete_file f

(* Remove trailing newline from a string *)
let chomp s = 
    if s.[(String.length s) - 1] = '\n' then
      String.sub s 0 ((String.length s) - 1)
    else
      s

(* Evaluate the supplied string as a OCaml top-level input *)
let eval txt = let lb = (Lexing.from_string txt) in
 let phr = !Toploop.parse_toplevel_phrase lb in
 Toploop.execute_phrase true Format.str_formatter phr 

(* Remove .tex suffix *)
let remove_tex_suffix file = global_replace (regexp "\\.tex$") "" file

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

(* Get the name of the program that was run. *)
let () = progname := basename (List.hd (command_line ()))

(* Set the name of the files used *)
let namefiles () =
    begin
      toocaml := (!jobname) ^ ".toml";
      fromocaml := (!jobname) ^ ".frml";
      toflag := (!jobname) ^ ".tfml";
      fromflag := (!jobname) ^ ".ffml";
      doneflag := (!jobname) ^ ".dfml";
      logfile := (!jobname) ^ ".lgml";
    end  

(* Create the separator string.

   Can probably use something shorter than 20, because probability of collision 
   with 1/26^20 is really low.*)
let () = Random.self_init ()
let rec genSeparator n = 
    match n with
        | 0 -> ""
        | n -> (Char.escaped (Char.chr ((Char.code 'A') + (Random.int 26)))) ^ genSeparator (n - 1) 
let separator = genSeparator 20


(* Remove temporary files, if they exist *)
let filelist () = [!toocaml; !fromocaml; !toflag; !fromflag; !doneflag]
let cleanup () = ignore (List.map delete_file_exists (filelist ()))
         
(* Prelude information sent to TeX *)
let firstcmd () = 
    "\\makeatletter" ^ 
    "\\def\\mlmac@tag{" ^ separator ^ "}" ^
    "\\def\\mlmac@tofile{" ^ (!toocaml) ^ "}" ^
    "\\def\\mlmac@fromfile{" ^ (!fromocaml) ^ "}" ^
    "\\def\\mlmac@toflag{" ^ (!toflag) ^ "}" ^ 
    "\\def\\mlmac@fromflag{" ^ (!fromflag) ^ "}" ^
    "\\def\\mlmac@doneflag{" ^ (!doneflag) ^ "}" ^ 
    "\\makeatother" ^
    "\\input{" ^ (!jobname) ^ "}"

(* Parses the command line arguments, handling the ones it understands,
   assumes the rest are file names to be processed *)
let rec parseArgs (args : string list) : string list =
  (match args with
   | []               -> []
   | "-h"::t          -> (showUsage := true); parseArgs t
   | "--h"::t         -> (showUsage := true); parseArgs t
   | "--help"::t      -> (showUsage := true); parseArgs t
   | (h::t)            
     when 
     (string_match (regexp "--latex=\\(.*\\)") h 0) -> 
      begin
        latexprog := (matched_group 1 h); 
        parseArgs t
      end
   | h::t            -> h :: (parseArgs t))

exception FileDone

(* Loop until the specified file exists, or the specified process exits *)
let rec awaitexists (file : string) pid =
    if (not (is_file_existing_fn file)) then 
      begin
        sleep 0;
        try (match (wait ~wflags:[WNOHANG] pid) with 
            | WEXITED pid' -> cleanup (); raise FileDone
            | WSIGNALED pid' -> ()
            | WSTOPPED pid' -> ())
        with Child_not_ready -> awaitexists file pid   
      end
    else
      ()

(* Process a message sent from TeX *)
let process_input (logCh : out_channel) (data : string) : string =  
    let chunks = List.map chomp (split (regexp separator) data) in
      match chunks with
          | optarg::tail ->
            (match optarg with
                 | ("USE"|"DEF") ->
                   (match tail with mname::tail ->
                     let mname = global_replace (regexp "^[^A-Za-z]+") "" mname in
                     let mname = global_replace (regexp "\\W") "_" mname in
                     let mname = "latex_" ^ mname in
                     (match optarg with
                          | "USE" -> 
                            let args = List.map (fun arg -> "\"" ^ arg ^ "\"") tail in
                            let args' = String.concat " " args in
                            let cmd = "let _ = evalData := " ^ mname ^ " " ^ args' ^ ";;" in
                            let () = write_string_out ~dst:logCh ("(* Using OCaml code *)\n\n" ^ cmd ^ 
                                                                "\n\n(********************)\n") in
                              ignore (eval cmd) ; !evalData
                          | "DEF" -> 
                            (match tail with
                                 | (args_string::tail') ->
                                   let args = split (regexp ",") args_string in
                                   let args' = List.map (fun arg -> "(" ^ arg ^ " : string)") args in
                                   let args'' = String.concat " " args' in
                                   let body = String.concat "" tail' in
                                   let cmd = "let rec " ^ mname ^ " " ^ args'' ^ " : string = " ^ body ^ ";;" in 
                                   let () = write_string_out ~dst:logCh 
                                            ("(* Defining OCaml function *)\n\n" ^ cmd ^ 
                                             "\n\n(***************************)\n") in
                                   ignore (eval cmd) ; ""
                                 | _ -> error ("Bad input stream:" ^ data))) 
                    | _ -> error ("Bad input stream:" ^ data))
                 | "EXEC" -> 
                    let code = String.concat "" tail in
                    let cmd = code ^ ";;" in
                    let () = write_string_out ~dst:logCh ("(* Execing OCaml code *)\n\n" ^ cmd ^ 
                                                          "\n\n(**********************)\n") in
                       ignore (eval cmd) ; ""
                 | _ -> error ("Bad opcode:" ^ data)) 
          | _ -> error ("Bad input stream:" ^ data)

(* Loop waiting for definition/evaluation requests *)
let rec main_loop (logCh : out_channel) pid =
    let () = awaitexists (!toflag) pid in
    let toocamlCh = open_input_file (!toocaml) in
    let entirefile = string_of_in_channel toocamlCh in
    let _ = close_in toocamlCh in 
    let result = process_input logCh entirefile in
    let result = result ^ "\\endinput" in
    let () = write_string_out ~dst:logCh 
                              ("(* OCaml evaluation result *)\n\n" ^ result ^ 
                               "\n\n(***************************)\n") in
    let () = delete_file_exists (!fromocaml) in
    let fromocamlCh = open_output_file ~flags:[O_CREAT] (!fromocaml) in
    let () = write_string_out ~dst:fromocamlCh result in
    let _ = close_out fromocamlCh in
    let () = delete_file_exists (!toflag) in
    let () = delete_file_exists (!toocaml) in
    let () = delete_file_exists (!doneflag) in
    let _ = close_out (open_output_file ~flags:[O_CREAT] (!fromflag)) in
    let () = awaitexists (!toocaml) pid in
    let () = delete_file_exists (!fromflag) in
    let _ = close_out (open_output_file ~flags:[O_CREAT] (!doneflag)) in
      main_loop logCh pid

(* Loop over the files supplied on the command-line *)
let rec file_loop files =
    match files with
        | [] -> exit 0
        | h::t -> 
          let () = jobname := remove_tex_suffix h in
          let () = namefiles () in
          (* Clean up temporary files that might be hanging around from a previous run. *)
          let () = cleanup () in
          let pid = (match fork () with
                   | Some(latexpid) ->  latexpid
                   | None -> exec_path (!latexprog) [firstcmd ()]) in
          let logfileCh = open_output_file ~flags:[O_CREAT] (!logfile) in
          let () = (try main_loop logfileCh pid with FileDone -> ()) in
          let _ = close_out logfileCh in
            file_loop t

(* Parse the command line and iterate over the input files *)
let main () =
      let () = print_endline 
               ("This is OCamlTeX, Version 0.5 (Copyright 2006 by Geoffrey Washburn)") in
      let args = parseArgs (List.tl (command_line ())) in
        if !showUsage then
          begin
            print_endline 
"Usage: ocamltex [-h|--h|-help|--help] [--latex=program] files...
Options:
ocamltex accepts the following command-line options

    -h,-help,--h,--help
        Display basic usage information.

    --latex=program
        Specify a program to use instead of latex. For example,
        \"--latex=pdflatex\" would typeset the given document using pdflatex
        instead of ordinary latex.
";
            exit 0
          end    
        else
          file_loop args

(* Set everything in motion *)
let _ = main ()