package wacc

import cats.syntax.all._
import cats.data.NonEmptyList
import cats.data.Chain

object typeChecker {
  import wacc.types._

  private enum Constraint {
    case Unconstrained
    // Allows weakening in one direction
    case Is(ty: SemType, msg: String)
    // Allows weakening in both directions, useful for array literals
    case IsSymmetricCompatible(ty: SemType, msg: String)
    // Does not allow weakening
    case IsUnweakenable(ty: SemType, msg: String)
    case IsEither(ty1: SemType, ty2: SemType, msg: String)
    case Never(msg: String)
  }

  extension (ty: SemType) {

    /** Check if a type satisfies a constraint.
      *
      * @param constraint
      *   Constraint to satisfy.
      * @param pos
      *   Position to pass to the error, if constraint was not satisfied.
      * @return
      *   The type if the constraint was satisfied, or ? if it was not.
      */
    private def satisfies(constraint: Constraint, pos: ast.Position): (SemType, Chain[Error]) =
      (ty, constraint) match {
        case (KnownType.Array(KnownType.Char), Constraint.Is(KnownType.String, _)) =>
          (KnownType.String, Chain.empty)
        case (
              KnownType.String,
              Constraint.IsSymmetricCompatible(KnownType.Array(KnownType.Char), _)
            ) =>
          (KnownType.String, Chain.empty)
        case (ty, Constraint.IsSymmetricCompatible(ty2, msg)) =>
          ty.satisfies(Constraint.Is(ty2, msg), pos)
        // Change to IsUnweakenable to disallow recursive weakening
        case (ty, Constraint.Is(ty2, msg)) => ty.satisfies(Constraint.IsUnweakenable(ty2, msg), pos)
        case (ty, Constraint.Unconstrained) => (ty, Chain.empty)
        case (ty, Constraint.Never(msg)) =>
          (?, Chain.one(Error.SemanticError(pos, msg)))
        case (ty, Constraint.IsEither(ty1, ty2, msg)) =>
          (ty moreSpecific ty1).orElse(ty moreSpecific ty2).map((_, Chain.empty)).getOrElse {
            (?, Chain.one(Error.TypeMismatch(pos, ty1, ty, msg)))
          }
        case (ty, Constraint.IsUnweakenable(ty2, msg)) =>
          (ty moreSpecific ty2).map((_, Chain.empty)).getOrElse {
            (?, Chain.one(Error.TypeMismatch(pos, ty2, ty, msg)))
          }
      }

    /** Tries to merge two types, returning the more specific one if possible.
      *
      * @param ty2
      *   The other type to merge with.
      * @return
      *   The more specific type if it could be determined, or None if the types are incompatible.
      */
    private infix def moreSpecific(ty2: SemType): Option[SemType] =
      (ty, ty2) match {
        case (ty, ?)                  => Some(ty)
        case (?, ty)                  => Some(ty)
        case (ty1, ty2) if ty1 == ty2 => Some(ty1)
        case (KnownType.Array(inn1), KnownType.Array(inn2)) =>
          (inn1 moreSpecific inn2).map(KnownType.Array(_))
        case (KnownType.Pair(fst1, snd1), KnownType.Pair(fst2, snd2)) =>
          (fst1 moreSpecific fst2, snd1 moreSpecific snd2).mapN(KnownType.Pair(_, _))
        case _ => None
      }
  }

  /** Type-check a function declaration.
    *
    * @param func
    *   The AST of the function to type-check.
    */
  def checkFuncDecl(func: ast.FuncDecl): (microWacc.FuncDecl, Chain[Error]) = {
    val ast.FuncDecl(_, name, params, stmts) = func
    val FuncType(retType, paramTypes) = name.ty.asInstanceOf[FuncType]
    val returnConstraint =
      if func.name.v == renamer.MAIN then Constraint.Never("main body must not return")
      else Constraint.Is(retType, s"function ${name.v} must return $retType")
    val (body, bodyErrors) = stmts.foldMap(checkStmt(_, returnConstraint))
    (
      microWacc.FuncDecl(
        microWacc.Ident(name.v, name.guid)(retType),
        params.zip(paramTypes).map { case (ast.Param(_, ident), ty) =>
          microWacc.Ident(ident.v, ident.guid)(ty)
        },
        body
      ),
      bodyErrors
    )
  }

