package main

import (
    "bytes"
    "encoding/binary"
    "flag"
    "fmt"
    "./qvm"
    "os"
)

var OpcodeTable = [...]string {
    "UNDEF     ", "IGNORE    ", "BREAK     ", "ENTER     ", "LEAVE     ", "CALL      ",
    "PUSH      ", "POP       ", "CONST     ", "LOCAL     ", "JUMP      ", "EQ        ",
    "NE        ", "LTI       ", "LEI       ", "GTI       ", "GEI       ", "LTU       ",
    "LEU       ", "GTU       ", "GEU       ", "EQF       ", "NEF       ", "LTF       ",
    "LEF       ", "GTF       ", "GEF       ", "LOAD1     ", "LOAD2     ", "LOAD4     ",
    "STORE1    ", "STORE2    ", "STORE4    ", "ARG       ", "BLOCK_COPY", "SEX8      ",
    "SEX16     ", "NEGI      ", "ADD       ", "SUB       ", "DIVI      ", "DIVU      ",
    "MODI      ", "MODU      ", "MULI      ", "MULU      ", "BAND      ", "BOR       ",
    "BXOR      ", "BCOM      ", "LSH       ", "RSHI      ", "RSHU      ", "NEGF      ",
    "ADDF      ", "SUBF      ", "DIVF      ", "MULF      ", "CVIF      ", "CVFI      "}

const ( OP_ENTER = 3
    OP_CALL = 5
    OP_CONST = 8
    OP_JUMP = 10 )

var ArgTable = []byte {
    0, 0, 0, 4, 4, 0,
    0, 0, 4, 4, 0, 4,
    4, 4, 4, 4, 4, 4,
    4, 4, 4, 4, 4, 4,
    4, 4, 4, 0, 0, 0,
    0, 0, 0, 1, 4, 0,
    0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0,
    0, 0, 0, 0, 0, 0}

type instruction struct {
    op uint8
    arg []byte
    valid bool
}

type procedure struct {
    name string
    startInstruction, offset int
    instructionCount int
    callees []*procedure
    callers []*procedure
}

var insns []instruction
var procs map[int]*procedure
var strs map[int]string

func printHeader(qvmFile *qvm.QVMFile) {
    version := "unknown"
    switch qvmFile.Header.Magic {
        case qvm.VM_MAGIC_VER1:
            version = "Ver1"
        case qvm.VM_MAGIC_VER2:
            version = "Ver2"
    }
    fmt.Printf("            Magic: 0x%x[%s]\n", qvmFile.Header.Magic, version)
    fmt.Printf("Instruction Count: 0x%x\n", qvmFile.Header.InstructionCount)
    fmt.Printf("      Code Offset: 0x%x\n", qvmFile.Header.CodeOffset)
    fmt.Printf("      Code Length: 0x%x\n", qvmFile.Header.CodeLength)
    fmt.Printf("      Data Offset: 0x%x\n", qvmFile.Header.DataOffset)
    fmt.Printf("      Data Length: 0x%x\n", qvmFile.Header.DataLength)
    fmt.Printf("       Lit Length: 0x%x\n", qvmFile.Header.LitLength)
    fmt.Printf("       Bss Length: 0x%x\n", qvmFile.Header.BssLength)
    if qvmFile.Header.Magic == qvm.VM_MAGIC_VER2 {
        fmt.Printf("Jump Table Length: 0x%x\n", qvmFile.Header.JumpTableLength)
    }
}

func parseInstructions(qvmFile *qvm.QVMFile) {
    codeOffset := int(qvmFile.Header.CodeOffset)
    for i, off := 0, 0; i < int(qvmFile.Header.InstructionCount); i++ {
        op := qvmFile.Code[off]
        if op < 0 || op > 59 {
            fmt.Println("Illegal opcode: %d at 0x%x\n", op, off + codeOffset)
            insns = append(insns, instruction{op, []byte{}, false})
            off++
            continue
        }
        switch ArgTable[op] {
            case 0:
                insns = append(insns, instruction{op, []byte{}, true})
                off++
            case 1:
                insns = append(insns, instruction{op, []byte{qvmFile.Code[off + 1]}, true})
                off += 2
            case 4:
                insns = append(insns, instruction{op, []byte{qvmFile.Code[off + 1],
                                         qvmFile.Code[off + 2],
                                         qvmFile.Code[off + 3],
                                         qvmFile.Code[off + 4]}, true})
                off += 5
        }
    }
}

