using System
using System.Text
using System.IO

using Nemerle.Logging
using Nemerle.Utility
using Nemerle.Collections
using Nemerle.Imperative

using Fx7.Parser.SExpr

using SCG = System.Collections.Generic

set namespace Fx7

public partial class Parser
  class EOF : System.Exception {}

  [ExtensionPattern (name ::: body = SExpr.Cons (SExpr.Name (name) :: body))] \
  variant SExpr
    | Name { name : string; }
    | Cons { body : list [SExpr]; }

    public line : int

    public this () { }
    public this (l : int)
      line = l

    public override ToString () : string
      match (this)
        | Name (name) => name
        | Cons (Name (name) :: _) => $ "($name ...)"
        | Cons (Cons (Name (name) :: _) :: _) => $ "(($name ...))"
        | Cons  => "(...)"

    [Nemerle.OverrideObjectEquals] \
    public Equals (other : SExpr) : bool
      match ((this, other))
        | (Name (n1), Name (n2)) => n1 == n2
        | (Cons (l1), Cons (l2)) =>
          l1.Length == l2.Length && List.ForAll2 (l1, l2, _.Equals (_))
        | _ => false

  static l (name : string, a1 : SExpr) : SExpr
    Cons (a1.line, [Name (a1.line, name), a1])

  static l (name : string, a1 : SExpr, a2 : SExpr) : SExpr
    Cons (a1.line, [Name (a1.line, name), a1, a2])

  static l (name : string, a1 : SExpr, a2 : SExpr, a3 : SExpr) : SExpr
    Cons (a1.line, [Name (a1.line, name), a1, a2, a3])

  static l (name : string, ch : list [SExpr]) : SExpr
    match (ch)
      | a1 :: _ => Cons (a1.line, Name (a1.line, name) :: ch)
      | [] => Cons ([Name (name)])

  static ParseAnnotations (sexprs : list [SExpr]) : list [string * option [SExpr]]
    def isann (s)
      s.StartsWith (":")
      
    def parse (acc, l : list [SExpr])
      match (l)
        | [] => acc.Rev ()
        | Name (name1) :: ((Name (name2) :: _) as xs) when isann (name2) && isann (name1) \
        | [Name (name1)] when isann (name1) with xs = [] =>
          parse ((name1, None ()) :: acc, xs)
        | Name (name1) :: arg :: xs when isann (name1) =>
          parse ((name1, Some (arg)) :: acc, xs)
        | x :: _ =>
          throw FatalError ($ "bad, bad annotation! $x")

    parse ([], sexprs)

  PrettyPrint (th : SExpr) : string
    def sb = StringBuilder ()
    PrettyPrint (sb, th)
    sb.ToString ()
    
  PrettyPrint (sb : StringBuilder, th : SExpr) : void
    def f (t)
      PrettyPrint (sb, t)
    def add (s)
      _ = sb.Append (s)
    match (th)
      | "explies" ::: [q, p] with sym = "⇒" \
      | "implies" ::: [p, q] with sym = "⇒" \
      | "iff" ::: [p, q] with sym = "⇔" \
        =>
        add ("(")
        f (p)
        add (" "); add (sym); add (" ")
        f (q)
        add (")")

      | "!=" ::: [p, q] with sym = "≠" \
      | "=" ::: [p, q] with sym = "=" \
      | "<" ::: [p, q] with sym = "<" \
      | ">" ::: [p, q] with sym = ">" \
      | ">=" ::: [p, q] with sym = "≥" \
      | "<=" ::: [p, q] with sym = "≤" \
      | "-" ::: [p, q] with sym = "-" \
      | "+" ::: [p, q] with sym = "+" \
        =>
        f (p)
        add (" "); add (sym); add (" ")
        f (q)

      | "xor" ::: ch with sym = "⊕" \
      | "and" ::: ch with sym = "∧" \
      | "or" ::: ch with sym = "∨" \
      | "+" ::: ch with sym = "+" \
        =>
        add ("(")
        f (ch.Head)
        foreach (c in ch.Tail)
          add (" " + sym + " ")
          f (c)
        add (")")

      | "select" ::: [p, q] =>
        f (p)
        add ("[")
        f (q)
        add ("]")

      | "store" ::: [p, q, r] =>
        add ("(")
        f (p)
        add ("["); f (q); add ("] ← ")
        f (r)
        add (")")

      | "not" ::: [p] with sym = "¬" \
        =>
        add (sym)
        f (p)
      
      | "distinct" ::: lst =>
        add (lst.ToString (" ≠ "))
        
      | "if_then_else" ::: [c, t, e] =>
        add ("if "); f (c); add (" then ")
        f (t); add (" else "); f (e)

      | "exists" ::: [v, t] with (p = null, sym = "∃") \
      | "exists" ::: [v, p, t] with sym = "∃" \
      | "forall" ::: [v, t] with (p = null, sym = "∀") \
      | "forall" ::: [v, p, t] with sym = "∀" =>
        add ("(" + sym)
        f (v)
        when (p != null)
          f (p)
        add (" • ")
        f (t)
        add (")")

      | Name (name) =>
        add (name)

      | name ::: [] =>
        add ("("); add (name); add (")")
        
      | name ::: children =>
        add ("("); add (name); add (" ")
        children.Iter (f)
        add (")")

      | s => throw FatalError ($ "evil sexpression $s")

  PlainSExprToTerm (th : SExpr, core : Core) : Term
    match (th)
      | Name (name) with children = [] \
      | name ::: children =>
        core.TermPool.Get (name, children.Map (PlainSExprToTerm (_, core)))

      | s => throw FatalError ($ "bad plain term s-expression: $s")

  SExprToTerm (th : SExpr, ctx : Map [string, Term], core : Core) : Term
    def f (t)
      SExprToTerm (t, ctx, core)
    match (th)
      | Name (name) when ctx.Contains (name) =>
        ctx.Get (name)

      | pred ::: args when renamings.Contains (pred) =>
        f (l (renamings [pred], args))

      | "flet" ::: (name ::: [val]) :: term :: _ \
      | "let"  ::: (name ::: [val]) :: term :: _ =>
        SExprToTerm (term, ctx.Add (name, SExprToTerm (val, ctx, core)), core)

      // absent in the standard
      | "explies" ::: [p, q] =>
        f (l ("implies", q, p))

      | "xor" ::: [p, q] =>
        f (l ("not", l ("iff", p, q)))

      /*
      | "iff" ::: [p, q] =>
        f (l ("or",  
          l ("and", p, q),
          l ("and", l ("not", p), l ("not", q))))

      | "implies" ::: [p, q] =>
        f (l ("or", l ("not", p), q))
      */

      | "if_then_else" ::: [c, t, e] =>
        f (l ("and",  
          l ("implies", l ("not", c), e),
          l ("implies", c, t)))

      | "!=" ::: [p, q] =>
        f (l ("not", l ("=", p, q)))

      | "=" ::: [Name ("$@true"), Name ("$@true")] =>
        f (Name ("true"))

      // formulas for Simplify sometimes contain stuff like
      // (EQ |@true| some_formula)
      // where some_formula is generated from DEFPRED
      | "=" ::: [p, Name ("$@true")] \
      | "=" ::: [Name ("$@true"), p] => 
        f (p)

      | "~" ::: [Name ("1")] =>
        f (Name ("-1"))
      
      | "LBL" ::: [_, x] when simplify_syntax \
      | "LBLPOS" ::: [_, x] when simplify_syntax \
      | "LBLNEG" ::: [_, x] when simplify_syntax =>
        f (x)

      | "ORDER" ::: [Name (lt), Name (le)] when simplify_syntax =>
        f (OrderAxiom (lt, le))

      | "forall" ::: [Cons (vars1), "forall" ::: [Cons (vars2), form]] =>
        f (l ("forall", Cons (vars1 + vars2), form))

      | qhead ::: (head ::: tail) :: rest when (qhead == "forall" || qhead == "exists") && head != "$@vars" =>
        f (Cons (Name (qhead) :: Cons (Name ("$@vars") :: Name (head) :: tail) :: rest))

      | Name (name) with children = [] \
      | name ::: children =>
        core.TermPool.Get (name, children.Map (f))

      | s => throw FatalError ($ "bad term s-expression: $s")

  SExprToCommand (th : SExpr) : Command
    def parse (f)
      core =>
        foreach ((from, to) in current_rewrite_rules)
          core.TermPool.AddRewriteRule (PlainSExprToTerm (from, core), 
                                        PlainSExprToTerm (to, core))

        SExprToTerm (f, Map (), core)
    match (th)
      | "benchmark" ::: Name (name) :: anns =>
        def anns = ParseAnnotations (anns)
        def args = anns.FoldLeft (current_axioms, fun ((name, val), acc) {
          match (val) {
            | Some (f) when name == ":formula" || name == ":assumption" =>
              f :: acc
            | _ => acc
          }
        })
        match (args)
          | [f] \
          | _ :: _ with f = l ("and", args) =>
            def cmd =
              match (List.Assoc (anns, ":status"))
                | Some (Some (Name ("sat"))) => 
                  Command.ExpectSAT (parse (f))
                | Some (Some (Name ("unsat"))) => 
                  Command.ExpectUnSAT (parse (f))
                | _ =>
                  Command.CheckSAT (parse (f))
            cmd.name = name
            cmd
          | _ => throw FatalError ("no :formula, no donut")

      | s => throw FatalError ($ "bad command s-expression: $s")

  mutable reader : TextReader
  mutable putback : int = -1;
  mutable current_line : int = 1

  DoRead () : char
    if (putback > 0)
      def ch = putback :> char
      putback = -1
      ch
    else
      def ch = reader.Read ()
      when (ch == '\n' :> int) current_line++
      if (ch < 0) throw EOF ()
      else
        //System.Console.Write (ch :> char)
        ch :> char
          
  Read () : char
    def ch = DoRead ()
    if (ch != ';') ch
    else
      while (DoRead () != '\n') {}
      '\n'

  ReadNonWS () : char
    def ch = Read ()
    if (Char.IsWhiteSpace (ch)) ReadNonWS ()
    else ch
 
  SimplifyFixup (_ : string) : string
    | s when !simplify_syntax => s
    | "AND" => "and"
    | "OR" => "or"
    | "NOT" => "not"
    | "IFF" => "iff"
    | "IMPLIES" => "implies"
    | "EXPLIES" => "explies"
    | "XOR" => "xor"
    | "DISTINCT" => "distinct"
    | "TRUE" => "true"
    | "FALSE" => "false"
    | "EQ" => "="
    | "NEQ" => "!="
    | "FORALL" => "forall"
    | "EXISTS" => "exists"
    | "|@true|" => "$@true"
    | "|@false|" => "$@false"
    | s => s
  
  SmtFixup (se : SExpr) : SExpr
    def is_smt_quant (_ : SExpr)
      | name ::: [Name (other)] => name.StartsWith ("?") && !other.StartsWith ("?")
      | _ => false
    def smt_quant_var (_ : SExpr)
      | name ::: [_] => SExpr.Name (name)
      | _ => assert (false)

    match (se)
      | "exists" ::: (a :: _) as args when is_smt_quant (a) with head = "exists" \
      | "forall" ::: (a :: _) as args when is_smt_quant (a) with head = "forall" =>
        def loop (_)
          | (acc, (Name (n) :: _) as rest) when n.StartsWith (":") =>
            (acc.Rev (), rest)
          | (acc, x :: xs) => loop (x :: acc, xs)
          | (acc, []) => (acc.Rev (), [])
        def (args, annot) = loop ([], args)
        def annot = ParseAnnotations (annot)
        def is_pat = annot.Exists ((x, _) => x == ":pat")
        def is_nopat = annot.Exists ((x, _) => x == ":nopat")
        if (is_pat && is_nopat)
          throw FatalError ("both :pat and :nopat given")
        else
          def parse (e)
            def e =
              if (e.StartsWith ("{") && e.EndsWith ("}"))
                e.Substring (1, e.Length - 2)
              else e
            try
              match (ParseSExprs (e))
                | [t] => t
                | lst =>
                  when (is_nopat) throw FatalError (":nopat expects a single term")
                  l ("MPAT", lst)
            catch
              | ex => throw FatalError ($"during parsing pattern '$e', got $(ex.Message)")
          def pats =
            annot.FoldLeft ([], _ => {
              | ((":nopat", Some (Name (e))), acc)
              | ((":pat", Some (Name (e))), acc) => parse (e) :: acc
              | ((":pat", _), acc) | ((":nopat", _), acc) =>
                if (Driver.IgnoreTriggers) acc
                else throw FatalError (":pat/:nopat requires an argument")
              | (_, acc) => acc
            })
          def (args, form) = args.DivideLast ()
          def vars = Cons (args.Map (smt_quant_var))
          if (pats is [])
            l (head, vars, form)
          else
            l (head, vars, l (if (is_pat) "PATS" else "NOPATS", pats), form)
      | _ => se


  GetSExpr (ch : char) : SExpr
    | '(' =>
      def loop (acc)
        def ch = ReadNonWS ()
        if (ch == ')') 
          if (simplify_syntax) 
            Cons (current_line, acc.Rev ())
          else SmtFixup (Cons (acc.Rev ()))
        else loop (GetSExpr (ch) :: acc)
      def res = loop ([])
      res
    | '|' when simplify_syntax with stop = '|' \
    | '{' with stop = '}' \
    | '"' with stop = '"' =>
      def sb = StringBuilder ()
      _ = sb.Append (ch)
      def loop ()
        def ch = Read ()
        _ = sb.Append (ch)
        if (ch == stop) Name (current_line, SimplifyFixup (sb.ToString ()))
        else loop ()
      loop ()
    | _ =>
      def sb = StringBuilder ()
      def loop (ch)
        if (ch == ')')
          putback = ch :> int
          Name (current_line, SimplifyFixup (sb.ToString ()))
        else if (Char.IsWhiteSpace (ch))
          Name (current_line, SimplifyFixup (sb.ToString ()))
        else 
          _ = sb.Append (ch)
          loop (Read ())
      loop (ch)

  
  ReadSExprs () : list [SExpr]
    def loop (acc)
      mutable first_char_read = false
      try
        def ch = ReadNonWS ()
        first_char_read = true
        def sexpr = GetSExpr (ch)
        when (ch == ')')
          throw FatalError ("unexpected ')'")
        loop (sexpr :: acc)
      catch
        | _ is EOF =>
          if (first_char_read)
            throw FatalError ("unexpected end of file")
          else acc.Rev ()
    loop ([])
 
  mutable simplify_syntax : bool
  mutable simpl_name_prefix = "simpl"

  [Accessor] plain_predicates : Hashtable [string, string] = Hashtable ()

  AddPredicate (name : SExpr) : void
    match (name)
      | name ::: _ \
      | Name (name) => plain_predicates [name] = name
      | _ => {}

  // modified by ReadAxioms
  mutable current_axioms : list [SExpr] = []
  mutable current_rewrite_rules : list [SExpr * SExpr] = []
  renamings : Hashtable [string, string] = Hashtable ()
  kill_list : Hashtable [string, list [SExpr]] = Hashtable ()
  mutable filename : string

  public this (filename : string)
    this.filename = filename
    def file = FileStream (filename, FileMode.Open, FileAccess.Read)
    reader = StreamReader (file, Encoding.UTF8)
    simplify_syntax = Driver.SimplifySyntax

  public this ()
    reader = System.Console.In
    simplify_syntax = Driver.SimplifySyntax


  OrderAxiom (lt : string, le : string) : SExpr
    plain_predicates [lt] = lt
    plain_predicates [le] = le
    def x = Name ("x")
    def y = Name ("y")
    def z = Name ("z")

    l ("and", [
      l ("forall", l ("x",[]), l ("not", l (lt, x, x))),
      l ("forall", l ("x",y,z), 
        l ("PATS", l ("MPAT", l (lt, x, y), l (lt, y, z))),
        l ("implies", l ("and", l (lt, x, y), l (lt, y, z)), l (lt, x, z))),
      l ("forall", l ("x",y), l ("iff", l ("or", l (lt, x, y), l ("=", x, y)), l (le, x, y)))
    ])

  DefPredAxiom (name : SExpr, val : SExpr) : SExpr
    match (name)
      | _ ::: Name (arg) :: args =>
        l ("forall", l (arg, args), l ("iff", name, val))
      | _ => throw FatalError ($"bad DEFPRED: $name")


  DefPredMapAxiom (name : SExpr, i : SExpr, val : SExpr) : SExpr
    match (name)
      | _ ::: Name (arg) :: args =>
        l ("forall", l (arg, args + [i]), l ("iff", l ("select", [name, i]), val))
      | _ => throw FatalError ($"bad DEFPREDMAP: $name")


  public ReadAxioms (filename : string) : void
    def backup = reader
    def backup_line = current_line
    try
      def file = FileStream (filename, FileMode.Open, FileAccess.Read)
      reader = StreamReader (file, Encoding.UTF8)
      def sexprs = ReadSExprs ()
      foreach (sexpr in sexprs)
        match (sexpr)
          | "DEFPRED" ::: [name] =>
            AddPredicate (name)
          | "DEFPRED" ::: [name, value] =>
            AddPredicate (name)
            current_axioms ::= DefPredAxiom (name, value)
          | "DEFPREDMAP" ::: [_, _] => {}
          | "DEFPREDMAP" ::: [pred, ind, form] =>
            current_axioms ::= DefPredMapAxiom (pred, ind, form)
          | "ORDER" ::: [Name (lt), Name (le)] =>
            current_axioms ::= OrderAxiom (lt, le)
          | "REWRITE" ::: [from, to] =>
            current_rewrite_rules ::= (from, to)
          | "RENAME" ::: [Name (from), Name (to)] =>
            renamings [from] = to

          // only for -smt converter
          | "FUN" ::: args =>
            dumper.AddFun (args)
          | "VFUN" ::: args =>
            dumper.AddFun (args, var_args = true)
          | "CAST" ::: args =>
            dumper.AddFun (args, is_cast = true)
          | "KILL" ::: [(name ::: _) as arg] =>
            if (kill_list.Contains (name))
              kill_list [name] ::= arg
            else
              kill_list [name] = [arg]

          // normal again
          | "BG_PUSH" ::: [formula] \
          | formula =>
            current_axioms ::= formula
    finally
      current_line = backup_line
      reader = backup

  ParseSExprs (sexpr : string) : list [SExpr]
    def backup = reader
    def backup_line = current_line
    try
      reader = StringReader (sexpr + " ")
      ReadSExprs ()
    finally
      current_line = backup_line
      reader = backup


  dumper : SmtDumper = SmtDumper (this)
  SmtDump (_core : Core, axioms : list [SExpr], formula : SExpr) : void
    dumper.Run (axioms, formula)

  SimplifyIface () : SCG.IEnumerable [Command]
    mutable count = 1
    mutable expect_next_invalid = false
    mutable axioms = current_axioms

    def parse (formula : SExpr, core)
      def formula =
        match (formula)
          | "forall" ::: [_, f] \
          | "forall" ::: [_, _, f] \
          | f => f
      def formula' = Cons (Name ("and") :: axioms + [l ("not", formula)])
      when (Driver.PrettyPrint)
        System.Console.WriteLine (PrettyPrint (formula'))
        System.Environment.Exit (0)

      foreach ((from, to) in current_rewrite_rules)
        core.TermPool.AddRewriteRule (PlainSExprToTerm (from, core), 
                                      PlainSExprToTerm (to, core))

      if (Driver.Translate)
        SmtDump (core, axioms, l ("not", formula))
        null
      else
        SExprToTerm (formula', Map (), core)

    def simpl_parse (sexpr : SExpr)
      match (sexpr)
        | "BG_PUSH" ::: [pred] =>
          axioms ::= pred
          None ()
          
        | "BG_PUSH" ::: pred =>
          axioms ::= l ("and", pred)
          None ()

        | "NAME_PREFIX" ::: [Name (n)] =>
          simpl_name_prefix = n
          None ()
          
        | "NEXT_INVALID" ::: [] =>
          expect_next_invalid = true
          None ()
          
        | "PROMPT_OFF" ::: [] \
        | "PROMPT_ON" ::: [] =>
          None ()
          
        | "BG_POP" ::: [] =>
          axioms = axioms.Tail
          None ()
          
        | "DEFPRED" ::: [name] =>
          AddPredicate (name)
          None ()
          
        | "DEFPRED" ::: [name, value] =>
          AddPredicate (name)
          axioms ::= DefPredAxiom (name, value)
          None ()

        | "DEFPREDMAP" ::: [_, _] =>
          None ()

        | "DEFPREDMAP" ::: [pred, ind, form] =>
          axioms ::= DefPredMapAxiom (pred, ind, form)
          None ()
          
        | "PING" ::: [] =>
          log (ANSWER, "PONG")
          None ()

        | "PRAGMA" ::: args =>
          Some (Command.Pragma (args.Map (_.ToString ())))
          
        | formula =>
          def parse = parse (formula, _)
          def cmd =
            if (expect_next_invalid)
              expect_next_invalid = false
              Command.ExpectSAT (parse)
            else if (Driver.ExpectAllValid)
              Command.ExpectUnSAT (parse)
            else if (Driver.ExpectAllInvalid)
              Command.ExpectSAT (parse)
            else
              Command.CheckSAT (parse)
          cmd.name = $"$simpl_name_prefix.$count"
          count++
          Some (cmd)

    def read_sexpr ()
      try
        def ch = ReadNonWS ()
        def sexpr = GetSExpr (ch)
        when (ch == ')')
          throw FatalError ("unexpected ')'")
        sexpr
      catch
        | _ is EOF => null

    while (true)
      def sx = read_sexpr ()
      if (sx == null) break
      else
        match (simpl_parse (sx))
          | None => 
            {}
            // log (INFO, "OK")
          | Some (cmd) => yield cmd
          

  public Run () : SCG.IEnumerable [Command]
    if (simplify_syntax)
      SimplifyIface ()
    else
      match (ReadSExprs ())
        | [sexpr] =>
          match (sexpr)
            | "list" ::: xs =>
              xs.Map (SExprToCommand (_))
            | "benchmark" ::: _ =>
              [SExprToCommand (sexpr)]
            | x => [Command.CheckSAT (SExprToTerm (x, Map (), _))]
        | _ =>
          throw FatalError ("expected only one sexpression in non-SIMPLIFY mode")