package wacc

object assemblyIR {

  sealed trait AsmLine
  sealed trait Operand
  sealed trait Src extends Operand // mem location, register and imm value
  sealed trait Dest extends Operand // mem location and register

  enum Size {
    case Q64, D32, W16, B8

    def toInt: Int = this match {
      case Q64 => 8
      case D32 => 4
      case W16 => 2
      case B8  => 1
    }

    private val ptr = "ptr "

    override def toString(): String = this match {
      case Q64 => "qword " + ptr
      case D32 => "dword " + ptr
      case W16 => "word " + ptr
      case B8  => "byte " + ptr
    }
  }

  enum RegName {
    case AX, BX, CX, DX, SI, DI, SP, BP, IP, R8, R9, R10, R11, R12, R13, R14, R15
  }

  case class Register(size: Size, name: RegName) extends Dest with Src {
    import RegName._

    if (size == Size.B8 && name == RegName.IP) {
      throw new IllegalArgumentException("Cannot have 8 bit register for IP")
    }
    override def toString = name match {
      case AX  => tradToString("ax", "al")
      case BX  => tradToString("bx", "bl")
      case CX  => tradToString("cx", "cl")
      case DX  => tradToString("dx", "dl")
      case SI  => tradToString("si", "sil")
      case DI  => tradToString("di", "dil")
      case SP  => tradToString("sp", "spl")
      case BP  => tradToString("bp", "bpl")
      case IP  => tradToString("ip", "#INVALID")
      case R8  => newToString(8)
      case R9  => newToString(9)
      case R10 => newToString(10)
      case R11 => newToString(11)
      case R12 => newToString(12)
      case R13 => newToString(13)
      case R14 => newToString(14)
      case R15 => newToString(15)
    }

    private def tradToString(base: String, byteName: String): String =
      size match {
        case Size.Q64 => "r" + base
        case Size.D32 => "e" + base
        case Size.W16 => base
        case Size.B8  => byteName
      }

    private def newToString(base: Int): String = {
      val b = base.toString
      "r" + (size match {
        case Size.Q64 => b
        case Size.D32 => b + "d"
        case Size.W16 => b + "w"
        case Size.B8  => b + "b"
      })
    }
  }

  // arguments
  enum CLibFunc extends Operand {
    case Scanf,
      Fflush,
      Exit,
      PrintF,
      Malloc,
      Free

    private val plt = "@plt"

    override def toString = this match {
      case Scanf  => "scanf" + plt
      case Fflush => "fflush" + plt
      case Exit   => "exit" + plt
      case PrintF => "printf" + plt
      case Malloc => "malloc" + plt
      case Free   => "free" + plt
    }
  }

  case class MemLocation(
      base: Register,
      offset: Int | LabelArg = 0,
      // scale 0 will make register irrelevant, no other reason as to why it's RAX
      scaledIndex: (Register, Int) = (Register(Size.Q64, RegName.AX), 0),
      opSize: Option[Size] = None
  ) extends Dest
      with Src {
    def copy(
        base: Register = this.base,
        offset: Int | LabelArg = this.offset,
        scaledIndex: (Register, Int) = this.scaledIndex,
        opSize: Option[Size] = this.opSize
    ): MemLocation = MemLocation(base, offset, scaledIndex, opSize)

    override def toString(): String = {
      val opSizeStr = opSize.map(_.toString).getOrElse("")
      val baseStr = base.toString
      val offsetStr = offset match {
        case 0   => ""
        case off => s" + $off"
      }
      val scaledIndexStr = scaledIndex match {
        case (reg, scale) if scale != 0 => s" + $reg * $scale"
        case _                          => ""
      }
      s"$opSizeStr[$baseStr$scaledIndexStr$offsetStr]"
    }
  }

  case class ImmediateVal(value: Int) extends Src {
    override def toString = value.toString
  }

  case class LabelArg(name: String) extends Operand {
    override def toString = name
  }

  abstract class Operation(ins: String, ops: Operand*) extends AsmLine {
    override def toString: String = s"\t$ins ${ops.mkString(", ")}"
  }

