(*
DISTRIBUTION STATEMENT A. Approved for public release. Distribution is unlimited.

This material is based upon work supported by the Under Secretary of Defense for
Research and Engineering under Air Force Contract No. FA8702-15-D-0001. Any
opinions, findings, conclusions or recommendations expressed in this material
are those of the author(s) and do not necessarily reflect the views of the Under
Secretary of Defense for Research and Engineering.

© 2022 Massachusetts Institute of Technology.

The software/firmware is provided to you on an As-Is basis

Delivered to the U.S. Government with Unlimited Rights, as defined in DFARS Part
252.227-7013 or 7014 (Feb 2014). Notwithstanding any copyright notice, U.S.
Government rights in this work are defined by DFARS 252.227-7013 or DFARS
252.227-7014 as detailed above. Use of this work other than as specifically
authorized by the U.S. Government may violate any copyrights that exist in this
work.
*)

module Commands

  open System.IO
  open Common

  // Stuff for processing commands goes here
  let tcam_text() = 
    eprintfn "\t\tAllowed params for tcam:"
    eprintfn "\t\t\tNo parameters specified: use defaults."
    eprintfn "\t\t\tunittcam: Run small TCAM query."
    eprintfn "\t\t\tfdd <dictionary> <stride>: build an FDD from TCAM using file <dictionary> with <stride>."
    eprintfn "\t\t\tsimulate <dictionary> <stride> <input>: build an FDD from TCAM using file <dictionary> with <stride> and simulate on <input>."
    eprintfn "\t\t\ttcam <dictionary> <stride>: build a TCAM using file <dictionary> with <stride>."
  let ddfa_text() = 
    eprintfn "\t\tAllowed params for ddfa:"
    eprintfn "\t\t\tNo parameters specified: use defaults."
    eprintfn "\t\t\tdict <dictionary> <filter> <stride>: build a dictionary from file <dictionary> with <filter> and <stride>."
    eprintfn "\t\t\tsearch <dictionary> <target>: build a dictionary from file <dictionary> and search for entries from file <target>."
    eprintfn "\t\t\tsearch <dictionary> <stride> <target>: build a dictionary from file <dictionary> with <filter> and search for entries from file <target>."
  let aho_text() = 
    eprintfn "\t\tAllowed params for aho:"
    eprintfn "\t\t\tNo parameters specified: use defaults."
    eprintfn "\t\t\tdict <dictionary> <filter>: build a dictionary from file <dictionary> with <filter>."
    eprintfn "\t\t\tsearch <dictionary> <target>: build a dictionary from file <dictionary> and search for entries from file <target>."
  let fdd_text() = 
    eprintfn "\t\tAllowed params for fdd:"
    eprintfn "\t\t\tNo parameters specified: use defaults."
    eprintfn "\t\t\tbitset: test bitset"
    eprintfn "\t\t\tbyterule: test byterule"
    eprintfn "\t\t\tfddproj: test FDD projections"
    eprintfn "\t\t\tfddbuild: build an FDD"  

  [<Literal>]
  let PRINT_ID = "print";
  [<Literal>]
  let GEN_ID = "generate";
  [<Literal>]
  let REGRESSION_ID = "test"
    
  // Runs intermediate stage verifications. 
  module Intermediate = begin
    let desc() =
        eprintfn "\t%s tcam|ddfa|ago|fdd [params]" PRINT_ID
        eprintfn "\t\tPrint out intermediate stage verifications to StdOut for specified stage."
        tcam_text()
        ddfa_text()
        aho_text()
        fdd_text()
        eprintfn "\t\tNote: this output can differ from the file output generated by the \"%s\" command" GEN_ID
    let handle(which, other_params) = 
      eprintfn "==== Running intermediate stage verification for %s with params: (%s)" which (other_params |> String.concat ",")
      let res = match which.ToLower() with
                 | "tcam" -> TCAM.Regression.run_stage(true, None, other_params)
                 | "ddfa" -> DDFA.Regression.run_stage(true, None, other_params)
                 | "aho" -> Aho_Corasick.Regression.run_stage(true, None, other_params)
                 | "fdd" -> FDD.Regression.run_stage(true, None, other_params)
                 | _ -> eprintfn "\tUnknown stage specified: %s" (which.ToLower()); desc(); false
      match res with 
      | true -> eprintfn "%s printing complete" (which.ToLower())
      | false -> eprintfn "%s printing failed" (which.ToLower())
      res
  end

  //  Generate intermediate stage verifications and capture to file
  module Generate = begin
    let desc() =
        eprintfn "\t%s tcam|ddfa|aho|fdd <fn> [params]" GEN_ID
        tcam_text()
        ddfa_text()
        aho_text()
        fdd_text()
        eprintfn "\t\tGenerate reference data for specified stage using [params] and write to <fn>."
        eprintfn "\t\tNote: this output can differ from the StdOut output generated by the \"%s\" command" PRINT_ID
    let handle(which, fn, other_params) = 
      eprintfn "==== Generating reference data for %s with params: (%s, %s)" which fn (other_params |> String.concat ",")
      try 
        let x = new StreamWriter(fn)
        let res = 
          match which.ToLower() with
          | "tcam" -> TCAM.Regression.run_stage(false, Some(x), other_params)
          | "ddfa" -> DDFA.Regression.run_stage(false, Some(x), other_params)
          | "aho" -> Aho_Corasick.Regression.run_stage(false, Some(x), other_params)
          | "fdd" -> FDD.Regression.run_stage(false, Some(x), other_params)
          | _ -> eprintfn "\tUnknown stage specified: %s" (which.ToLower()); desc(); false
        x.Flush()
        x.Close()
        match res with 
          | true -> eprintfn "%s generation complete" (which.ToLower())
          | false -> eprintfn "%s generation failed" (which.ToLower())
        res
      with
         | _ -> eprintfn "%s generation failed on error" (which.ToLower()); false
  end

  // Runs individual regression tests from the sub-projects.
  module Regression = begin
    let desc() =    
      eprintfn "\t%s tcam|ddfa|aho|fdd <ref_fn> [params]" REGRESSION_ID
      eprintfn "\t\tRun regression test for specified stage and compare results to <ref_fn>."
      tcam_text()
      ddfa_text()
      aho_text()
      fdd_text()
    let handle(which, ref_fn, other_params) =
      eprintfn "==== Running regression test for %s with params: (%s %s)" which ref_fn (other_params |> String.concat ",")
      // TODO on fail, print out the other_params
      // If fails with line -1, means something went wrong before the test started
      // generate a tmp file with current results that we will compare to the reference
      let tmp = Path.GetTempFileName()
      try
        let x = new StreamWriter(tmp)
        let res = 
          match which.ToLower() with
          | "tcam" -> TCAM.Regression.run_stage(false, Some(x), other_params)
          | "ddfa" -> DDFA.Regression.run_stage(false, Some(x), other_params)
          | "aho" -> Aho_Corasick.Regression.run_stage(false, Some(x), other_params)
          | "fdd" -> FDD.Regression.run_stage(false, Some(x), other_params)
          | _ -> eprintfn "\tUnknown stage specified: %s" (which.ToLower()); desc(); false
        //let res = run_stage(false, Some(x), other_params)
        x.Flush()
        x.Close()
        let comp_res = 
          match res with 
          | true -> compare_result_files(tmp, ref_fn)
          | false -> (false,0) // something went wrong with generating temp file
        match comp_res with 
        | true,_ -> eprintfn "==== %s test passed" which
        | false,x -> eprintfn "==== %s test failed at line %i" which (x-1)
      with
        | _ -> eprintfn "%s test failed" (which.ToLower()); ()
    end


  module GenTCAMTables = begin
    [<Literal>]
    let GENTABLES = "gen_tcam_tables"
    let desc() =
      eprintfn "\t%s <stride> <dict> [-f filename]" GENTABLES
      eprintfn "\t\tGenerate TCAM tables from a provided dictionary and stride length"
      eprintfn "\t\tDictionaries are a path to a text file containing one word per line."
      eprintfn "\t\tIf the -f flag is provided, write to file instead of StdOut"

    let gen_cam (dictionary : string[]) (filter) (stride) (fd: Option<StreamWriter>) = 
      let example =  Aho_Corasick.AC_Table<_>.Create (dictionary |> Seq.filter (fun i -> i.Length > filter) |> Seq.map (fun i -> (i,i))) |> DDFA.AhoCorasick.toDDFA
      let replacements = DDFA.Utilities.Optimzer.State_Expander.DeferChainReplacement (example)
    
      let ((startId,startMask),tcam) = TCAM.Utilities.TCAMConvert.toTCAM (example, stride)
    
      let tst = tcam.Rows |> Seq.head
      let stateBanks = tst.bank.Length - stride
      let headers = [|
        for idx = 0 to stateBanks - 1 do yield sprintf "state %i" idx
        for idx = 0 to stride - 1 do yield sprintf "Data %i" idx
        yield "match"
        yield "advance"
        for idx = 0 to stateBanks - 1 do yield sprintf "nextState %i" idx
      |]
      let msg = sprintf "%s" (headers |> String.concat ",")

      let out : TextWriter = 
        match fd with
        | Some(a) -> a 
        | None -> System.Console.Out
        
      do out.WriteLine(msg)

      for r in tcam.Rows do
        let s = seq{
          for b in r.bank do yield b.ToString()
          let (m,a,n) = r.decision
          yield match m with Some x -> x | None -> ""
          yield a.ToString()
          for b in n do yield b.ToString()
        } 
        out.WriteLine(s |> String.concat ",")
        //let l2 = (sprintf "Start ID : [%s], StartMask : [%s]" (startId |> Seq.map string |> String.concat ",") (startMask |> Seq.map string |> String.concat ","))
        //match fd with
        //  | Some(a) -> a.WriteLine(l1 + "\n" + l2)
        //  | None -> printfn "%s" (l1 + "\n" + l2)
      printfn "Start ID : [%s], StartMask : [%s]" (startId |> Seq.map string |> String.concat ",") (startMask |> Seq.map string |> String.concat ",")

    let handle(stride : int, dict_file : string, output_file : Option<string> ) =
      let input_path =
        if Path.IsPathRooted(dict_file) then
          dict_file
        else
          Path.Combine(__SOURCE_DIRECTORY__, dict_file)
      let dict = System.IO.File.ReadAllLines(input_path)
      
      let output_path =
        match output_file with
          | Some(a) -> Some(new StreamWriter(if Path.IsPathRooted(a) then a else Path.Combine(__SOURCE_DIRECTORY__, a)))
          | None -> None
      gen_cam dict 0 stride output_path|> ignore
      ()
  end 

  module Help = begin
    [<Literal>]
    let HELP = "help"
    let desc() =
      eprintfn "Usage: Stardust-Chicken <command>"
      eprintfn "Allowed Commands:"
      // list all the possible commands here
      eprintfn "\t%s" HELP
      eprintfn "\t\tPrints this."
      GenTCAMTables.desc()
      Intermediate.desc()
      Generate.desc()
      Regression.desc()
      //TestAll.desc()
  end

  // This parses a single command and evaluates it
  let commandParse args =
    match args |> Array.toList with
      [] -> 
        Help.desc()
        0
      | Help.HELP::_ -> 
        Help.desc()
        0
      | PRINT_ID::which::tail ->
        Intermediate.handle(which, tail) |> ignore
        0
      | GEN_ID::which::fn::tail ->
        Generate.handle(which, fn, tail) |> ignore
        0
      | REGRESSION_ID::which::fn::tail ->
        Regression.handle(which, fn, tail)
        0
      | GenTCAMTables.GENTABLES::stride::dict::tail ->
        let flag = 
          match (tail |> String.concat " ") with
          | "" -> None
          | Prefix "-f " filename ->  Some(filename)
          | other -> eprintfn "Option: %s\nnot recognized.\n" other; exit(-1)
        GenTCAMTables.handle(stride |> int, dict, flag) |> ignore
        0
      // For all unrecognized commands
      | x -> 
        // Print some help to the command line
        eprintfn "Command: %s\nnot recognized.\n" (x |> Seq.map (fun i -> "\"" + i + "\"") |> String.concat " ")
        Help.desc()
        0


