package wacc

import cats.syntax.all._
import scala.collection.mutable
import cats.data.NonEmptyList

object typeChecker {
  import wacc.types._

  case class TypeCheckerCtx(
      errors: mutable.Builder[Error, List[Error]]
  ) {
    def error(err: Error): SemType =
      errors += err
      ?
  }

  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)(using
        ctx: TypeCheckerCtx
    ): SemType =
      (ty, constraint) match {
        case (KnownType.Array(KnownType.Char), Constraint.Is(KnownType.String, _)) =>
          KnownType.String
        case (
              KnownType.String,
              Constraint.IsSymmetricCompatible(KnownType.Array(KnownType.Char), _)
            ) =>
          KnownType.String
        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
        case (ty, Constraint.Never(msg)) =>
          ctx.error(Error.SemanticError(pos, msg))
        case (ty, Constraint.IsEither(ty1, ty2, msg)) =>
          (ty moreSpecific ty1).orElse(ty moreSpecific ty2).getOrElse {
            ctx.error(Error.TypeMismatch(pos, ty1, ty, msg))
          }
        case (ty, Constraint.IsUnweakenable(ty2, msg)) =>
          (ty moreSpecific ty2).getOrElse {
            ctx.error(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 WACC program.
    *
    * @param prog
    *   The AST of the program to type-check.
    * @param ctx
    *   The type checker context which includes the global names and functions, and an errors
    *   builder.
    */
  def check(prog: ast.Program, errors: mutable.Builder[Error, List[Error]]): microWacc.Program =
    given ctx: TypeCheckerCtx = TypeCheckerCtx(errors)
    microWacc.Program(
      // Ignore function syntax types for return value and params, since those have been converted
      // to SemTypes by the renamer.
      prog.funcs.map { case ast.FuncDecl(_, name, params, stmts) =>
        val FuncType(retType, paramTypes) = name.ty.asInstanceOf[FuncType]
        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)
          },
          stmts.toList
            .flatMap(
              checkStmt(_, Constraint.Is(retType, s"function ${name.v} must return $retType"))
            )
        )
      },
      prog.main.toList
        .flatMap(checkStmt(_, Constraint.Never("main function must not return")))
    )

  /** 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)(using
      ctx: TypeCheckerCtx
  ): List[microWacc.Stmt] = 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 = checkValue(
        value,
        Constraint.Is(
          expectedTy.asInstanceOf[SemType],
          s"variable ${name.v} must be assigned a value of type $expectedTy"
        )
      )
      List(
        microWacc.Assign(
          microWacc.Ident(name.v, name.guid)(expectedTy.asInstanceOf[SemType]),
          typedValue
        )
      )
    case ast.Assign(lhs, rhs) =>
      val lhsTyped = checkLValue(lhs, Constraint.Unconstrained)
      val rhsTyped =
        checkValue(rhs, Constraint.Is(lhsTyped.ty, s"assignment must have type ${lhsTyped.ty}"))
      (lhsTyped.ty, rhsTyped.ty) match {
        case (?, ?) =>
          ctx.error(
            Error.SemanticError(lhs.pos, "assignment with both sides of unknown type is illegal")
          )
        case _ => ()
      }
      List(microWacc.Assign(lhsTyped, rhsTyped))
    case ast.Read(dest) =>
      val destTyped = checkLValue(dest, Constraint.Unconstrained)
      val destTy = destTyped.ty match {
        case ? =>
          ctx.error(
            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
          )
      }
      List(
        microWacc.Assign(
          destTyped,
          microWacc.Call(
            microWacc.Builtin.Read,
            List(
              destTy match {
                case KnownType.Int      => " %d".toMicroWaccCharArray
                case KnownType.Char | _ => " %c".toMicroWaccCharArray
              },
              destTyped
            )
          )
        )
      )
    case ast.Free(lhs) =>
      List(
        microWacc.Call(
          microWacc.Builtin.Free,
          List(
            checkValue(
              lhs,
              Constraint.IsEither(
                KnownType.Array(?),
                KnownType.Pair(?, ?),
                "free must be applied to an array or pair"
              )
            )
          )
        )
      )
    case ast.Return(expr) =>
      List(microWacc.Return(checkValue(expr, returnConstraint)))
    case ast.Exit(expr) =>
      List(
        microWacc.Call(
          microWacc.Builtin.Exit,
          List(checkValue(expr, Constraint.Is(KnownType.Int, "exit value must be int")))
        )
      )
    case ast.Print(expr, newline) =>
      // This constraint should never fail, the scope-checker should have caught it already
      val exprTyped = 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) =>
        List(
          microWacc.Call(
            func,
            List(
              s"$exprFormat${if newline then "\n" else ""}".toMicroWaccCharArray,
              value
            )
          )
        )
      }
      exprTyped.ty match {
        case KnownType.Bool =>
          List(
            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)
      }
    case ast.If(cond, thenStmt, elseStmt) =>
      List(
        microWacc.If(
          checkValue(cond, Constraint.Is(KnownType.Bool, "if condition must be a bool")),
          thenStmt.toList.flatMap(checkStmt(_, returnConstraint)),
          elseStmt.toList.flatMap(checkStmt(_, returnConstraint))
        )
      )
    case ast.While(cond, body) =>
      List(
        microWacc.While(
          checkValue(cond, Constraint.Is(KnownType.Bool, "while condition must be a bool")),
          body.toList.flatMap(checkStmt(_, returnConstraint))
        )
      )
    case ast.Block(body)   => body.toList.flatMap(checkStmt(_, returnConstraint))
    case skip @ ast.Skip() => List.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)(using
      ctx: TypeCheckerCtx
  ): microWacc.Expr = value match {
    case l @ ast.IntLiter(v) =>
      KnownType.Int.satisfies(constraint, l.pos)
      microWacc.IntLiter(v)
    case l @ ast.BoolLiter(v) =>
      KnownType.Bool.satisfies(constraint, l.pos)
      microWacc.BoolLiter(v)
    case l @ ast.CharLiter(v) =>
      KnownType.Char.satisfies(constraint, l.pos)
      microWacc.CharLiter(v)
    case l @ ast.StrLiter(v) =>
      KnownType.String.satisfies(constraint, l.pos)
      v.toMicroWaccCharArray
    case l @ ast.PairLiter() =>
      microWacc.NullLiter()(KnownType.Pair(?, ?).satisfies(constraint, l.pos))
    case ast.Parens(expr) => checkValue(expr, constraint)
    case l @ ast.ArrayLiter(elems) =>
      val (elemTy, elemsTyped) = elems.mapAccumulate[SemType, microWacc.Expr](?) {
        case (acc, elem) =>
          val elemTyped = checkValue(
            elem,
            Constraint.IsSymmetricCompatible(acc, s"array elements must have the same type")
          )
          (elemTyped.ty, elemTyped)
      }
      val arrayTy = 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)
    case l @ ast.NewPair(fst, snd) =>
      val fstTyped = checkValue(fst, Constraint.Unconstrained)
      val sndTyped = checkValue(snd, Constraint.Unconstrained)
      microWacc.ArrayLiter(List(fstTyped, sndTyped))(
        KnownType.Pair(fstTyped.ty, sndTyped.ty).satisfies(constraint, l.pos)
      )
    case ast.Call(id, args) =>
      val funcTy @ FuncType(retTy, paramTys) = id.ty.asInstanceOf[FuncType]
      if (args.length != paramTys.length) {
        ctx.error(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 argsTyped = args.zip(paramTys).map { case (arg, paramTy) =>
        checkValue(arg, Constraint.Is(paramTy, s"argument type mismatch in function ${id.v}"))
      }
      microWacc.Call(microWacc.Ident(id.v, id.guid)(retTy.satisfies(constraint, id.pos)), argsTyped)

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

    // 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")
      microWacc.BinaryOp(
        checkValue(op.x, operand),
        checkValue(op.y, operand),
        microWacc.BinaryOperator.fromAst(op)
      )(KnownType.Int.satisfies(constraint, op.pos))
    case op: (ast.Eq | ast.Neq) =>
      val xTyped = checkValue(op.x, Constraint.Unconstrained)
      microWacc.BinaryOp(
        xTyped,
        checkValue(
          op.y,
          Constraint
            .Is(xTyped.ty, s"${op.name} operator must be applied to values of the same type")
        ),
        microWacc.BinaryOperator.fromAst(op)
      )(KnownType.Bool.satisfies(constraint, op.pos))
    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 = 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")
      }
      microWacc.BinaryOp(
        xTyped,
        checkValue(op.y, yConstraint),
        microWacc.BinaryOperator.fromAst(op)
      )(KnownType.Bool.satisfies(constraint, op.pos))
    case op: (ast.And | ast.Or) =>
      val operand = Constraint.Is(KnownType.Bool, s"${op.name} operator must be applied to a bool")
      microWacc.BinaryOp(
        checkValue(op.x, operand),
        checkValue(op.y, operand),
        microWacc.BinaryOperator.fromAst(op)
      )(KnownType.Bool.satisfies(constraint, op.pos))

    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)(using
      ctx: TypeCheckerCtx
  ): microWacc.LValue = value match {
    case id @ ast.Ident(name, guid, ty) =>
      microWacc.Ident(name, guid)(ty.asInstanceOf[SemType].satisfies(constraint, id.pos))
    case ast.ArrayElem(id, indices) =>
      val arrayTy = id.ty.asInstanceOf[SemType]
      val (elemTy, indicesTyped) = indices.mapAccumulate(arrayTy.asInstanceOf[SemType]) {
        (acc, elem) =>
          val idxTyped =
            checkValue(elem, Constraint.Is(KnownType.Int, "array index must be an int"))
          val next = acc match {
            case KnownType.Array(innerTy) => innerTy
            case ?                        => ? // we can keep indexing an unknown type
            case nonArrayTy =>
              ctx.error(
                Error.TypeMismatch(
                  elem.pos,
                  KnownType.Array(?),
                  acc,
                  "cannot index into a non-array"
                )
              )
              ?
          }
          (next, idxTyped)
      }
      val firstArrayElem = microWacc.ArrayElem(
        microWacc.Ident(id.v, id.guid)(arrayTy),
        indicesTyped.head
      )(elemTy.satisfies(constraint, value.pos))
      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
      microWacc.ArrayElem(arrayElem.value, arrayElem.index)(elemTy.satisfies(constraint, value.pos))
    case ast.Fst(elem) =>
      val elemTyped = checkLValue(
        elem,
        Constraint.Is(KnownType.Pair(?, ?), "fst must be applied to a pair")
      )
      microWacc.ArrayElem(
        elemTyped,
        microWacc.IntLiter(0)
      )(elemTyped.ty match {
        case KnownType.Pair(left, _) =>
          left.satisfies(constraint, elem.pos)
        case _ => ctx.error(Error.InternalError(elem.pos, "fst must be applied to a pair"))
      })
    case ast.Snd(elem) =>
      val elemTyped = checkLValue(
        elem,
        Constraint.Is(KnownType.Pair(?, ?), "snd must be applied to a pair")
      )
      microWacc.ArrayElem(
        elemTyped,
        microWacc.IntLiter(1)
      )(elemTyped.ty match {
        case KnownType.Pair(_, right) =>
          right.satisfies(constraint, elem.pos)
        case _ => ctx.error(Error.InternalError(elem.pos, "snd must be applied to a pair"))
      })
  }

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