import scala.collection.mutable.ArraySeq

object StateOOInterp {
  sealed abstract class Expr

  case class Id(id: Symbol) extends Expr
  case class Num(n: Float) extends Expr
  case class Add(lhs: Expr, rhs: Expr) extends Expr
  case class Mult(lhs: Expr, rhs: Expr) extends Expr
  case class Let(id: Symbol, namedExpr: Expr, body: Expr) extends Expr

  case class New(className: Symbol, fields: Expr*) extends Expr
  case class GetField(objExpr: Expr, fieldName: Symbol) extends Expr
  case class Call(objExpr: Expr, methodName: Symbol, args: Expr*)
    extends Expr

  case class SetField(objExpr: Expr, fieldName: Symbol, e: Expr) extends Expr
  case class Set(id: Symbol, e: Expr) extends Expr
  case class Seqn(exprs: Expr*) extends Expr

  case class Class(
    superClass: Symbol,
    fields: Seq[Symbol],
    methods: Map[Symbol, Method])

  case class Method(body: Expr, params: Symbol*)

  sealed abstract class Value
  // Note how Object now holds Locations instead of Values
  case class Object(className: Symbol, fields: Seq[Location]) extends Value
  case class NumV(n: Float) extends Value

  type Env = Map[Symbol, Location]
  type Location = Int

  class Store(size: Int) extends scala.collection.mutable.Map[Location, Value] {
    val memory = new ArraySeq[Value](size)
    private var _nextFreeLocation = 0
    private var _free = size

    override def +=(kv: (Location, Value)) = {
      val (k, v) = kv
      val oldVal = memory(k)
      memory(k) = v
      if (oldVal == null) _free -= 1
      this
    }

    override def -=(k: Location) = {
      val oldVal = memory(k)
      memory(k) = null
      if (oldVal != null) _free += 1
      this
    }

    override def get(k: Location) = {
      val v = memory(k)
      if (v == null) None else Some(v)
    }

    override def iterator = new Iterator[(Location, Value)] {
      private var _i = 0
      private def findNext(l: Location) = {
        var j = l
        while (j < memory.size && memory(j) == null) j += 1
        j
      }
      override def hasNext = findNext(_i) < memory.size
      override def next = {
        _i = findNext(_i)
        val k = _i
        val v = memory(_i)
        _i += 1
        (k, v)
      }
    }

    def malloc(v: Value): Location = {
      val loc = findNextFreeLocation()
      this += loc -> v
      loc
    }

    private def findNextFreeLocation(): Location = {
      if (_free <= 0) sys.error("Store is full")
      while (memory(_nextFreeLocation) != null) {
        _nextFreeLocation += 1
        if (_nextFreeLocation >= size) _nextFreeLocation = 0
      }
      _nextFreeLocation
    }
  }

  /**
   * Returns all fields in the path from the root to className in the
   * inheritance tree
   */
  def fieldsInPath(
    className: Symbol,
    classes: Map[Symbol, Class]): List[Symbol] = className match {
    case 'Object => List()
    case _ => {
      val clazz = classes.getOrElse(
        className,
        sys.error("Unknown class %s".format(className)))
      fieldsInPath(clazz.superClass, classes) ++ clazz.fields
    }
  }

  /**
   * Returns the first method found in the path from className to the root in
   * the inheritance tree or None
   */
  def lookupMethod(
    methodName: Symbol,
    className: Symbol,
    classes: Map[Symbol, Class]): Option[Method] = className match {
    case 'Object => None
    case _ => {
      val clazz = classes.getOrElse(
        className,
        sys.error("Unknown class %s".format(className)))
      if (clazz.methods.contains(methodName)) Some(clazz.methods(methodName))
      else lookupMethod(methodName, clazz.superClass, classes)
    }
  }

  def interp(
    e: Expr,
    env: Env,
    store: Store,
    classes: Map[Symbol, Class]): Location = e match {
    case Num(n) => store.malloc(NumV(n))

    case Add(lhs, rhs) => {
      val lhsV = store(interp(lhs, env, store, classes))
      val rhsV = store(interp(rhs, env, store, classes))
      (lhsV, rhsV) match {
        case (NumV(a), NumV(b)) => store.malloc(NumV(a + b))
        case _ => sys.error("Can only add numbers")
      }
    }

    case Mult(lhs, rhs) => {
      val lhsV = store(interp(lhs, env, store, classes))
      val rhsV = store(interp(rhs, env, store, classes))
      (lhsV, rhsV) match {
        case (NumV(a), NumV(b)) => store.malloc(NumV(a * b))
        case _ => sys.error("Can only multiply numbers")
      }
    }

    case Let(id, namedExpr, body) => {
      val loc = interp(namedExpr, env, store, classes)
      interp(body, env + (id -> loc), store, classes)
    }

    case Seqn() => sys.error("Need at least one expression for Seqn")
    case Seqn(exprs @ _*) => exprs.foldLeft[Location](-1) {
      case (_, e) => interp(e, env, store, classes)
    }

    case Id(id) => {
      env.get(id) match {
        case Some(loc) => loc
        case None => sys.error("%s is not found".format(id))
      }
    }

    case New(className, args @ _*) => {
      if(!classes.contains(className)){
        sys.error("class not found")
      }
      store.malloc(Object(className, args.map(expr => interp(expr, env, store, classes))))
    }

    case GetField(objExpr, fieldName) => {
      val loc = interp(objExpr, env, store, classes)
      store.get(loc) match {
        case Some(x) => x match{
          case Object(className, fields) => {
            val fieldIndex = fieldsInPath(className, classes) lastIndexOf fieldName
            if(fieldIndex == -1 ){
              sys.error("the field %s does not exist".format(fieldName))
            }
            fields(fieldIndex)
          }
          case v => sys.error("can only access field of object but got: "+v)
        }
        case None => sys.error("location is empty")
      }
    }

    case Call(objExpr, methodName, args @ _*) => {
      val loc = interp(objExpr, env, store, classes)
      store.get(loc) match {
        case Some(x) => x match{
          case Object(className, fields) => {
            val method = lookupMethod(methodName, className, classes) getOrElse
              sys.error("methodname is unknown")
            val evaledArgs = args map (expr => interp(expr, env, store, classes))
            val argBindings = method.params zip evaledArgs
            val fieldBindings = fieldsInPath(className,classes) zip fields
            val thisBinding = 'this -> loc
            interp(method.body, env ++ fieldBindings ++ argBindings + thisBinding, store, classes)
          }
          case v => sys.error("can only access method of object but got: "+v)
        }
        case None => sys.error("location is empty")
      }
    }

    case SetField(objExpr, fieldName, e) => {
      val loc = interp(objExpr, env, store, classes)
      store.get(loc) match {
        case Some(x) => x match {
          case Object(className, fields) => {
            val fieldIndex = fieldsInPath(className, classes) indexOf fieldName
            if(fieldIndex == -1 ){
              sys.error("the field %s does not exist".format(fieldName))
            }
            val eLoc = interp(e, env, store, classes)
            //store +=(fields(fieldIndex), store.get(eLoc).get)
            store(fields(fieldIndex)) = store.get(eLoc).get
            store -= eLoc
            fields(fieldIndex)
          }
          case v => sys.error("can only set field of object but got: " + v)
        }
        case None => sys.error("location is empty")
      }
    }

    case Set(id, e) => {
      val eLoc = interp(e, env, store, classes)
      env get id match {
        case Some(loc) => {
          store(loc) = store.get(eLoc).get
          store -= eLoc
          loc
        }
        case None => sys.error("id %s not found".format(id))
      }
    }
  }
}

//