// CW3

import $file.lexer

import lexer._ 

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)

}

/*

// atomic parser for (particular) strings

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

  def parse(sb: String) = {

    val (prefix, suffix) = sb.splitAt(s.length)

    if (prefix == s) Set((prefix, suffix)) else Set()

  }

}

extension (sc: StringContext) 

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

*/

// 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)

}

// New parser that takes as input a list of tokens

case class TokenParser(t: Token) extends Parser[List[Token], Token] {

    def parse(in: List[Token]) = {

      // an example of an atomic parser for characters

      if (!in.isEmpty && in.head == t) Set((t, in.tail)) else Set()

    }

}   

case class TokenListParser(ts: List[Token]) extends Parser[List[Token], List[Token]] {

    def parse(tsb: List[Token]) = {

        val (prefix, suffix) = tsb.splitAt(ts.length)

        if (prefix == ts) Set((prefix, suffix)) else Set()

    }

}

// Implicit definitions to go from a token 

// or a list of tokens to a TokenListParser

implicit def token2parser(t: Token) : Parser[List[Token], Token] = 

  TokenParser(t)

extension (t: Token) {

    def || (q : => Parser[List[Token], Token]) = 

      new AltParser[List[Token], Token](t, q)

    def ~[S](q : => Parser[List[Token], S]) = 

      new SeqParser[List[Token], Token, S](t, q)  

}

// Abstract Syntax Trees

abstract class Stmt

abstract class AExp

abstract class BExp

type Block = List[Stmt]

case object Skip extends Stmt

case class If(a: BExp, bl1: Block, bl2: Block) extends Stmt

case class While(b: BExp, bl: Block) extends Stmt

case class Assign(s: String, a: AExp) extends Stmt

case class Read(s: String) extends Stmt

case class WriteId(s: String) extends Stmt  // for printing values of variables

case class WriteString(s: String) extends Stmt  // for printing words

case class Var(s: String) extends AExp

case class Num(i: Int) extends AExp

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

case object True extends BExp

case object False extends BExp

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

case class And(b1: BExp, b2: BExp) extends BExp

case class Or(b1: BExp, b2: BExp) extends BExp

case class IdParser() extends Parser[List[Token], String] {

    def parse(tsb: List[Token]) = tsb match {

        case T_ID(id) :: rest => Set((id, rest))

        case _ => Set()

    }

}

case class NumParser() extends Parser[List[Token], Int] {

    def parse(tsb: List[Token]) = tsb match {

        case T_NUM(n) :: rest => Set((n, rest))

        case _ => Set()

    }

}

case class StringParser() extends Parser[List[Token], String] {

    def parse(tsb: List[Token]) = tsb match {

        case T_STRING(s) :: rest => Set((s, rest))

        case _ => Set()

    }

}

// WHILE Language Parsing

lazy val AExp: Parser[List[Token], AExp] = 

  (Te ~ T_OP("+") ~ AExp).map{ case x ~ _ ~ z => Aop("+", x, z): AExp } ||

  (Te ~ T_OP("-") ~ AExp).map{ case x ~ _ ~ z => Aop("-", x, z): AExp } || Te

lazy val Te: Parser[List[Token], AExp] = 

  (Fa ~ T_OP("*") ~ Te).map{ case x ~ _ ~ z => Aop("*", x, z): AExp } || 

  (Fa ~ T_OP("/") ~ Te).map{ case x ~ _ ~ z => Aop("/", x, z): AExp } || 

  (Fa ~ T_OP("%") ~ Te).map{ case x ~ _ ~ z => Aop("%", x, z): AExp } || Fa  

lazy val Fa: Parser[List[Token], AExp] = 

   (T_PAREN("(") ~ AExp ~ T_PAREN(")")).map{ case _ ~ y ~ _ => y } || 

   IdParser().map{Var(_)}  || 

   NumParser().map{Num(_)}

lazy val BExp: Parser[List[Token], BExp] = 

   (AExp ~ T_OP("==") ~ AExp).map{ case x ~ _ ~ z => Bop("==", x, z): BExp } || 

   (AExp ~ T_OP("!=") ~ AExp).map{ case x ~ _ ~ z => Bop("!=", x, z): BExp } || 

   (AExp ~ T_OP("<") ~ AExp).map{ case x ~ _ ~ z => Bop("<", x, z): BExp } || 

   (AExp ~ T_OP(">") ~ AExp).map{ case x ~ _ ~ z => Bop(">", x, z): BExp } ||

   (T_PAREN("(") ~ BExp ~ T_PAREN(")") ~ T_OP("&&") ~ BExp).map{ case _ ~ y ~ _ ~ _ ~ v => And(y, v): BExp } ||

   (T_PAREN("(") ~ BExp ~ T_PAREN(")") ~ T_OP("||") ~ BExp).map{ case _ ~ y ~ _ ~ _ ~ v => Or(y, v): BExp } ||

   (T_KEYWORD("true").map(_ => True: BExp )) || 

   (T_KEYWORD("false").map(_ => False: BExp )) ||

   (T_PAREN("(") ~ BExp ~ T_PAREN(")")).map{ case _ ~ x ~ _ => x }

