package wacc

import java.io.File
import scala.collection.mutable
import cats.effect.IO
import cats.implicits._
import cats.data.Chain
import cats.data.NonEmptyList
import parsley.{Failure, Success}

object renamer {
  import ast._
  import types._

  val MAIN = "$main"

  enum IdentType {
    case Func
    case Var
  }

  case class ScopeKey(path: String, name: String, identType: IdentType)
  case class ScopeValue(id: Ident, public: Boolean)

  class Scope(
      private val current: mutable.Map[ScopeKey, ScopeValue],
      private val parent: Map[ScopeKey, ScopeValue],
      guidStart: Int = 0,
      val guidInc: Int = 1
  ) {
    private var guid = guidStart
    private var immutable = false

    private def nextGuid(): Int = {
      val id = guid
      guid += guidInc
      id
    }

    private def verifyMutable(): Unit = {
      if (immutable) throw new IllegalStateException("Cannot modify an immutable scope")
    }

    /** Create a new scope with the current scope as its parent.
      *
      * To be used for single-threaded applications.
      *
      * @return
      *   A new scope with an empty current scope, and this scope flattened into the parent scope.
      */
    def withSubscope[T](f: Scope => T): T = {
      val subscope =
        Scope(mutable.Map.empty, Map.empty.withDefault(current.withDefault(parent)), guid, guidInc)
      immutable = true
      val result = f(subscope)
      guid = subscope.guid // Sync GUID
      immutable = false
      result
    }

    /** Create new scopes with the current scope as its parent and GUID numbering adjusted
      * correctly.
      *
      * This will permanently mark the current scope as immutable, for thread safety.
      *
      * To be used for multi-threaded applications.
      *
      * @return
      *   New scopes with an empty current scope, and this scope flattened into the parent scope.
      */
    def subscopes(n: Int): Seq[Scope] = {
      verifyMutable()
      immutable = true
      (0 until n).map { i =>
        Scope(
          mutable.Map.empty,
          Map.empty.withDefault(current.withDefault(parent)),
          guid + i * guidInc,
          guidInc * n
        )
      }
    }

    /** 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 name
      *   The name of the identifier.
      * @return
      *   An error, if one occurred.
      */
    def add(name: Ident, public: Boolean = false): Chain[Error] = {
      verifyMutable()
      val path = name.pos.file.getCanonicalPath
      val identType = name.ty match {
        case _: SemType  => IdentType.Var
        case _: FuncType => IdentType.Func
      }
      val key = ScopeKey(path, name.v, identType)
      current.get(key) match {
        case Some(ScopeValue(Ident(_, id, _), _)) =>
          name.guid = id
          Chain.one(Error.DuplicateDeclaration(name))
        case None =>
          name.guid = nextGuid()
          current(key) = ScopeValue(name, public)
          Chain.empty
      }
    }

    /** Attempt to add a new identifier as an alias to another to the existing scope.
      *
      * @param alias
      *   The (new) alias identifier.
      * @param orig
      *   The (existing) original identifier.
      *
      * @return
      *   An error, if one occurred.
      */
    def addAlias(alias: Ident, orig: ScopeValue, public: Boolean = false): Chain[Error] = {
      verifyMutable()
      val path = alias.pos.file.getCanonicalPath
      val identType = alias.ty match {
        case _: SemType  => IdentType.Var
        case _: FuncType => IdentType.Func
      }
      val key = ScopeKey(path, alias.v, identType)
      current.get(key) match {
        case Some(ScopeValue(Ident(_, id, _), _)) =>
          alias.guid = id
          Chain.one(Error.DuplicateDeclaration(alias))
        case None =>
          alias.guid = nextGuid()
          current(key) = ScopeValue(orig.id, public)
          Chain.empty
      }
    }

    def get(path: String, name: String, identType: IdentType): Option[ScopeValue] =
      // 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)(ScopeKey(path, name, identType)))
      } catch {
        case _: NoSuchElementException => None
      }