  /** Type-check an AST statement node.
    *
    * @param stmt
    *   The statement to type-check.
    * @param returnConstraint
    *   The constraint that any `return <expr>` statements must satisfy.
    */
  private def checkStmt(
      stmt: ast.Stmt,
      returnConstraint: Constraint
  ): (Chain[microWacc.Stmt], Chain[Error]) = stmt match {
    // Ignore the type of the variable, since it has been converted to a SemType by the renamer.
    case ast.VarDecl(_, name, value) =>
      val expectedTy = name.ty
      val (typedValue, valueErrors) = checkValue(
        value,
        Constraint.Is(
          expectedTy.asInstanceOf[SemType],
          s"variable ${name.v} must be assigned a value of type $expectedTy"
        )
      )
      (
        Chain.one(
          microWacc.Assign(
            microWacc.Ident(name.v, name.guid)(expectedTy.asInstanceOf[SemType]),
            typedValue
          )
        ),
        valueErrors
      )
    case ast.Assign(lhs, rhs) =>
      val (lhsTyped, lhsErrors) = checkLValue(lhs, Constraint.Unconstrained)
      val (rhsTyped, rhsErrors) =
        checkValue(rhs, Constraint.Is(lhsTyped.ty, s"assignment must have type ${lhsTyped.ty}"))
      val unknownError = (lhsTyped.ty, rhsTyped.ty) match {
        case (?, ?) =>
          Chain.one(
            Error.SemanticError(lhs.pos, "assignment with both sides of unknown type is illegal")
          )
        case _ => Chain.empty
      }
      (Chain.one(microWacc.Assign(lhsTyped, rhsTyped)), lhsErrors ++ rhsErrors ++ unknownError)
    case ast.Read(dest) =>
      val (destTyped, destErrors) = checkLValue(dest, Constraint.Unconstrained)
      val (destTy, destTyErrors) = destTyped.ty match {
        case ? =>
          (
            ?,
            Chain.one(
              Error.SemanticError(dest.pos, "cannot read into a destination with an unknown type")
            )
          )
        case destTy =>
          destTy.satisfies(
            Constraint.IsEither(
              KnownType.Int,
              KnownType.Char,
              "read must be applied to an int or char"
            ),
            dest.pos
          )
      }
      (
        Chain.one(
          microWacc.Assign(
            destTyped,
            microWacc.Call(
              microWacc.Builtin.Read,
              List(
                destTy match {
                  case KnownType.Int      => " %d".toMicroWaccCharArray
                  case KnownType.Char | _ => " %c".toMicroWaccCharArray
                },
                destTyped
              )
            )
          )
        ),
        destErrors ++ destTyErrors
      )
    case ast.Free(lhs) =>
      val (lhsTyped, lhsErrors) = checkValue(
        lhs,
        Constraint.IsEither(
          KnownType.Array(?),
          KnownType.Pair(?, ?),
          "free must be applied to an array or pair"
        )
      )
      (Chain.one(microWacc.Call(microWacc.Builtin.Free, List(lhsTyped))), lhsErrors)
    case ast.Return(expr) =>
      val (exprTyped, exprErrors) = checkValue(expr, returnConstraint)
      (Chain.one(microWacc.Return(exprTyped)), exprErrors)
    case ast.Exit(expr) =>
      val (exprTyped, exprErrors) =
        checkValue(expr, Constraint.Is(KnownType.Int, "exit value must be int"))
      (Chain.one(microWacc.Call(microWacc.Builtin.Exit, List(exprTyped))), exprErrors)
    case ast.Print(expr, newline) =>
      // This constraint should never fail, the scope-checker should have caught it already
      val (exprTyped, exprErrors) = checkValue(expr, Constraint.Unconstrained)
      val exprFormat = exprTyped.ty match {
        case KnownType.Bool | KnownType.String             => "%s"
        case KnownType.Array(KnownType.Char)               => "%.*s"
        case KnownType.Char                                => "%c"
        case KnownType.Int                                 => "%d"
        case KnownType.Pair(_, _) | KnownType.Array(_) | ? => "%p"
      }
      val printfCall = { (func: microWacc.Builtin, value: microWacc.Expr) =>
        Chain.one(
          microWacc.Call(
            func,
            List(
              s"$exprFormat${if newline then "\n" else ""}".toMicroWaccCharArray,
              value
            )
          )
        )
      }
      (
        exprTyped.ty match {
          case KnownType.Bool =>
            Chain.one(
              microWacc.If(
                exprTyped,
                printfCall(microWacc.Builtin.Printf, "true".toMicroWaccCharArray),
                printfCall(microWacc.Builtin.Printf, "false".toMicroWaccCharArray)
              )
            )
          case KnownType.Array(KnownType.Char) =>
            printfCall(microWacc.Builtin.PrintCharArray, exprTyped)
          case _ => printfCall(microWacc.Builtin.Printf, exprTyped)
        },
        exprErrors
      )
    case ast.If(cond, thenStmt, elseStmt) =>
      val (condTyped, condErrors) =
        checkValue(cond, Constraint.Is(KnownType.Bool, "if condition must be a bool"))
      val (thenStmtTyped, thenErrors) = thenStmt.foldMap(checkStmt(_, returnConstraint))
      val (elseStmtTyped, elseErrors) = elseStmt.foldMap(checkStmt(_, returnConstraint))
      (
        Chain.one(microWacc.If(condTyped, thenStmtTyped, elseStmtTyped)),
        condErrors ++ thenErrors ++ elseErrors
      )
    case ast.While(cond, body) =>
      val (condTyped, condErrors) =
        checkValue(cond, Constraint.Is(KnownType.Bool, "while condition must be a bool"))
      val (bodyTyped, bodyErrors) = body.foldMap(checkStmt(_, returnConstraint))
      (Chain.one(microWacc.While(condTyped, bodyTyped)), condErrors ++ bodyErrors)
    case ast.Block(body)   => body.foldMap(checkStmt(_, returnConstraint))
    case skip @ ast.Skip() => (Chain.empty, Chain.empty)
  }

