#encoding=utf-8

# Copyright (C) 2021 Collabora, Ltd.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation
# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice (including the next
# paragraph) shall be included in all copies or substantial portions of the
# Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL
# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

import argparse
import sys
import struct
from valhall import instructions, enums, immediates, typesize

LINE = ''

class ParseError(Exception):
    def __init__(self, error):
        self.error = error

class FAUState:
    def __init__(self, mode):
        self.mode = mode
        self.uniform_slot = None
        self.special = None
        self.buffer = set()

    def push(self, s):
        self.buffer.add(s)
        die_if(len(self.buffer) > 2, "Overflowed FAU buffer")

    def push_special(self, s):
        die_if(self.special is not None and self.special != s,
                'Multiple special immediates')
        self.special = s
        self.push(s)

    def descriptor(self, s):
        die_if(self.mode != 'none', f'Expected no modifier with {s}')
        self.push_special(s)

    def uniform(self, v):
        slot = v >> 1

        die_if(self.mode != 'none',
                'Expected uniform with default immediate mode')
        die_if(self.uniform_slot is not None and self.uniform_slot != slot,
                'Overflowed uniform slots')
        self.uniform_slot = slot
        self.push(f'uniform{v}')

    def id(self, s):
        die_if(self.mode != 'id',
                'Expected .id modifier with thread storage pointer')

        self.push_special(f'id{s}')

    def ts(self, s):
        die_if(self.mode != 'ts',
                'Expected .ts modifier with thread pointer')
        self.push_special(f'ts{s}')

    def constant(self, cons):
        self.push(cons)

# When running standalone, exit with the error since we're dealing with a
# human. Otherwise raise a Python exception so the test harness can handle it.
def die(s):
    if __name__ == "__main__":
        print(LINE)
        print(s)
        sys.exit(1)
    else:
        raise ParseError(s)

def die_if(cond, s):
    if cond:
        die(s)

def parse_int(s, minimum, maximum):
    try:
        number = int(s, base = 0)
    except ValueError:
        die(f"Expected number {s}")

    if number > maximum or number < minimum:
        die(f"Range error on {s}")

    return number

def encode_source(op, fau):
    if op == 'atest_datum':
        fau.descriptor(op)
        return 0x2A | 0xC0
    elif op.startswith('blend_descriptor_'):
        fau.descriptor(op)
        fin = op[len('blend_descriptor_'):]
        die_if(len(fin) != 3, 'Bad syntax')
        die_if(fin[1] != '_', 'Bad syntax')
        die_if(fin[2] not in ['x', 'y'], 'Bad component')

        rt = parse_int(fin[0], 0, 7)
        hi = 1 if (fin[2] == 'y') else 0
        return (0x30 | (2*rt) + hi) | 0xC0
    elif op[0] == '`':
        die_if(op[1] != 'r', f"Expected register after discard {op}")
        return parse_int(op[2:], 0, 63) | 0x40
    elif op[0] == 'r':
        return parse_int(op[1:], 0, 63)
    elif op[0] == 'u':
        val = parse_int(op[1:], 0, 63)
        fau.uniform(val)
        return val | 0x80
    elif op[0] == 'i':
        return int(op[3:]) | 0xC0
    elif op in enums['thread_storage_pointers'].bare_values:
        fau.ts(op)
        idx = 32 + enums['thread_storage_pointers'].bare_values.index(op)
        return idx | 0xC0
    elif op in enums['thread_identification'].bare_values:
        fau.id(op)
        idx = 32 + enums['thread_identification'].bare_values.index(op)
        return idx | 0xC0
    elif op.startswith('0x'):
        try:
            val = int(op, base=0)
        except ValueError:
            die('Expected value')

        die_if(val not in immediates, 'Unexpected immediate value')
        fau.constant(val)
        return immediates.index(val) | 0xC0
    else:
        die('Invalid operand')

def encode_dest(op):
    die_if(op[0] != 'r', f"Expected register destination {op}")

    parts = op.split(".")
    reg = parts[0]

    # Default to writing in full
    wrmask = 0x3

    if len(parts) > 1:
        WMASKS = ["h0", "h1"]
        die_if(len(parts) > 2, "Too many modifiers")
        mask = parts[1];
        die_if(mask not in WMASKS, "Expected a write mask")
        wrmask = 1 << WMASKS.index(mask)

    return parse_int(reg[1:], 0, 63) | (wrmask << 6)