func parseProcedures() {
    procs = make(map[int]*procedure, 0)
    lastIndex := 0
    for i, off:= 0, 0; i < len(insns); i++ {
        if insns[i].op == OP_ENTER {
            procs[i] = &procedure{fmt.Sprintf("sub_%08x", off), i, off, 0, nil, nil}
            procs[lastIndex].instructionCount = i - procs[lastIndex].startInstruction
            lastIndex = i
        }
        off = off + 1 + int(ArgTable[insns[i].op])
    }
    procs[lastIndex].instructionCount = len(insns) - procs[lastIndex].startInstruction
}

func parseCodeXRefs() {
    for _, proc := range procs {
        for i := proc.startInstruction; i < proc.startInstruction + proc.instructionCount; i++ {
            if insns[i].op == OP_CONST && insns[i+1].op == OP_CALL {
                var target int32
                bufArg := bytes.NewBuffer(insns[i].arg)
                binary.Read(bufArg, binary.LittleEndian, &target)
                if target > 0 {
                    proc.callees = append(proc.callees, procs[int(target)])
                    procs[int(target)].callers = append(procs[int(target)].callers, proc)
                }
            }
        }
    }
}

func parseStrings(qvmFile *qvm.QVMFile) {
    strs = make(map[int]string, 0)
    firstChar := 0
    foundFirstChar := false
    for i, char := range qvmFile.Lit {
        if !foundFirstChar && char != '\x00' {
            firstChar = i
            foundFirstChar = true
        } else if foundFirstChar && char == '\x00' {
            str := string(qvmFile.Lit[firstChar:i])
            strs[firstChar + int(qvmFile.Header.DataOffset + qvmFile.Header.DataLength)] = str
            firstChar = i
            foundFirstChar = false
        }
    }
}

func disassemble() {
    for _, proc := range procs {
        for j := proc.startInstruction; j < proc.startInstruction + proc.instructionCount; j++ {
            arg := ""
            comment := ""
            mnemonic := ""
            if insns[j].valid {
                mnemonic = OpcodeTable[insns[j].op]
                switch ArgTable[insns[j].op] {
                    case 1:
                        arg = fmt.Sprintf("0x%02x", insns[j].arg[0])
                    case 4:
                        var target int32
                        bufArg := bytes.NewBuffer(insns[j].arg)
                        binary.Read(bufArg, binary.LittleEndian, &target)
                        arg = fmt.Sprintf("0x%08x", uint32(target))
                        if insns[j].op == OP_CONST && insns[j+1].op == OP_CALL {
                            if target < 0 {
                                //Syscall...
                                comment = fmt.Sprintf("; Syscall: %d", target)
                            } else {
                                //Code xref...
                                comment = fmt.Sprintf("; %s", procs[int(target)].name)
                            }
                        }
                }
            } else {
                //Invalid opcode...
                mnemonic = fmt.Sprint(".byte 0x%02x", insns[j].op)
            }
            fmt.Printf("<0x%08x>: %s %s%s\n", j, mnemonic, arg, comment);
        }
    }
}

func main() {
    var syscallsFile string
    var commentsFile string
    flag.StringVar(&syscallsFile, "s", "", "Specify a file defining the system calls for this QVM.")
    flag.StringVar(&commentsFile, "c", "", "Specify a file containing comments from a previous disassembly.")
    flag.Parse()

    if flag.NArg() != 1 {
        fmt.Println("Must specify one QVM!")
    }
    file, err := os.Open(flag.Arg(0), os.O_RDONLY, 0444)
    if err != nil {
        fmt.Println(err.String())
        os.Exit(-1)
    }
    qvmFile, err := qvm.NewQVMFile(file)
    parseInstructions(qvmFile)
    parseProcedures()
    parseCodeXRefs()
    parseStrings(qvmFile)
    printHeader(qvmFile)
    disassemble()
}