  /** Type-check an AST LValue, RValue or Expr node. This function does all 3 since these traits
    * overlap in the AST.
    *
    * @param value
    *   The value to type-check.
    * @param constraint
    *   The type constraint that the value must satisfy.
    * @return
    *   The most specific type of the value if it could be determined, or ? if it could not.
    */
  private def checkValue(
      value: ast.LValue | ast.RValue | ast.Expr,
      constraint: Constraint
  ): (microWacc.Expr, Chain[Error]) = value match {
    case l @ ast.IntLiter(v) =>
      val (_, errors) = KnownType.Int.satisfies(constraint, l.pos)
      (microWacc.IntLiter(v), errors)
    case l @ ast.BoolLiter(v) =>
      val (_, errors) = KnownType.Bool.satisfies(constraint, l.pos)
      (microWacc.BoolLiter(v), errors)
    case l @ ast.CharLiter(v) =>
      val (_, errors) = KnownType.Char.satisfies(constraint, l.pos)
      (microWacc.CharLiter(v), errors)
    case l @ ast.StrLiter(v) =>
      val (_, errors) = KnownType.String.satisfies(constraint, l.pos)
      (v.toMicroWaccCharArray, errors)
    case l @ ast.PairLiter() =>
      val (ty, errors) = KnownType.Pair(?, ?).satisfies(constraint, l.pos)
      (microWacc.NullLiter()(ty), errors)
    case ast.Parens(expr) => checkValue(expr, constraint)
    case l @ ast.ArrayLiter(elems) =>
      val ((elemTy, elemsErrors), elemsTyped) =
        elems.mapAccumulate[(SemType, Chain[Error]), microWacc.Expr]((?, Chain.empty)) {
          case ((acc, errors), elem) =>
            val (elemTyped, elemErrors) = checkValue(
              elem,
              Constraint.IsSymmetricCompatible(acc, s"array elements must have the same type")
            )
            ((elemTyped.ty, errors ++ elemErrors), elemTyped)
        }
      val (arrayTy, arrayErrors) = KnownType
        // Start with an unknown param type, make it more specific while checking the elements.
        .Array(elemTy)
        .satisfies(constraint, l.pos)
      (microWacc.ArrayLiter(elemsTyped)(arrayTy), elemsErrors ++ arrayErrors)
    case l @ ast.NewPair(fst, snd) =>
      val (fstTyped, fstErrors) = checkValue(fst, Constraint.Unconstrained)
      val (sndTyped, sndErrors) = checkValue(snd, Constraint.Unconstrained)
      val (pairTy, pairErrors) =
        KnownType.Pair(fstTyped.ty, sndTyped.ty).satisfies(constraint, l.pos)
      (microWacc.ArrayLiter(List(fstTyped, sndTyped))(pairTy), fstErrors ++ sndErrors ++ pairErrors)
    case ast.Call(id, args) =>
      val funcTy @ FuncType(retTy, paramTys) = id.ty.asInstanceOf[FuncType]
      val lenError =
        if (args.length == paramTys.length) then Chain.empty
        else Chain.one(Error.FunctionParamsMismatch(id, paramTys.length, args.length, funcTy))
      // Even if the number of arguments is wrong, we still check the types of the arguments
      // in the best way we can (by taking a zip).
      val (argsErrors, argsTyped) =
        args.zip(paramTys).mapAccumulate(Chain.empty[Error]) { case (errors, (arg, paramTy)) =>
          val (argTyped, argErrors) =
            checkValue(arg, Constraint.Is(paramTy, s"argument type mismatch in function ${id.v}"))
          (errors ++ argErrors, argTyped)
        }
      val (retTyChecked, retErrors) = retTy.satisfies(constraint, id.pos)
      (
        microWacc.Call(microWacc.Ident(id.v, id.guid)(retTyChecked), argsTyped),
        lenError ++ argsErrors ++ retErrors
      )

    // Unary operators
    case ast.Negate(x) =>
      val (argTyped, argErrors) =
        checkValue(x, Constraint.Is(KnownType.Int, "negation must be applied to an int"))
      val (retTy, retErrors) = KnownType.Int.satisfies(constraint, x.pos)
      (microWacc.UnaryOp(argTyped, microWacc.UnaryOperator.Negate)(retTy), argErrors ++ retErrors)
    case ast.Not(x) =>
      val (argTyped, argErrors) =
        checkValue(x, Constraint.Is(KnownType.Bool, "logical not must be applied to a bool"))
      val (retTy, retErrors) = KnownType.Bool.satisfies(constraint, x.pos)
      (microWacc.UnaryOp(argTyped, microWacc.UnaryOperator.Not)(retTy), argErrors ++ retErrors)
    case ast.Len(x) =>
      val (argTyped, argErrors) =
        checkValue(x, Constraint.Is(KnownType.Array(?), "len must be applied to an array"))
      val (retTy, retErrors) = KnownType.Int.satisfies(constraint, x.pos)
      (microWacc.UnaryOp(argTyped, microWacc.UnaryOperator.Len)(retTy), argErrors ++ retErrors)
    case ast.Ord(x) =>
      val (argTyped, argErrors) =
        checkValue(x, Constraint.Is(KnownType.Char, "ord must be applied to a char"))
      val (retTy, retErrors) = KnownType.Int.satisfies(constraint, x.pos)
      (microWacc.UnaryOp(argTyped, microWacc.UnaryOperator.Ord)(retTy), argErrors ++ retErrors)
    case ast.Chr(x) =>
      val (argTyped, argErrors) =
        checkValue(x, Constraint.Is(KnownType.Int, "chr must be applied to an int"))
      val (retTy, retErrors) = KnownType.Char.satisfies(constraint, x.pos)
      (microWacc.UnaryOp(argTyped, microWacc.UnaryOperator.Chr)(retTy), argErrors ++ retErrors)

    // Binary operators
    case op: (ast.Add | ast.Sub | ast.Mul | ast.Div | ast.Mod) =>
      val operand = Constraint.Is(KnownType.Int, s"${op.name} operator must be applied to an int")
      val (xTyped, xErrors) = checkValue(op.x, operand)
      val (yTyped, yErrors) = checkValue(op.y, operand)
      val (retTy, retErrors) = KnownType.Int.satisfies(constraint, op.pos)
      (
        microWacc.BinaryOp(xTyped, yTyped, microWacc.BinaryOperator.fromAst(op))(retTy),
        xErrors ++ yErrors ++ retErrors
      )
    case op: (ast.Eq | ast.Neq) =>
      val (xTyped, xErrors) = checkValue(op.x, Constraint.Unconstrained)
      val (yTyped, yErrors) = checkValue(
        op.y,
        Constraint.Is(xTyped.ty, s"${op.name} operator must be applied to values of the same type")
      )
      val (retTy, retErrors) = KnownType.Bool.satisfies(constraint, op.pos)
      (
        microWacc.BinaryOp(xTyped, yTyped, microWacc.BinaryOperator.fromAst(op))(retTy),
        xErrors ++ yErrors ++ retErrors
      )
    case op: (ast.Less | ast.LessEq | ast.Greater | ast.GreaterEq) =>
      val xConstraint = Constraint.IsEither(
        KnownType.Int,
        KnownType.Char,
        s"${op.name} operator must be applied to an int or char"
      )
      val (xTyped, xErrors) = checkValue(op.x, xConstraint)
      // If x type-check failed, we still want to check y is an Int or Char (rather than ?)
      val yConstraint = xTyped.ty match {
        case ? => xConstraint
        case xTy =>
          Constraint.Is(xTy, s"${op.name} operator must be applied to values of the same type")
      }
      val (yTyped, yErrors) = checkValue(op.y, yConstraint)
      val (retTy, retErrors) = KnownType.Bool.satisfies(constraint, op.pos)
      (
        microWacc.BinaryOp(xTyped, yTyped, microWacc.BinaryOperator.fromAst(op))(retTy),
        xErrors ++ yErrors ++ retErrors
      )
    case op: (ast.And | ast.Or) =>
      val operand = Constraint.Is(KnownType.Bool, s"${op.name} operator must be applied to a bool")
      val (xTyped, xErrors) = checkValue(op.x, operand)
      val (yTyped, yErrors) = checkValue(op.y, operand)
      val (retTy, retErrors) = KnownType.Bool.satisfies(constraint, op.pos)
      (
        microWacc.BinaryOp(xTyped, yTyped, microWacc.BinaryOperator.fromAst(op))(retTy),
        xErrors ++ yErrors ++ retErrors
      )

    case lvalue: ast.LValue => checkLValue(lvalue, constraint)
  }

