package ixee.re.disasm

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

import scala.collection.mutable.Map

sealed trait AddressingMode {
  def name: String
  def wide: Boolean
  override def toString: String = name
  def nameifyRegister(reg: Int): String = reg match {
    case 0 => "pc"
    case 1 => "sp"
    case 2 => "sr"
    case 3 => "cg"
    case num @ _ => s"r$num"
  }
}

sealed trait NormMode extends AddressingMode {
  def wide = false
  def stringifyWith(reg: Int): String
}

sealed trait WideMode extends AddressingMode {
  def wide = true
  def stringifyWith(reg: Int, word: Int): String
}

sealed trait OneOpAddressing extends AddressingMode
sealed trait TwoOpAddressing extends AddressingMode
sealed trait TwoOpSourceAddressing extends TwoOpAddressing
sealed trait TwoOpDestAddressing extends TwoOpAddressing

object TwoOpSourceAddressingModes {
  // just alias the apply over to make names look right at call sites.
  // think hard about refactoring this
  def apply(register: Int, mode: Int): Option[TwoOpSourceAddressing] = OneOpAddressingModes(register, mode)
}

object OneOpAddressingModes {
  trait AddressingMode extends OneOpAddressing with TwoOpSourceAddressing

  def apply(register: Int, mode: Int): Option[AddressingMode] = register match {
    case 0 => PCModes(mode) orElse General(mode)
    case 2 => SRModes(mode) orElse General(mode)
    case 3 => CGModes(mode) orElse General(mode)
    case _ => General(mode)
  }

  trait Modes {
    private val modeMap: Map[Int, AddressingMode] = Map()

    case class Mode(as: Int, name: String, stringifier: String => String) extends AddressingMode with NormMode {
      modeMap += as -> this
      def stringifyWith(reg: Int): String = stringifier(nameifyRegister(reg))
    }
    case class ExtraByteMode(as: Int, name: String, stringifier: (String, Int) => String) extends AddressingMode with WideMode {
      modeMap += as -> this
      def stringifyWith(reg: Int, const: Int): String = stringifier(nameifyRegister(reg), const)
    }

    def apply(x: Int): Option[AddressingMode] = modeMap.get(x)
  }

  object PCModes extends Modes {
    ExtraByteMode(3, "NextWord", (reg, byte) => s"${if (byte < 0) "-" else ""}0x${"%x".format(Math.abs(byte))}")
  }
  object SRModes extends Modes {
    Mode(2, "Const4", reg => "#4")
    Mode(3, "Const8", reg => "#8")
  }
  object CGModes extends Modes {
    Mode(0, "Const0", reg => "#0")
    Mode(1, "Const1", reg => "#1")
    Mode(2, "Const2", reg => "#2")
    Mode(3, "ConstNeg1", reg => "#-1")
  }
  object General extends Modes {
    Mode(0, "RegisterDirect", ident)
    ExtraByteMode(1, "Indexed", (reg, byte) => s"${if (byte < 0) "-" else ""}0x${"%x".format(Math.abs(byte))}($reg)")
    Mode(2, "RegisterIndirect", reg => s"@$reg")
    Mode(3, "IndirectAutoInc", reg => s"@$reg+")
  }

  PCModes
  SRModes
  CGModes
  General
}


object TwoOpDestAddressingModes {
  trait AddressingMode extends TwoOpDestAddressing

  def apply(register: Int, mode: Int): Option[AddressingMode] = register match {
    case 0 => PCModes(mode) orElse General(mode)
    case 2 => SRModes(mode) orElse General(mode)
    case _ => General(mode)
  }

  trait Modes {
    private val modeMap: Map[Int, AddressingMode] = Map()

    case class Mode(as: Int, name: String, stringifier: String => String) extends AddressingMode with NormMode {
      modeMap += as -> this
      def stringifyWith(reg: Int): String = stringifier(nameifyRegister(reg))
    }
    case class ExtraByteMode(as: Int, name: String, stringifier: (String, Int) => String) extends AddressingMode with WideMode {
      modeMap += as -> this
      def stringifyWith(reg: Int, const: Int): String = stringifier(nameifyRegister(reg), const)
    }

    def apply(x: Int): Option[AddressingMode] = modeMap.get(x)
  }

  object PCModes extends Modes {
    Mode(1, "Symbolic", reg => s"<??> (symbolic? x(PC)?")
  }
  object SRModes extends Modes {
    ExtraByteMode(1, "Absolute", (reg, byte) => s"&${"%4x".format(byte)}")
  }
  object General extends Modes {
    Mode(0, "RegisterDirect", reg => s"$reg")
    ExtraByteMode(1, "Indexed", (reg, byte) => s"${if (byte < 0) "-" else ""}0x${"%x".format(Math.abs(byte))}($reg)")
  }

  PCModes
  SRModes
  General
}