package wacc

import scala.collection.mutable

object renamer {
  import ast._
  import types._

  enum IdentType {
    case Func
    case Var
  }

  private case class Scope(
      current: mutable.Map[(String, IdentType), Ident],
      parent: Map[(String, IdentType), Ident]
  ) {
    def subscope: Scope =
      Scope(mutable.Map.empty, Map.empty.withDefault(current.withDefault(parent)))

    def add(semType: SemType, name: Ident, identType: IdentType)(using
        globalNames: mutable.Map[Ident, SemType],
        globalNumbering: mutable.Map[String, Int],
        errors: mutable.Builder[Error, List[Error]]
    ) = {
      if (current.contains((name.v, identType))) {
        errors += Error.DuplicateDeclaration(name)
      } else {
        val uid = globalNumbering.getOrElse(name.v, 0)
        name.uid = uid
        current((name.v, identType)) = name

        globalNames(name) = semType
        globalNumbering(name.v) = uid + 1
      }
    }
  }

  def rename(prog: Program)(using
      errors: mutable.Builder[Error, List[Error]]
  ): Map[Ident, SemType] =
    given globalNames: mutable.Map[Ident, SemType] = mutable.Map.empty
    given globalNumbering: mutable.Map[String, Int] = mutable.Map.empty
    rename(Scope(mutable.Map.empty, Map.empty))(prog)
    globalNames.toMap

  private def rename(scope: Scope)(
      node: Program | FuncDecl | Ident | Stmt | LValue | RValue
  )(using
      globalNames: mutable.Map[Ident, SemType],
      globalNumbering: mutable.Map[String, Int],
      errors: mutable.Builder[Error, List[Error]]
  ): Unit = node match {
    case Program(funcs, main) => {
      funcs.foreach(rename(scope))
      main.toList.foreach(rename(scope))
    }
    case FuncDecl(retType, name, params, body) => {
      val functionScope = scope.subscope
      val paramTypes = params.map { param =>
        val paramType = SemType(param.paramType)
        functionScope.add(paramType, param.name, IdentType.Var)
        paramType
      }
      scope.add(KnownType.Func(SemType(retType), paramTypes), name, IdentType.Func)
      body.toList.foreach(rename(functionScope))
    }
    case VarDecl(synType, name, value) => {
      // Order matters here. Variable isn't declared until after the value is evaluated.
      rename(scope)(value)
      scope.add(SemType(synType), name, IdentType.Var)
    }
    case Assign(lhs, value) => {
      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 If(cond, thenStmt, elseStmt) => {
      rename(scope)(cond)
      thenStmt.toList.foreach(rename(scope.subscope))
      elseStmt.toList.foreach(rename(scope.subscope))
    }
    case While(cond, body) => {
      rename(scope)(cond)
      body.toList.foreach(rename(scope.subscope))
    }
    case Block(body) => body.toList.foreach(rename(scope.subscope))
    case NewPair(fst, snd) => {
      rename(scope)(fst)
      rename(scope)(snd)
    }
    case Call(name, args) => {
      renameIdent(scope, name, IdentType.Func)
      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 => renameIdent(scope, id, IdentType.Var)
    case IntLiter(_) | BoolLiter(_) | CharLiter(_) | StrLiter(_) | PairLiter | Skip => ()
  }

  private def renameIdent(scope: Scope, ident: Ident, identType: IdentType)(using
      globalNames: mutable.Map[Ident, SemType],
      globalNumbering: mutable.Map[String, Int],
      errors: mutable.Builder[Error, List[Error]]
  ): Unit = {
    scope.current.withDefault(scope.parent).get((ident.v, identType)) match {
      case Some(Ident(_, uid)) => ident.uid = uid
      case None => {
        errors += Error.UndefinedIdentifier(ident, identType)
        scope.add(?, ident, identType)
      }
    }
  }
}