// Author: Zhuo Ying Jiang Li

// Starting code by Dr Christian Urban

// parser: convert sequence of tokens to AST

//

// Use this command to print parsed AST:

// amm fun_parser.sc <name>.fun

//

import $file.fun_tokens, fun_tokens._

// more convenience for the map parsers later on;

// it allows writing nested patterns as

// case x ~ y ~ z => ...

case class ~[+A, +B](x: A, y: B)

// parser combinators

abstract class Parser[I, T](using is: I => Seq[_])  {

  def parse(in: I): Set[(T, I)]  

  def parse_all(in: I) : Set[T] =

    for ((hd, tl) <- parse(in); 

        if is(tl).isEmpty) yield hd

}

// alternative parser

class AltParser[I, T](p: => Parser[I, T], 

                      q: => Parser[I, T])(using I => Seq[_]) extends Parser[I, T] {

  def parse(in: I) = p.parse(in) ++ q.parse(in)   

}

// sequence parser

class SeqParser[I, T, S](p: => Parser[I, T], 

                         q: => Parser[I, S])(using I => Seq[_]) extends Parser[I, ~[T, S]] {

  def parse(in: I) = 

    for ((hd1, tl1) <- p.parse(in); 

         (hd2, tl2) <- q.parse(tl1)) yield (new ~(hd1, hd2), tl2)

}

// map parser

class MapParser[I, T, S](p: => Parser[I, T], 

                         f: T => S)(using I => Seq[_]) extends Parser[I, S] {

  def parse(in: I) = for ((hd, tl) <- p.parse(in)) yield (f(hd), tl)

}

// more convenient syntax for parser combinators

extension [I, T](p: Parser[I, T])(using I => Seq[_]) {

  def ||(q : => Parser[I, T]) = new AltParser[I, T](p, q)

  def ~[S] (q : => Parser[I, S]) = new SeqParser[I, T, S](p, q)

  def map[S](f: => T => S) = new MapParser[I, T, S](p, f)

}

// -------------------------------------------------

// atomic parsers

// atomic parser for types

case class TypeParser(ty: Set[String]) extends Parser[Tokens, String] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "type" && ty.contains(tk._2) => Set((tk._2, tkns))

    case _ => Set()

  }

}

// atomic parser for global ids

case object GlobalIdParser extends Parser[Tokens, String] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "global" => Set((tk._2, tkns))

    case _ => Set()

  }

}

// atomic parser for ids

case object IdParser extends Parser[Tokens, String] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "id" => Set((tk._2, tkns))

    case _ => Set()

  }

}

// atomic parser for doubles (I use Float because that's what is used in the AST structures given in CW5)

case object DoubleParser extends Parser[Tokens, Float] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "double" => Set((tk._2.toFloat, tkns))

    case _ => Set()

  }

}

// atomic parser for integers

case object IntParser extends Parser[Tokens, Int] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "int" => Set((tk._2.toInt, tkns))

    case _ => Set()

  }

}

// atomic parser for operators

case class OpParser(ops: Set[String]) extends Parser[Tokens, String] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "op" && ops.contains(tk._2) => Set((tk._2, tkns))

    case _ => Set()

  }

}

// atomic parser for character

case object CharParser extends Parser[Tokens, Char] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._1 == "ch" => {

      val stripped = tk._2.slice(1, tk._2.length-1)  // strip off single quotes

      stripped match {

        case "\\n" => Set(('\n', tkns))

        case "\\t" => Set(('\t', tkns))

        case "\\r" => Set(('\r', tkns))

        case c => Set((c(0), tkns))

      }

    }

    case _ => Set()

  }

}

// parser for list of arguments

def ListParser[I, T, S](p: => Parser[I, T], 

                        q: => Parser[I, S])(implicit ev: I => Seq[_]): Parser[I, List[T]] = {

  (p ~ q ~ ListParser(p, q)).map{ case x ~ _ ~ z => x :: z : List[T] } ||

  (p.map((s) => List(s)))

}

// I may want to write string interpolations for:

// keywords, semicolon, colon, comma, parentheses

case class StrParser(s: String) extends Parser[Tokens, String] {

  def parse(tokens: Tokens) = tokens match {

    case Nil => Set()

    case tk::tkns if tk._2 == s => Set((s, tkns))

    case _ => Set()

  }

}

extension (sc: StringContext) {

   def p(args: Any*) = StrParser(sc.s(args:_*))

}

// the AST datastructures for the FUN language

abstract class Exp

abstract class BExp

abstract class Decl

case class Def(name: String, args: List[(String, String)], ty: String, body: Exp) extends Decl

case class Main(e: Exp) extends Decl

case class Const(name: String, v: Int) extends Decl

case class FConst(name: String, x: Float) extends Decl

case class Call(name: String, args: List[Exp]) extends Exp

case class If(a: BExp, e1: Exp, e2: Exp) extends Exp

case class Var(s: String) extends Exp

case class Num(i: Int) extends Exp  // integer numbers

case class FNum(i: Float) extends Exp  // float numbers

case class ChConst(c: Int) extends Exp  // character constants

case class Aop(o: String, a1: Exp, a2: Exp) extends Exp

case class Sequence(e1: Exp, e2: Exp) extends Exp  // expressions separated by semicolons

case class Bop(o: String, a1: Exp, a2: Exp) extends BExp

lazy val Exps: Parser[Tokens, Exp] =

  (Exp ~ p";" ~ Exps).map[Exp]{ case x ~ _ ~ z => Sequence(x, z) } ||

  Exp

lazy val Exp: Parser[Tokens, Exp] =

  (p"if" ~ BExp ~ p"then" ~ Exp ~ p"else" ~ Exp).map[Exp]{ case _ ~ x ~ _ ~ y ~ _ ~ z => If(x, y, z) } ||

  M