  /** Type-check an AST LValue node. Separate because microWacc keeps LValues
    *
    * @param value
    *   The value to type-check.
    * @param constraint
    *   The type constraint that the value must satisfy.
    * @param ctx
    *   The type checker context which includes the global names and functions, and an errors
    *   builder.
    * @return
    *   The most specific type of the value if it could be determined, or ? if it could not.
    */
  private def checkLValue(
      value: ast.LValue,
      constraint: Constraint
  ): (microWacc.LValue, Chain[Error]) = value match {
    case id @ ast.Ident(name, guid, ty) =>
      val (idTy, idErrors) = ty.asInstanceOf[SemType].satisfies(constraint, id.pos)
      (microWacc.Ident(name, guid)(idTy), idErrors)
    case ast.ArrayElem(id, indices) =>
      val arrayTy = id.ty.asInstanceOf[SemType]
      val ((elemTy, elemErrors), indicesTyped) =
        indices.mapAccumulate((arrayTy.asInstanceOf[SemType], Chain.empty[Error])) {
          case ((acc, errors), elem) =>
            val (idxTyped, idxErrors) =
              checkValue(elem, Constraint.Is(KnownType.Int, "array index must be an int"))
            val (next, nextError) = acc match {
              case KnownType.Array(innerTy) => (innerTy, Chain.empty)
              case ? => (?, Chain.empty) // we can keep indexing an unknown type
              case nonArrayTy =>
                (
                  ?,
                  Chain.one(
                    Error.TypeMismatch(
                      elem.pos,
                      KnownType.Array(?),
                      acc,
                      "cannot index into a non-array"
                    )
                  )
                )
            }
            ((next, errors ++ idxErrors ++ nextError), idxTyped)
        }
      val (retTy, retErrors) = elemTy.satisfies(constraint, value.pos)
      val firstArrayElem = microWacc.ArrayElem(
        microWacc.Ident(id.v, id.guid)(arrayTy),
        indicesTyped.head
      )(retTy)
      val arrayElem = indicesTyped.tail.foldLeft(firstArrayElem) { (acc, idx) =>
        microWacc.ArrayElem(acc, idx)(KnownType.Array(acc.ty))
      }
      // Need to type-check the final arrayElem with the constraint
      // TODO: What
      (microWacc.ArrayElem(arrayElem.value, arrayElem.index)(retTy), elemErrors ++ retErrors)
    case ast.Fst(elem) =>
      val (elemTyped, elemErrors) = checkLValue(
        elem,
        Constraint.Is(KnownType.Pair(?, ?), "fst must be applied to a pair")
      )
      val (retTy, retErrors) = elemTyped.ty match {
        case KnownType.Pair(left, _) => left.satisfies(constraint, elem.pos)
        case _ => (?, Chain.one(Error.InternalError(elem.pos, "fst must be applied to a pair")))
      }
      (microWacc.ArrayElem(elemTyped, microWacc.IntLiter(0))(retTy), elemErrors ++ retErrors)
    case ast.Snd(elem) =>
      val (elemTyped, elemErrors) = checkLValue(
        elem,
        Constraint.Is(KnownType.Pair(?, ?), "snd must be applied to a pair")
      )
      val (retTy, retErrors) = elemTyped.ty match {
        case KnownType.Pair(_, right) => right.satisfies(constraint, elem.pos)
        case _ => (?, Chain.one(Error.InternalError(elem.pos, "snd must be applied to a pair")))
      }
      (microWacc.ArrayElem(elemTyped, microWacc.IntLiter(1))(retTy), elemErrors ++ retErrors)
  }

  extension (s: String) {
    def toMicroWaccCharArray: microWacc.ArrayLiter =
      microWacc.ArrayLiter(s.map(microWacc.CharLiter(_)).toList)(KnownType.String)
  }
}