    def getVar(name: Ident): Option[Ident] =
      get(name.pos.file.getCanonicalPath, name.v, IdentType.Var).map(_.id)
    def getFunc(name: Ident): Option[Ident] =
      get(name.pos.file.getCanonicalPath, name.v, IdentType.Func).map(_.id)
  }

  def prepareGlobalScope(
      partialProg: PartialProgram
  )(using scope: Scope): IO[(FuncDecl, Chain[FuncDecl], Chain[Error])] = {
    def readImportFile(file: File): IO[String] =
      IO.blocking(os.read(os.Path(file.getCanonicalPath)))

    def prepareImport(contents: String, file: File)(using
        scope: Scope
    ): IO[(Chain[FuncDecl], Chain[Error])] = {
      parser.parse(contents) match {
        case Failure(msg) =>
          IO.pure(Chain.empty, Chain.one(Error.SyntaxError(file, msg)))
        case Success(fn) =>
          val partialProg = fn(file)
          for {
            (main, chunks, errors) <- prepareGlobalScope(partialProg)
          } yield (main +: chunks, errors)
      }
    }

    def addImportsToScope(importFile: File, funcs: NonEmptyList[ImportedFunc])(using
        scope: Scope
    ): Chain[Error] =
      funcs.foldMap { case ImportedFunc(srcName, aliasName) =>
        scope.get(importFile.getCanonicalPath, srcName.v, IdentType.Func) match {
          case Some(src) if src.public =>
            aliasName.ty = src.id.ty
            scope.addAlias(aliasName, src)
          case _ =>
            Chain.one(Error.UndefinedFunction(srcName))
        }
      }

    val PartialProgram(imports, prog) = partialProg

    // First prepare this file's functions...
    val Program(funcs, main) = prog
    val (funcChunks, funcErrors) = funcs.foldLeft((Chain.empty[FuncDecl], Chain.empty[Error])) {
      case ((chunks, errors), func @ FuncDecl(retType, name, params, body)) =>
        val paramTypes = params.map { param =>
          val paramType = SemType(param.paramType)
          param.name.ty = paramType
          paramType
        }
        name.ty = FuncType(SemType(retType), paramTypes)
        (chunks :+ func, errors ++ scope.add(name, public = true))
    }
    // ...and main body.
    val mainBodyIdent = Ident(MAIN, ty = FuncType(?, Nil))(prog.pos)
    val mainBodyErrors = scope.add(mainBodyIdent, public = false)
    val mainBodyChunk = FuncDecl(IntType()(prog.pos), mainBodyIdent, Nil, main)(prog.pos)

    // Now handle imports
    val file = prog.pos.file
    val preparedImports = imports.foldLeftM[IO, (Chain[FuncDecl], Chain[Error])](
      (Chain.empty[FuncDecl], Chain.empty[Error])
    ) { case ((chunks, errors), Import(name, funcs)) =>
      val importFile = File(file.getParent, name.v)
      if (!importFile.exists()) {
        IO.pure(
          (
            chunks,
            errors :+ Error.SemanticError(
              name.pos,
              s"File not found: ${importFile.getCanonicalPath}"
            )
          )
        )
      } else if (!importFile.canRead()) {
        IO.pure(
          (
            chunks,
            errors :+ Error.SemanticError(
              name.pos,
              s"File not readable: ${importFile.getCanonicalPath}"
            )
          )
        )
      } else if (importFile.getCanonicalPath == file.getCanonicalPath) {
        IO.pure(
          (
            chunks,
            errors :+ Error.SemanticError(
              name.pos,
              s"Cannot import self: ${importFile.getCanonicalPath}"
            )
          )
        )
      } else if (scope.get(importFile.getCanonicalPath, MAIN, IdentType.Func).isDefined) {
        IO.pure(chunks, errors ++ addImportsToScope(importFile, funcs))
      } else {
        for {
          contents <- readImportFile(importFile)
          (importChunks, importErrors) <- prepareImport(contents, importFile)
          importAliasErrors = addImportsToScope(importFile, funcs)
        } yield (chunks ++ importChunks, errors ++ importErrors)
      }
    }

    for {
      (importChunks, importErrors) <- preparedImports
      allChunks = importChunks ++ funcChunks
      allErrors = importErrors ++ funcErrors ++ mainBodyErrors
    } yield (mainBodyChunk, allChunks, allErrors)
  }

  /** Check scoping of all variables and flatten a program. Also generates semantic types and parses
    * any imported files.
    *
    * @param partialProg
    *   AST of the program
    * @return
    *   (flattenedProg, errors)
    */
  def renameFunction(funcScopePair: (FuncDecl, Scope)): IO[Chain[Error]] = {
    val (FuncDecl(_, _, params, body), subscope) = funcScopePair
    val paramErrors = params.foldMap(param => subscope.add(param.name))
    IO(subscope.withSubscope { s => body.foldMap(rename(s)) })
      .map(bodyErrors => paramErrors ++ bodyErrors)
  }

  /** Check scoping of all identifies in a given AST node.
    *
    * @param scope
    *   The current scope and flattened parent scope.
    * @param node
    *   The AST node.
    */
  private def rename(scope: Scope)(node: Ident | Stmt | LValue | RValue | Expr): Chain[Error] =
    node match {
      // These cases are more interes/globting 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.
        val errors = rename(scope)(value)
        // Attempt to add the new variable to the current scope.
        name.ty = SemType(synType)
        errors ++ scope.add(name)
      }
      case If(cond, thenStmt, elseStmt) => {
        val condErrors = rename(scope)(cond)
        // then and else both have their own scopes
        val thenErrors = scope.withSubscope(s => thenStmt.foldMap(rename(s)))
        val elseErrors = scope.withSubscope(s => elseStmt.foldMap(rename(s)))
        condErrors ++ thenErrors ++ elseErrors
      }
      case While(cond, body) => {
        val condErrors = rename(scope)(cond)
        // while bodies have their own scopes
        val bodyErrors = scope.withSubscope(s => body.foldMap(rename(s)))
        condErrors ++ bodyErrors
      }
      // begin-end blocks have their own scopes
      case Block(body) => scope.withSubscope(s => body.foldMap(rename(s)))

      // 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) => {
        val nameErrors = scope.getFunc(name) match {
          case Some(Ident(realName, guid, ty)) =>
            name.v = realName
            name.ty = ty
            name.guid = guid
            Chain.empty
          case None =>
            name.ty = FuncType(?, args.map(_ => ?))
            scope.add(name)
            Chain.one(Error.UndefinedFunction(name))
        }
        val argsErrors = args.foldMap(rename(scope))
        nameErrors ++ argsErrors
      }
      case Fst(elem)         => rename(scope)(elem)
      case Snd(elem)         => rename(scope)(elem)
      case ArrayLiter(elems) => elems.foldMap(rename(scope))
      case ArrayElem(name, indices) => {
        val nameErrors = rename(scope)(name)
        val indicesErrors = indices.foldMap(rename(scope))
        nameErrors ++ indicesErrors
      }
      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) match {
          case Some(Ident(_, guid, ty)) =>
            id.ty = ty
            id.guid = guid
            Chain.empty
          case None =>
            id.ty = ?
            scope.add(id)
            Chain.one(Error.UndeclaredVariable(id))
        }
      }
      // These literals cannot contain identifies, exit immediately.
      case IntLiter(_) | BoolLiter(_) | CharLiter(_) | StrLiter(_) | PairLiter() | Skip() =>
        Chain.empty
    }
}
