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 RegSize {
    case R64
    case E32
    case Byte

    override def toString = this match {
      case R64  => "r"
      case E32  => "e"
      case Byte => ""
    }
  }

  enum RegName {
    case AX, AL, BX, CX, DX, SI, DI, SP, BP, IP, Reg8, Reg9, Reg10, Reg11, Reg12, Reg13, Reg14,
      Reg15
    override def toString = this match {
      case AX    => "ax"
      case AL    => "al"
      case BX    => "bx"
      case CX    => "cx"
      case DX    => "dx"
      case SI    => "si"
      case DI    => "di"
      case SP    => "sp"
      case BP    => "bp"
      case IP    => "ip"
      case Reg8  => "8"
      case Reg9  => "9"
      case Reg10 => "10"
      case Reg11 => "11"
      case Reg12 => "12"
      case Reg13 => "13"
      case Reg14 => "14"
      case Reg15 => "15"
    }
  }

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

    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
    }
  }

  // TODO register naming conventions are wrong
  case class Register(size: RegSize, name: RegName) extends Dest with Src {
    override def toString = s"${size}${name}"
  }
  case class MemLocation(pointer: Long | Register, opSize: SizeDir = SizeDir.Unspecified)
      extends Dest
      with Src {
    override def toString = pointer match {
      case hex: Long     => opSize.toString + f"[0x$hex%X]"
      case reg: Register => opSize.toString + s"[$reg]"
    }
  }
  case class IndexAddress(
      base: Register,
      offset: Int | LabelArg,
      opSize: SizeDir = SizeDir.Unspecified
  ) extends Dest
      with Src {
    override def toString = s"$opSize[$base + $offset]"
  }

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

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

  // TODO Check if dest and src are not both memory locations
  abstract class Operation(ins: String, ops: Operand*) extends AsmLine {
    override def toString: String = s"\t$ins ${ops.mkString(", ")}"
  }
  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)

  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)

  // stack operations
  case class Push(op1: Src) extends Operation("push", op1)
  case class Pop(op1: Src) extends Operation("pop", op1)
  case class Call(op1: CLibFunc | LabelArg) extends Operation("call", op1)

  case class Move(op1: Dest, op2: Src) extends Operation("mov", op1, op2)
  case class Load(op1: Register, op2: MemLocation | IndexAddress)
      extends Operation("lea ", op1, op2)
  case class CDQ() extends Operation("cdq")

  case class Return() extends Operation("ret")

  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
    case Global(name: String)
    case Int(value: scala.Int)
    case Asciz(string: String)

    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".int $value"
      case Asciz(string) => s".asciz \"$string\""
    }
  }

  enum PrintFormat {
    case Int, Char, String

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

  enum SizeDir {
    case Byte, Word, DWord, Unspecified

    private val ptr = "ptr "

    override def toString(): String = this match {
      case Byte        => "byte " + ptr
      case Word        => "word " + ptr // TODO check word/doubleword/quadword
      case DWord       => "dword " + ptr
      case Unspecified => ""
    }
  }
}