  // arithmetic operations
  case class Add(op1: Dest, op2: Src) extends Operation("add", op1, op2)
  case class Subtract(op1: Dest, op2: Src) extends Operation("sub", op1, op2)
  case class Multiply(ops: Operand*) extends Operation("imul", ops*)
  case class Divide(op1: Src) extends Operation("idiv", op1)
  case class Negate(op: Dest) extends Operation("neg", op)
  // bitwise operations
  case class And(op1: Dest, op2: Src) extends Operation("and", op1, op2)
  case class Or(op1: Dest, op2: Src) extends Operation("or", op1, op2)
  case class Xor(op1: Dest, op2: Src) extends Operation("xor", op1, op2)
  case class Compare(op1: Dest, op2: Src) extends Operation("cmp", op1, op2)
  case class CDQ() extends Operation("cdq")
  // stack operations
  case class Push(op1: Src) extends Operation("push", op1)
  case class Pop(op1: Src) extends Operation("pop", op1)
  // move operations
  case class Move(op1: Dest, op2: Src) extends Operation("mov", op1, op2)
  case class Load(op1: Register, op2: MemLocation) extends Operation("lea ", op1, op2)

  // function call operations
  case class Call(op1: CLibFunc | LabelArg) extends Operation("call", op1)
  case class Return() extends Operation("ret")

  // conditional operations
  case class Jump(op1: LabelArg, condition: Cond = Cond.Always)
      extends Operation(s"j${condition.toString}", op1)
  case class Set(op1: Dest, condition: Cond = Cond.Always)
      extends Operation(s"set${condition.toString}", op1)

  case class LabelDef(name: String) extends AsmLine {
    override def toString = s"$name:"
  }

  case class Comment(comment: String) extends AsmLine {
    override def toString =
      comment.split("\n").map(line => s"# ${line}").mkString("\n")
  }

  enum Cond {
    case Equal,
      NotEqual,
      Greater,
      GreaterEqual,
      Less,
      LessEqual,
      Overflow,
      Always
    override def toString(): String = this match {
      case Equal        => "e"
      case NotEqual     => "ne"
      case Greater      => "g"
      case GreaterEqual => "ge"
      case Less         => "l"
      case LessEqual    => "le"
      case Overflow     => "o"
      case Always       => "mp"
    }
  }

  enum Directive extends AsmLine {
    case IntelSyntax, RoData, Text, EndFunc
    case Global(name: String)
    case Int(value: scala.Int)
    case Asciz(string: String)
    case File(no: scala.Int, file: String)
    case Location(fileNo: scala.Int, lineNo: scala.Int, colNo: Option[scala.Int])
    case Func(name: String, label: LabelDef)
    case Type(label: LabelDef, symbolType: SymbolType)
    case Size(label: LabelDef, expr: SizeExpr)

    override def toString(): String = this match {
      case IntelSyntax   => ".intel_syntax noprefix"
      case Global(name)  => s".globl $name"
      case Text          => ".text"
      case RoData        => ".section .rodata"
      case Int(value)    => s"\t.int $value"
      case Asciz(string) => s"\t.asciz \"$string\""
      case File(no, file) => s".file $no \"${file}\""
      case Location(fileNo, lineNo, colNo) =>
        s"\t.loc $fileNo $lineNo" + colNo.map(c => s" $c").getOrElse("")
      case Func(name, label) =>
        s".func $name, ${label.name}"
      case EndFunc => ".endfunc"
      case Type(label, symbolType) =>
        s".type ${label.name}, @${symbolType.toString}"
      case Directive.Size(label, expr) =>
        s".size ${label.name}, ${expr.toString}"
    }
  }

  enum SymbolType {
    case Function

    override def toString(): String = this match {
      case Function => "function"
    }
  }

  enum SizeExpr {
    case Relative(label: LabelDef)

    override def toString(): String = this match {
      case Relative(label) => s".-${label.name}"
    }
  }

  enum PrintFormat {
    case Int, Char, String

    override def toString(): String = this match {
      case Int    => "%d"
      case Char   => "%c"
      case String => "%s"
    }
  }

  object commonRegisters {
    import Size._
    import RegName._

    val RAX = Register(Q64, AX)
    val EAX = Register(D32, AX)
    val RDI = Register(Q64, DI)
    val RIP = Register(Q64, IP)
    val RBP = Register(Q64, BP)
    val RSI = Register(Q64, SI)
    val RDX = Register(Q64, DX)
    val RCX = Register(Q64, CX)
    val ECX = Register(D32, CX)
  }
}