package wacc

import scala.collection.mutable

object renamer {
  import ast._
  import types._

  private enum IdentType {
    case Func
    case Var
  }

  private class Scope(
      val current: mutable.Map[(String, IdentType), Ident],
      val parent: Map[(String, IdentType), Ident]
  ) {

    /** Create a new scope with the current scope as its parent.
      *
      * @return
      *   A new scope with an empty current scope, and this scope flattened into the parent scope.
      */
    def subscope: Scope =
      Scope(mutable.Map.empty, Map.empty.withDefault(current.withDefault(parent)))

    /** Attempt to add a new identifier to the current scope. If the identifier already exists in
      * the current scope, add an error to the error list.
      *
      * @param ty
      *   The semantic type of the variable identifier, or function identifier type.
      * @param name
      *   The name of the identifier.
      * @param globalNames
      *   The global map of identifiers to semantic types - the identifier will be added to this
      *   map.
      * @param globalNumbering
      *   The global map of identifier names to the number of times they have been declared - will
      *   used to rename this identifier, and will be incremented.
      * @param errors
      *   The list of errors to append to.
      */
    def add(ty: SemType | FuncType, name: Ident)(using
        globalNames: mutable.Map[Ident, SemType],
        globalFuncs: mutable.Map[Ident, FuncType],
        globalNumbering: mutable.Map[String, Int],
        errors: mutable.Builder[Error, List[Error]]
    ) = {
      val identType = ty match {
        case _: SemType  => IdentType.Var
        case _: FuncType => IdentType.Func
      }
      current.get((name.v, identType)) match {
        case Some(Ident(_, uid)) =>
          errors += Error.DuplicateDeclaration(name)
          name.uid = uid
        case None =>
          val uid = globalNumbering.getOrElse(name.v, 0)
          name.uid = uid
          current((name.v, identType)) = name

          ty match {
            case semType: SemType =>
              globalNames(name) = semType
            case funcType: FuncType =>
              globalFuncs(name) = funcType
          }
          globalNumbering(name.v) = uid + 1
      }
    }

    private def get(name: String, identType: IdentType): Option[Ident] =
      // Unfortunately map defaults only work with `.apply()`, which throws an error when the key is not found.
      // Neither is there a way to check whether a default exists, so we have to use a try-catch.
      try {
        Some(current.withDefault(parent)((name, identType)))
      } catch {
        case _: NoSuchElementException => None
      }

    def getVar(name: String): Option[Ident] = get(name, IdentType.Var)
    def getFunc(name: String): Option[Ident] = get(name, IdentType.Func)
  }

  /** Check scoping of all variables and functions in the program. Also generate semantic types for
    * all identifiers.
    *
    * @param prog
    *   AST of the program
    * @param errors
    *   List of errors to append to
    * @return
    *   Map of all (renamed) identifies to their semantic types
    */
  def rename(prog: Program)(using
      errors: mutable.Builder[Error, List[Error]]
  ): (Map[Ident, SemType], Map[Ident, FuncType]) = {
    given globalNames: mutable.Map[Ident, SemType] = mutable.Map.empty
    given globalFuncs: mutable.Map[Ident, FuncType] = mutable.Map.empty
    given globalNumbering: mutable.Map[String, Int] = mutable.Map.empty
    val scope = Scope(mutable.Map.empty, Map.empty)
    val Program(funcs, main) = prog
    funcs
      // First add all function declarations to the scope
      .map { case FuncDecl(retType, name, params, body) =>
        val paramTypes = params.map { param =>
          val paramType = SemType(param.paramType)
          paramType
        }
        scope.add(FuncType(SemType(retType), paramTypes), name)
        (params zip paramTypes, body)
      }
      // Only then rename the function bodies
      // (functions can call one-another regardless of order of declaration)
      .foreach { case (params, body) =>
        val functionScope = scope.subscope
        params.foreach { case (param, paramType) =>
          functionScope.add(paramType, param.name)
        }
        body.toList.foreach(rename(functionScope.subscope)) // body can shadow function params
      }
    main.toList.foreach(rename(scope))
    (globalNames.toMap, globalFuncs.toMap)
  }

  /** Check scoping of all identifies in a given AST node.
    *
    * @param scope
    *   The current scope and flattened parent scope.
    * @param node
    *   The AST node.
    * @param globalNames
    *   The global map of identifiers to semantic types - renamed identifiers will be added to this
    *   map.
    * @param globalNumbering
    *   The global map of identifier names to the number of times they have been declared - used and
    *   updated during identifier renaming.
    * @param errors
    */
  private def rename(scope: Scope)(
      node: Ident | Stmt | LValue | RValue | Expr
  )(using
      globalNames: mutable.Map[Ident, SemType],
      globalFuncs: mutable.Map[Ident, FuncType],
      globalNumbering: mutable.Map[String, Int],
      errors: mutable.Builder[Error, List[Error]]
  ): Unit = node match {
    // These cases are more interesting because the involve making subscopes
    // or modifying the current scope.
    case VarDecl(synType, name, value) => {
      // Order matters here. Variable isn't declared until after the value is evaluated.
      rename(scope)(value)
      // Attempt to add the new variable to the current scope.
      scope.add(SemType(synType), name)
    }
    case If(cond, thenStmt, elseStmt) => {
      rename(scope)(cond)
      // then and else both have their own scopes
      thenStmt.toList.foreach(rename(scope.subscope))
      elseStmt.toList.foreach(rename(scope.subscope))
    }
    case While(cond, body) => {
      rename(scope)(cond)
      // while bodies have their own scopes
      body.toList.foreach(rename(scope.subscope))
    }
    // begin-end blocks have their own scopes
    case Block(body) => body.toList.foreach(rename(scope.subscope))

    // These cases are simpler, mostly just recursive calls to rename()
    case Assign(lhs, value) => {
      // Variables may be reassigned with their value in the rhs, so order doesn't matter here.
      rename(scope)(lhs)
      rename(scope)(value)
    }
    case Read(lhs)      => rename(scope)(lhs)
    case Free(expr)     => rename(scope)(expr)
    case Return(expr)   => rename(scope)(expr)
    case Exit(expr)     => rename(scope)(expr)
    case Print(expr, _) => rename(scope)(expr)
    case NewPair(fst, snd) => {
      rename(scope)(fst)
      rename(scope)(snd)
    }
    case Call(name, args) => {
      scope.getFunc(name.v) match {
        case Some(Ident(_, uid)) => name.uid = uid
        case None =>
          errors += Error.UndefinedFunction(name)
          scope.add(FuncType(?, args.map(_ => ?)), name)
      }
      args.foreach(rename(scope))
    }
    case Fst(elem)         => rename(scope)(elem)
    case Snd(elem)         => rename(scope)(elem)
    case ArrayLiter(elems) => elems.foreach(rename(scope))
    case ArrayElem(name, indices) => {
      rename(scope)(name)
      indices.toList.foreach(rename(scope))
    }
    case Parens(expr) => rename(scope)(expr)
    case op: UnaryOp  => rename(scope)(op.x)
    case op: BinaryOp => {
      rename(scope)(op.x)
      rename(scope)(op.y)
    }
    // Default to variables. Only `call` uses IdentType.Func.
    case id: Ident => {
      scope.getVar(id.v) match {
        case Some(Ident(_, uid)) => id.uid = uid
        case None =>
          errors += Error.UndeclaredVariable(id)
          scope.add(?, id)
      }
    }
    // These literals cannot contain identifies, exit immediately.
    case IntLiter(_) | BoolLiter(_) | CharLiter(_) | StrLiter(_) | PairLiter() | Skip() => ()
  }
}