lazy val Stmt: Parser[List[Token], Stmt] =

    T_KEYWORD("skip").map(_ => Skip: Stmt) ||

    (IdParser() ~ T_OP(":=") ~ AExp).map{ case id ~ _ ~ z => Assign(id, z): Stmt } ||

    (T_KEYWORD("if") ~ BExp ~ T_KEYWORD("then") ~ Block ~ T_KEYWORD("else") ~ Block).map{ case _ ~ y ~ _ ~ u ~ _ ~ w => If(y, u, w): Stmt } ||

    (T_KEYWORD("while") ~ BExp ~ T_KEYWORD("do") ~ Block).map{ case _ ~ y ~ _ ~ w => While(y, w) : Stmt } ||

    (T_KEYWORD("read") ~ IdParser()).map{ case _ ~ id => Read(id): Stmt} ||

    (T_KEYWORD("write") ~ IdParser()).map{ case _ ~ id => WriteId(id): Stmt} ||

    (T_KEYWORD("write") ~ StringParser()).map{ case _ ~ s => WriteString(s): Stmt} || 

    (T_KEYWORD("write") ~ T_PAREN("(") ~ IdParser() ~ T_PAREN(")")).map{ case _ ~ _ ~ id ~ _ => WriteId(id): Stmt} ||

    (T_KEYWORD("write") ~  T_PAREN("(") ~ StringParser() ~ T_PAREN(")")).map{ case _ ~ _ ~ s ~ _ => WriteString(s): Stmt}

lazy val Stmts: Parser[List[Token], Block] =

    (Stmt ~ T_SEMI ~ Stmts).map{ case x ~ _ ~ z => x :: z : Block } ||

    (Stmt.map(s => List(s) : Block))

lazy val Block: Parser[List[Token], Block] =

    (T_PAREN("{") ~ Stmts ~ T_PAREN("}")).map{ case x ~ y ~ z => y} ||

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

// Testing with programs 2 & 3

//println("Fibonacci")

//println(Stmts.parse_all(tokenise(os.read(os.pwd / "fib.while"))))

//println("Loops")

//println(Stmts.parse_all(tokenise(os.read(os.pwd / "loops.while"))))

//println("Collatz")

//println(Stmts.parse_all(tokenise(os.read(os.pwd / "collatz2.while"))))

// Interpreter

// Environment to store values of variables

type Env = Map[String, Int]

def eval_aexp(a: AExp, env: Env) : Int = a match {

  case Num(i) => i

  case Var(s) => env(s)

  case Aop("+", a1, a2) => eval_aexp(a1, env) + eval_aexp(a2, env)

  case Aop("-", a1, a2) => eval_aexp(a1, env) - eval_aexp(a2, env)

  case Aop("*", a1, a2) => eval_aexp(a1, env) * eval_aexp(a2, env)

  case Aop("/", a1, a2) => eval_aexp(a1, env) / eval_aexp(a2, env)

  case Aop("%", a1, a2) => eval_aexp(a1, env) % eval_aexp(a2, env)

}

def eval_bexp(b: BExp, env: Env) : Boolean = b match {

  case True => true

  case False => false

  case Bop("==", a1, a2) => eval_aexp(a1, env) == eval_aexp(a2, env)

  case Bop("!=", a1, a2) => !(eval_aexp(a1, env) == eval_aexp(a2, env))

  case Bop(">", a1, a2) => eval_aexp(a1, env) > eval_aexp(a2, env)

  case Bop("<", a1, a2) => eval_aexp(a1, env) < eval_aexp(a2, env)

  case And(b1, b2) => eval_bexp(b1, env) && eval_bexp(b2, env)

  case Or(b1, b2) => eval_bexp(b1, env) || eval_bexp(b2, env)

}

// Import needed to take int as input from the user

import scala.io.StdIn.readInt

def eval_stmt(s: Stmt, env: Env) : Env = s match {

  case Skip => env

  case Assign(x, a) => env + (x -> eval_aexp(a, env))

  case If(b, bl1, bl2) => if (eval_bexp(b, env)) eval_bl(bl1, env) else eval_bl(bl2, env) 

  case While(b, bl) => 

    if (eval_bexp(b, env)) eval_stmt(While(b, bl), eval_bl(bl, env))

    else env

  case WriteId(x) => { print(env(x)) ; env }

  case WriteString(x) => {

        print(x.replaceAll("\"", "").replaceAll("""\\n""", "\n")) ;

        env

       }

  case Read(x) => { 

      println("Enter an integer and press ENTER:") ; 

      val n = readInt() ; // Note: Does not work when using the REPL

      eval_stmt(Assign(x, Num(n)), env) 

  }

}

def eval_bl(bl: Block, env: Env) : Env = bl match {

  case Nil => env

  case s::bl => eval_bl(bl, eval_stmt(s, env))

}

def eval(bl: Block) : Env = eval_bl(bl, Map())

@main

def run(file: String) = {

  val contents = os.read(os.pwd / file)

  println(s"Lex $file: ")

  println(tokenise(contents))

  println(s"Parse $file: ")

  println(Stmts.parse_all(tokenise(contents)).head)

  println(s"Eval $file: ")

  println(eval(Stmts.parse_all(tokenise(contents)).head))

}

@main

def test(file: String) = {

  val contents = os.read(os.pwd / file)

  println(s"Lex $file: ")

  println(tokenise(contents))

  println(s"Parse $file: ")

  println(Stmts.parse_all(tokenise(contents)).head)

}

/*

println("Loops eval")

val start = System.nanoTime()

val end = System.nanoTime()

println("Time taken in seconds: ")

println((end - start)/(1.0e9))

*/