package analysis

import com.microsoft.z3._
import com.typesafe.scalalogging.StrictLogging
import constraints.Z3Constraint
import dex.{DexApp, DexClass, DexMethod}
import domain.AbstractDomainsCollection
import org.jf.dexlib2.Opcode
import org.jf.dexlib2.iface.instruction.Instruction
import org.jf.dexlib2.iface.instruction.formats._

import scala.collection.mutable.ArrayBuffer

class Analyzer(
    ctx: Context,
    domains: AbstractDomainsCollection,
    app: DexApp,
    preAnalysisResult: PreAnalysisResult
) extends StrictLogging {

  /**
    * Methods to analyze. This collection is populated while analyzing each instruction.
    */
  private var methodsToAnalyze: ArrayBuffer[(DexClass, DexMethod)] =
    ArrayBuffer[(DexClass, DexMethod)]()

  /**
    * Analyze the program.
    * @return the z3 constraints.
    */
  def analyze: Seq[Z3Constraint] = {
    /*
      Add all the entry points (from both manifest and callbacks)
      to the methods to analyze.
     */
    val entryPoints = app.manifest.entryPoints ++ app.entryPoints
    methodsToAnalyze = ArrayBuffer(entryPoints: _*)
    val env = new Environment(ctx, domains, app, preAnalysisResult)
    analyzeAux(env, 0, List())
  }

  /**
    * Analyze the methods, starting from the given index, until there are no more methods to analyze.
    */
  def analyzeAux(env: Environment,
                 index: Int,
                 constraints: Seq[Z3Constraint]): Seq[Z3Constraint] = {
    if (index >= methodsToAnalyze.length) {
      // No more methods to analyze.
      return constraints
    }
    val (c, m) = methodsToAnalyze(index)
    logger.debug(s"Analyzing method '${m.name}' of class '${c.name}'")
    val newConstraints: Seq[Z3Constraint] =
      m.instructions.zipWithIndex.flatMap {
        case (i: Instruction, pc: Int) =>
          logger.info(s"$pc - ${i.getOpcode}")
          val pp = PP(c, m, pc)
          i.getOpcode match {
            case Opcode.CONST_4 =>
              InstructionAnalyzer.const4(
                i.asInstanceOf[Instruction11n],
                pp,
                env
              )
            case Opcode.CONST_STRING =>
              InstructionAnalyzer.constString(
                i.asInstanceOf[Instruction21c],
                pp,
                env
              )
            case Opcode.CONST_STRING_JUMBO =>
              InstructionAnalyzer.constStringJumbo(
                i.asInstanceOf[Instruction31c],
                pp,
                env
              )
            case Opcode.IPUT =>
              InstructionAnalyzer.iput(
                i.asInstanceOf[Instruction22c],
                pp,
                env
              )
            case Opcode.SGET_OBJECT =>
              InstructionAnalyzer.sgetObject(
                i.asInstanceOf[Instruction21c],
                pp,
                env
              )
            case Opcode.SPUT_OBJECT =>
              InstructionAnalyzer.sputObject(
                i.asInstanceOf[Instruction21c],
                pp,
                env
              )
            case Opcode.MOVE_OBJECT =>
              InstructionAnalyzer.move(
                i.asInstanceOf[Instruction12x],
                pp,
                env
              )
            case Opcode.MOVE_RESULT | Opcode.MOVE_RESULT_OBJECT =>
              InstructionAnalyzer.moveResult(
                i.asInstanceOf[Instruction11x],
                pp,
                env
              )
            case Opcode.IF_EQ =>
              InstructionAnalyzer.ifEq(i.asInstanceOf[Instruction22t], pp, env)
            case Opcode.IF_NE =>
              InstructionAnalyzer.ifEq(i.asInstanceOf[Instruction22t], pp, env)
            case Opcode.GOTO =>
              InstructionAnalyzer.goto(i.asInstanceOf[Instruction10t], pp, env)
            case Opcode.NEW_INSTANCE =>
              InstructionAnalyzer.newInstance(
                i.asInstanceOf[Instruction21c],
                pp,
                env
              )
            case Opcode.INVOKE_DIRECT =>
              InstructionAnalyzer.invokeDirect(
                i.asInstanceOf[Instruction35c],
                pp,
                env,
                (c, m) => this.addMethodToAnalyze(c, m)
              )
            case Opcode.INVOKE_VIRTUAL =>
              InstructionAnalyzer.invokeVirtual(
                i.asInstanceOf[Instruction35c],
                pp,
                env,
                (c, m) => this.addMethodToAnalyze(c, m)
              )
            case Opcode.RETURN_VOID =>
              InstructionAnalyzer.retVoid(
                i.asInstanceOf[Instruction10x],
                pp,
                env
              )
            case Opcode.RETURN =>
              InstructionAnalyzer.ret(i.asInstanceOf[Instruction11x], pp, env)
            case Opcode.CHECK_CAST | Opcode.MOVE_EXCEPTION =>
              InstructionAnalyzer.nop(i, pp, env)
            case _ =>
              logger.warn(s"Not implemented: ${i.getOpcode}, changed to NOP")
              InstructionAnalyzer.nop(i, pp, env)
          }
      }
    analyzeAux(env, index + 1, constraints ++ newConstraints)
  }

  private def addMethodToAnalyze(dexClass: DexClass,
                                 dexMethod: DexMethod): Unit =
    if (!methodsToAnalyze.contains((dexClass, dexMethod))) {
      methodsToAnalyze += ((dexClass, dexMethod))
    }

}