def parse_asm(line):
    global LINE
    LINE = line # For better errors
    encoded = 0

    # Figure out mnemonic
    head = line.split(" ")[0]
    opts = [ins for ins in instructions if head.startswith(ins.name)]
    opts = sorted(opts, key=lambda x: len(x.name), reverse=True)

    if len(opts) == 0:
        die(f"No known mnemonic for {head}")

    if len(opts) > 1 and len(opts[0].name) == len(opts[1].name):
        print(f"Ambiguous mnemonic for {head}")
        print(f"Options:")
        for ins in opts:
            print(f"  {ins}")
        sys.exit(1)

    ins = opts[0]

    # Split off modifiers
    if len(head) > len(ins.name) and head[len(ins.name)] != '.':
        die(f"Expected . after instruction in {head}")

    mods = head[len(ins.name) + 1:].split(".")
    modifier_map = {}
    immediate_mode = 'none'

    for mod in mods:
        if mod in enums['immediate_mode'].bare_values:
            die_if(immediate_mode != 'none', 'Multiple immediate modes specified')
            immediate_mode = mod

    tail = line[(len(head) + 1):]
    operands = [x.strip() for x in tail.split(",") if len(x.strip()) > 0]
    expected_op_count = len(ins.srcs) + len(ins.dests) + len(ins.immediates) + len(ins.staging)
    if len(operands) != expected_op_count:
        die(f"Wrong number of operands in {line}, expected {expected_op_count}, got {len(operands)} {operands}")

    # Encode each operand
    for i, (op, sr) in enumerate(zip(operands, ins.staging)):
        die_if(op[0] != '@', f'Expected staging register, got {op}')
        parts = op[1:].split(':')

        die_if(any([x[0] != 'r' for x in parts]), f'Expected registers, got {op}')
        regs = [parse_int(x[1:], 0, 63) for x in parts]

        sr_count = len(regs)
        die_if(sr_count < 1, f'Expected staging register, got {op}')
        die_if(sr_count > 7, f'Too many staging registers {sr_count}')

        base = regs[0]
        die_if(any([reg != (base + i) for i, reg in enumerate(regs)]),
                'Expected consecutive staging registers, got {op}')

        if sr.count == 0:
            modifier_map["staging_register_count"] = sr_count
        else:
            die_if(sr_count != sr.count, f"Expected 4 staging registers, got {sr_count}")

        encoded |= ((sr.encoded_flags | base) << sr.start)
    operands = operands[len(ins.staging):]

    for op, dest in zip(operands, ins.dests):
        encoded |= encode_dest(op) << 40
    operands = operands[len(ins.dests):]

    if len(ins.dests) == 0 and len(ins.staging) == 0:
        # Set a placeholder writemask to prevent encoding faults
        encoded |= (0xC0 << 40)

    fau = FAUState(immediate_mode)

    for i, (op, src) in enumerate(zip(operands, ins.srcs)):
        parts = op.split('.')
        encoded |= encode_source(parts[0], fau) << (i * 8)

        # Has a swizzle been applied yet?
        swizzled = False

        for mod in parts[1:]:
            # Encode the modifier
            if mod in src.offset and src.bits[mod] == 1:
                encoded |= (1 << src.offset[mod])
            elif mod in enums[f'swizzles_{src.size}_bit'].bare_values and (src.widen or src.lanes):
                die_if(swizzled, "Multiple swizzles specified")
                swizzled = True
                val = enums[f'swizzles_{src.size}_bit'].bare_values.index(mod)
                encoded |= (val << src.offset['widen'])
            elif src.lane and mod in enums[f'lane_{src.size}_bit'].bare_values:
                die_if(swizzled, "Multiple swizzles specified")
                swizzled = True
                val = enums[f'lane_{src.size}_bit'].bare_values.index(mod)
                encoded |= (val << src.offset['lane'])
            elif src.size == 32 and mod in enums['widen'].bare_values:
                die_if(not src.swizzle, "Instruction doesn't take widens")
                die_if(swizzled, "Multiple swizzles specified")
                swizzled = True
                val = enums['widen'].bare_values.index(mod)
                encoded |= (val << src.offset['swizzle'])
            elif src.size == 16 and mod in enums['swizzles_16_bit'].bare_values:
                die_if(not src.swizzle, "Instruction doesn't take swizzles")
                die_if(swizzled, "Multiple swizzles specified")
                swizzled = True
                val = enums['swizzles_16_bit'].bare_values.index(mod)
                encoded |= (val << src.offset['swizzle'])
            elif mod in enums['lane_8_bit'].bare_values:
                die_if(not src.lane, "Instruction doesn't take a lane")
                die_if(swizzled, "Multiple swizzles specified")
                swizzled = True
                val = enums['lane_8_bit'].bare_values.index(mod)
                encoded |= (val << src.lane)
            else:
                die(f"Unknown modifier {mod}")

        # Encode the identity if a swizzle is required but not specified
        if src.swizzle and not swizzled and src.size == 16:
            mod = enums['swizzles_16_bit'].default
            val = enums['swizzles_16_bit'].bare_values.index(mod)
            encoded |= (val << src.offset['swizzle'])
        elif src.widen and not swizzled and src.size == 16:
            die_if(swizzled, "Multiple swizzles specified")
            mod = enums['swizzles_16_bit'].default
            val = enums['swizzles_16_bit'].bare_values.index(mod)
            encoded |= (val << src.offset['widen'])

    operands = operands[len(ins.srcs):]

    for i, (op, imm) in enumerate(zip(operands, ins.immediates)):
        if op[0] == '#':
            die_if(imm.name != 'constant', "Wrong syntax for immediate")
            parts = [imm.name, op[1:]]
        else:
            parts = op.split(':')
            die_if(len(parts) != 2, f"Wrong syntax for immediate, wrong number of colons in {op}")
            die_if(parts[0] != imm.name, f"Wrong immediate, expected {imm.name}, got {parts[0]}")

        if imm.signed:
            minimum = -(1 << (imm.size - 1))
            maximum = +(1 << (imm.size - 1)) - 1
        else:
            minimum = 0
            maximum = (1 << imm.size) - 1

        val = parse_int(parts[1], minimum, maximum)

        if val < 0:
            # Sign extends
            val = (1 << imm.size) + val

        encoded |= (val << imm.start)

    operands = operands[len(ins.immediates):]

    # Encode the operation itself
    encoded |= (ins.opcode << 48)
    encoded |= (ins.opcode2 << ins.secondary_shift)

    # Encode modifiers
    has_action = False
    for mod in mods:
        if len(mod) == 0:
            continue

        if mod in enums['action'].bare_values:
            die_if(has_action, "Multiple actions specified")
            has_action = True
            encoded |= (enums['action'].bare_values.index(mod) << 59)
            encoded |= (1 << 62) # Action, not wait
        elif mod.startswith('wait'):
            die_if(has_action, "Multiple actions specified")
            has_action = True

            slots = mod[len('wait'):]
            try:
                slots = set([int(x) for x in slots])
            except ValueError:
                die(f"Expected slots in {mod}")

            known_slots = set([0, 1, 2])
            die_if(not slots.issubset(known_slots), f"Unknown slots in {mod}")

            if 0 in slots:
                encoded |= (1 << 59)
            if 1 in slots:
                encoded |= (1 << 60)
            if 2 in slots:
                encoded |= (1 << 61)
        elif mod in enums['immediate_mode'].bare_values:
            pass # handled specially
        else:
            candidates = [c for c in ins.modifiers if mod in c.bare_values]

            die_if(len(candidates) == 0, f"Invalid modifier {mod} used")
            assert(len(candidates) == 1) # No ambiguous modifiers
            opts = candidates[0]

            value = opts.bare_values.index(mod)
            assert(value is not None)

            die_if(opts.name in modifier_map, f"{opts.name} specified twice")
            modifier_map[opts.name] = value

    for mod in ins.modifiers:
        value = modifier_map.get(mod.name, mod.default)
        die_if(value is None, f"Missing required modifier {mod.name}")

        assert(value < (1 << mod.size))
        encoded |= (value << mod.start)

    encoded |= (enums['immediate_mode'].bare_values.index(immediate_mode) << 57)
    return encoded

if __name__ == "__main__":
    # Provide commandline interface
    parser = argparse.ArgumentParser(description='Assemble Valhall shaders')
    parser.add_argument('infile', nargs='?', type=argparse.FileType('r'),
                        default=sys.stdin)
    parser.add_argument('outfile', type=argparse.FileType('wb'))
    args = parser.parse_args()

    lines = args.infile.read().strip().split('\n')
    lines = [l for l in lines if len(l) > 0 and l[0] != '#']

    packed = b''.join([struct.pack('<Q', parse_asm(ln)) for ln in lines])
    args.outfile.write(packed)