package wacc

import parsley.Result
import parsley.Parsley
import parsley.Parsley.{atomic, many, pure}
import parsley.combinator.{countSome, sepBy}
import parsley.expr.{precedence, SOps, InfixL, InfixN, InfixR, Prefix, Atoms}
import parsley.errors.combinator._
import parsley.cats.combinator.{sepBy1, some}
import cats.data.NonEmptyList

object parser {
  import lexer.implicits.implicitSymbol
  import lexer.{ident, integer, charLit, stringLit}
  import ast._

  def parse(input: String): Result[String, Program] = parser.parse(input)
  private val parser = lexer.fully(`<program>`)

  // Expressions
  private lazy val `<expr>`: Parsley[Expr] = precedence {
    SOps(InfixR)(Or from "||") +:
      SOps(InfixR)(And from "&&") +:
      SOps(InfixN)(Eq from "==", Neq from "!=") +:
      SOps(InfixN)(
        Less from "<",
        LessEq from "<=",
        Greater from ">",
        GreaterEq from ">="
      ) +:
      SOps(InfixL)(Add from "+", Sub from "-") +:
      SOps(InfixL)(Mul from "*", Div from "/", Mod from "%") +:
      SOps(Prefix)(
        Not from "!",
        Negate from "-",
        Len from "len",
        Ord from "ord",
        Chr from "chr"
      ) +:
      `<atom>`
  }

  // Atoms
  private lazy val `<atom>`: Atoms[Expr6] = Atoms(
    IntLiter(integer),
    BoolLiter(("true" as true) | ("false" as false)),
    CharLiter(charLit),
    StrLiter(stringLit),
    PairLiter from "null",
    `<ident-or-array-elem>`,
    Parens("(" ~> `<expr>` <~ ")")
  )
  private val `<ident>` = Ident(ident)
  private lazy val `<ident-or-array-elem>` =
    `<ident>` <**> (`<array-indices>` </> identity)
  private val `<array-indices>` =
    some("[" ~> `<expr>` <~ "]") map { indices =>
      ArrayElem((_: Ident), indices)
    }

  // Types
  private lazy val `<type>`: Parsley[Type] =
    (`<base-type>` | (`<pair-type>` ~> `<pair-elems-type>`)) <**> (`<array-type>` </> identity)
  private val `<base-type>` =
    (IntType from "int") | (BoolType from "bool") | (CharType from "char") | (StringType from "string")
  private lazy val `<array-type>` =
    countSome("[" ~> "]") map { cnt => ArrayType((_: Type), cnt) }
  private val `<pair-type>` = "pair"
  private val `<pair-elems-type>`: Parsley[PairType] = PairType(
    "(" ~> `<pair-elem-type>` <~ ",",
    `<pair-elem-type>` <~ ")"
  )
  private lazy val `<pair-elem-type>` =
    (`<base-type>` <**> (`<array-type>` </> identity)) |
      `<pair-type>` ~> ((`<pair-elems-type>` <**> `<array-type>`) </> UntypedPairType)

  // Statements
  private lazy val `<program>` = Program(
    "begin" ~> many(atomic(`<type>` <~> `<ident>` <~ "(") <**> `<partial-func-decl>`),
    `<stmt>` <~ "end"
  )
  private lazy val `<partial-func-decl>` =
    (sepBy(`<param>`, ",") <~ ")" <~ "is" <~> `<stmt>`.guardAgainst {
      case stmts if !stmts.isReturning => Seq("All functions must end in a returning statement")
    } <~ "end") map { (params, stmt) =>
      (FuncDecl((_: Type), (_: Ident), params, stmt)).tupled
    }
  private lazy val `<param>` = Param(`<type>`, `<ident>`)
  private lazy val `<stmt>`: Parsley[NonEmptyList[Stmt]] =
    sepBy1(`<basic-stmt>`, ";")
  private lazy val `<basic-stmt>` =
    (Skip from "skip")
      | Read("read" ~> `<lvalue>`)
      | Free("free" ~> `<expr>`)
      | Return("return" ~> `<expr>`)
      | Exit("exit" ~> `<expr>`)
      | Print("print" ~> `<expr>`, pure(false))
      | Print("println" ~> `<expr>`, pure(true))
      | If(
        "if" ~> `<expr>` <~ "then",
        `<stmt>` <~ "else",
        `<stmt>` <~ "fi"
      )
      | While("while" ~> `<expr>` <~ "do", `<stmt>` <~ "done")
      | Block("begin" ~> `<stmt>` <~ "end")
      | VarDecl(`<type>`, `<ident>` <~ "=", `<rvalue>`)
      | Assign(`<ident>` <~ "=", `<rvalue>`)
  private lazy val `<lvalue>`: Parsley[LValue] =
    `<pair-elem>` | `<ident-or-array-elem>`
  private lazy val `<rvalue>`: Parsley[RValue] =
    `<array-liter>` |
      NewPair(
        "newpair" ~> "(" ~> `<expr>` <~ ",",
        `<expr>` <~ ")"
      ) |
      `<pair-elem>` |
      Call(
        "call" ~> `<ident>` <~ "(",
        sepBy(`<expr>`, ",") <~ ")"
      ) | `<expr>`
  private lazy val `<pair-elem>` =
    Fst("fst" ~> `<lvalue>`) | Snd("snd" ~> `<lvalue>`)
  private lazy val `<array-liter>` = ArrayLiter(
    "[" ~> sepBy(`<expr>`, ",") <~ "]"
  )

  extension (stmts: NonEmptyList[Stmt]) {
    def isReturning: Boolean = stmts.last match {
      case Return(_) | Exit(_)       => true
      case If(_, thenStmt, elseStmt) => thenStmt.isReturning && elseStmt.isReturning
      case While(_, body)            => body.isReturning
      case Block(body)               => body.isReturning
      case _                         => false
    }
  }
}