lazy val M: Parser[Tokens, Exp] = 

  (T ~ OpParser(Set("+", "-")) ~ M).map[Exp]{ case x ~ y ~ z => Aop(y, x, z) } ||

  T

lazy val T: Parser[Tokens, Exp] = 

  (U ~ OpParser(Set("*", "/", "%")) ~ T).map[Exp]{ case x ~ y ~ z => Aop(y, x, z) } ||

  U

// includes negative factor

// a + - b CAN be recognised

// - - - b CAN be recognised

lazy val U: Parser[Tokens, Exp] =

  (OpParser(Set("-")) ~ U).map[Exp]{ case _ ~ y => Aop("*", Num(-1), y) } ||

  (OpParser(Set("+")) ~ U).map[Exp]{ case _ ~ y => y } ||

  F

lazy val F: Parser[Tokens, Exp] = 

  (p"(" ~ Exp ~ p")").map[Exp]{ case _ ~ y ~ _ => y } ||

  (p"skip").map(_ => Call("skip", Nil)) ||  // hardcoded

  (p"skip" ~ p"(" ~ p")").map(_ => Call("skip", Nil)) ||  // hardcoded

  (IdParser ~ p"(" ~ ListParser(Exp, p",") ~ p")").map[Exp]{ case id ~ _ ~ args ~ _ => Call(id, args) } ||

  (IdParser ~ p"(" ~ p")").map[Exp]{ case id ~ _ ~ _ => Call(id, Nil) } ||  // NOTE: empty args are also accepted!

  (IdParser || GlobalIdParser).map(x => Var(x)) ||

  IntParser.map(x => Num(x)) ||

  DoubleParser.map(x => FNum(x)) ||

  CharParser.map(x => ChConst(x.toInt)) ||

  (p"{" ~ Exps ~ p"}").map[Exp]{ case _ ~ x ~ _ => x }

lazy val BExp: Parser[Tokens, BExp] = 

  (Exp ~ OpParser(Set("==", "!=", "<", ">", "<=", ">=")) ~ Exp).map[BExp]{ case x ~ y ~ z => Bop(y, x, z) } ||

  (p"(" ~ BExp ~ p")").map[BExp]{ case _ ~ y ~ _ => y }

lazy val TypedIdParser: Parser[Tokens, (String, String)] =

  (IdParser ~ p":" ~ TypeParser(Set("Int", "Double"))).map{ case n ~ _ ~ t => (n, t) }

lazy val Defn: Parser[Tokens, Decl] =

  (p"def" ~ IdParser ~ p"(" ~ ListParser(TypedIdParser, p",") ~ p")" ~ p":" ~ TypeParser(Set("Int", "Double", "Void")) ~ OpParser(Set("=")) ~ Exp).map[Decl]{

    case _ ~ y ~ _ ~ w ~ _ ~ _ ~ t ~ _ ~ b => Def(y, w, t, b)

  } ||

  (p"def" ~ IdParser ~ p"(" ~ p")" ~ p":" ~ TypeParser(Set("Int", "Double", "Void")) ~ OpParser(Set("=")) ~ Exp).map[Decl]{

    case _ ~ y ~ _ ~ _ ~ _ ~ t ~ _ ~ b => Def(y, Nil, t, b)

  }

lazy val Constp: Parser[Tokens, Decl] = 

  (p"val" ~ GlobalIdParser ~ p":" ~ TypeParser(Set("Int")) ~ OpParser(Set("=")) ~ IntParser).map[Decl]{  // IntParser? Not Exp? For this AST, impossible to define Exp

    case _ ~ id ~ _ ~ _ ~ _ ~ n => Const(id, n)

  } ||

  (p"val" ~ GlobalIdParser ~ p":" ~ TypeParser(Set("Int")) ~ OpParser(Set("=")) ~ OpParser(Set("-")) ~ IntParser).map[Decl]{  // IntParser? Not Exp? For this AST, impossible to define Exp

    case _ ~ id ~ _ ~ _ ~ _ ~ _ ~ n => Const(id, -n)

  }

// Int can be converted to Double but not viceversa

lazy val FConstp: Parser[Tokens, Decl] =

  (p"val" ~ GlobalIdParser ~ p":" ~ TypeParser(Set("Double")) ~ OpParser(Set("=")) ~ (DoubleParser || IntParser.map[Float](i => i.toFloat))).map[Decl]{

    case _ ~ id ~ _ ~ _ ~ _ ~ n => FConst(id, n)

  } ||

  (p"val" ~ GlobalIdParser ~ p":" ~ TypeParser(Set("Double")) ~ OpParser(Set("=")) ~ OpParser(Set("-")) ~ (DoubleParser || IntParser.map[Float](i => i.toFloat))).map[Decl]{

    case _ ~ id ~ _ ~ _ ~ _ ~ _ ~ n => FConst(id, -n)

  }

// Prog consists of global const declarations, f(x) defs, and exp in ANY order

// restricted to main body at the bottom

lazy val Prog: Parser[Tokens, List[Decl]] = 

  (Defn ~ p";" ~ Prog).map[List[Decl]]{ case x ~ _ ~ z => x :: z } ||

  (Constp ~ p";" ~ Prog).map[List[Decl]]{ case x ~ _ ~ z => x :: z } ||

  (FConstp ~ p";" ~ Prog).map[List[Decl]]{ case x ~ _ ~ z => x :: z } ||

  Exp.map[List[Decl]](s => List(Main(s)))

def parse_tks(tokens: Tokens) = Prog.parse_all(tokens)

import scala.io.Source._

@main

def parse(filename: String) = {

  val fun_code = fromFile(filename).getLines.mkString("\n")

  // print the AST list to screen

  println(parse_tks(tokenise(fun_code)))

}