package ixee.re.disasm

import ixee.cryptopals.utils.ByteUtils._
import ixee.cryptopals.utils.FunctionUtils._

import scala.collection.mutable.{ArrayBuffer, Map}

sealed trait Instruction {
  def isDone: Boolean
  def giveWord(word: Int): Unit
}

case class NoOpInst(code: NoOpCode, offset: Int) extends Instruction {
  override def toString: String = {
    s"$code -> $offset"
  }
  def giveWord(word: Int) = throw new RuntimeException("Too many bytes!")
  def isDone = true
}
case class OneOpInst(code: OneOpCode, byte: Boolean, dest: Int, addressMode: OneOpAddressing) extends Instruction {
  private var extraWord: Option[Int] = None
  def neededWords = Seq(addressMode).count(_.wide)
  def giveWord(word: Int) =
    extraWord.fold(extraWord = Some(word))(_ => throw new RuntimeException("Too many bytes!"))
  def isDone = neededWords == Seq(extraWord).flatten.length

  override def toString: String = {
    val byteStr = if (byte) ".B" else "  "
    s"$code$byteStr ${addressMode match {
      case mode: NormMode => mode.stringifyWith(dest)
      case mode: WideMode => mode.stringifyWith(dest, extraWord.get)
    }}"
  }
}
case class TwoOpInst(code: TwoOpCode, byte: Boolean, src: Int, srcMode: TwoOpSourceAddressing, dest: Int, dstMode: TwoOpDestAddressing) extends Instruction {
  private var extraWord1: Option[Int] = None
  private var extraWord2: Option[Int] = None
  def neededWords = Seq(srcMode, dstMode).count(_.wide)
  def giveWord(word: Int) =
    extraWord1.fold(extraWord1 = Some(word))(_ =>
      extraWord2.fold(extraWord2 = Some(word))(_ => throw new RuntimeException("Too many bytes!"))
    )
  def isDone = neededWords == Seq(extraWord1, extraWord2).flatten.length

  override def toString: String = {
    val byteStr = if (byte) ".B" else "  "
    s"$code$byteStr ${srcMode match {
      case mode: NormMode => mode.stringifyWith(src)
      case mode: WideMode => mode.stringifyWith(src, extraWord1.get)
    }}, ${dstMode match {
      case mode: NormMode => mode.stringifyWith(dest)
      case mode: WideMode => mode.stringifyWith(dest, extraWord2.orElse(if (srcMode.wide) None else extraWord1).get)
    }}"
  }
}

object NoOp {
  def info(bytes: (Byte, Byte)): (Int, Int) = {
    val op = (bytes._1 & 0xfc) >> 2
    val offset = ((bytes._1 & 0x03) << 8) | bytes._2
    (op, offset)
  }

  def apply(bytes: (Byte, Byte)): Option[NoOpInst] = {
    val (opVal, offset) = info(bytes)
    for {
      opcode <- Opcodes.noOp(opVal << 10)
    } yield NoOpInst(opcode, offset)
  }
}

object OneOp {
  def info(bytes: (Byte, Byte)): (Int, Int, Int, Int) = {
    val dest = (bytes._2 & 0x0f)
    val destMode = ((bytes._2 & 0x30) >> 4)
    val byteBit = ((bytes._2 & 0x40) >> 6)
    val opcode = ((bytes._1 & 0xff) << 1) | ((bytes._2 & 0x80) >> 7)
    (opcode, byteBit, dest, destMode)
  }

  def apply(bytes: (Byte, Byte)): Option[OneOpInst] = {
    val (opVal, byteInstr, dest, rawDestMode) = info(bytes)
    for {
      opcode <- Opcodes.oneOp(opVal << 7)
      isByteOp = if (byteInstr == 1) true else false
      destMode <- OneOpAddressingModes(dest, rawDestMode)
    } yield OneOpInst(opcode, isByteOp, dest, destMode)
  }
}

object TwoOp {
  def info(bytes: (Byte, Byte)): (Int, Int, Int, Int, Int, Int) = {
    val dest = (bytes._2 & 0x0f)
    val sourceMode = (bytes._2 & 0x30) >> 4
    val byteBit = (bytes._2 & 0x40) >> 6
    val destMode = (bytes._2 & 0x80) >> 7
    val source = (bytes._1 & 0x0f)
    val opcode = (bytes._1 & 0xf0) >> 4
    (opcode, byteBit, source, destMode, dest, sourceMode)
  }

  def apply(bytes: (Byte, Byte)): Option[TwoOpInst] = {
    val (opVal, byteInstr, source, rawDestMode, dest, rawSourceMode) = info(bytes)

    for {
      opcode <- Opcodes.twoOp(opVal << 12)
      isByteOp = if (byteInstr == 1) true else false
      destMode <- TwoOpDestAddressingModes(dest, rawDestMode)
      sourceMode <- TwoOpSourceAddressingModes(source, rawSourceMode)
    } yield TwoOpInst(opcode, isByteOp, source, sourceMode, dest, destMode)
  }
}

object Instruction {
  def apply(instruction: (Byte, Byte)): Option[Instruction] = NoOp(instruction) orElse OneOp(instruction) orElse TwoOp(instruction)
  def littleEndian(instruction: String): Option[Instruction] = {
    val instr = Integer.parseInt(instruction, 16)
    val hi = (instr & 0xff00) >> 8
    val lo = instr & 0x00ff
    apply((hi.toByte, lo.toByte))
  }
  def apply(instruction: String): Option[Instruction] = {
    val instr = Integer.parseInt(instruction, 16)
    val hi = (instr & 0xff00) >> 8
    val lo = instr & 0x00ff
    apply((lo.toByte, hi.toByte))
  }

  def noOp(instruction: (Byte, Byte)): Option[NoOpInst] = NoOp.apply(instruction)

  def oneOp(instruction: (Byte, Byte)): Option[OneOpInst] = OneOp.apply(instruction)

  def twoOp(instruction: (Byte, Byte)): Option[TwoOpInst] = TwoOp.apply(instruction)

  def isNoOp(instruction: (Byte, Byte)) = noOp(instruction).isDefined

  def isOneOp(instruction: (Byte, Byte)) = oneOp(instruction).isDefined

  def isTwoOp(instruction: (Byte, Byte)) = twoOp(instruction).isDefined
}

class InstructionParser(bytes: Seq[Byte]) {
  val instructions = scala.collection.mutable.ArrayBuffer[Instruction]()
  val bytePair = scala.collection.mutable.ArrayBuffer[Byte]()

  for (byte <- bytes) {
    bytePair match {
      case ArrayBuffer() => bytePair += byte
      case ArrayBuffer(first) => {
        val second = byte
        def parseInstr = {
          instructions += Instruction((second, first)).getOrElse { throw new RuntimeException(s"Parse failure! Bytes = $first, $second") }
          bytePair.clear
        }

        val lastIsDone = instructions.lastOption.map(_.isDone).getOrElse(true)

        instructions.lastOption.map { instr =>
          if (instr.isDone) {
            parseInstr
          } else {
            instr.giveWord(second * 256 + first)
            bytePair.clear
          }
        } getOrElse { parseInstr }
      }
      case buffer @ _ => throw new RuntimeException("Byte array overflow: " + buffer)
    }
  }

  instructions.lastOption match {
    case Some(instr) if !instr.isDone => throw new RuntimeException("Unable to parse instruction stream, last instruction is missing data.")
    case _ => Unit
  }
}

object InstructionParser {
  def apply(bytes: Seq[String]) = new InstructionParser(bytes.map(java.lang.Integer.parseInt(_, 16).toByte))
}