#!/usr/bin/env python3
"""
Microcode Assembler for microcoded CPU
Parses .src files and generates Verilog hex files
"""

import sys
import re

class Field:
    """Represents a microcode field definition"""
    def __init__(self, name, width, default, enums):
        self.name = name
        self.width = width
        self.default = default  # Can be int or str
        self.enums = enums      # Dict of name -> value
        self.position = 0        # Bit position (set later)
    
    def encode(self, value):
        """Encode a value for this field"""
        if isinstance(value, str):
            if value in self.enums:
                return self.enums[value]
            else:
                raise ValueError(f"Unknown enum '{value}' for field '{self.name}'. Valid: {list(self.enums.keys())}")
        return int(value)

class MicrocodeAssembler:
    def __init__(self):
        self.fields = []
        self.field_map = {}
        self.subroutines = []
        self.sub_addresses = {}
        
        # Config
        self.microcode_length = 512
        self.microcode_width = 36
        self.microcode_output = 'ucode.hex'
        self.mapping_length = 32
        self.mapping_output = 'umap.hex'
        self.mapping_entries = [0] * 32  # Default all to 0 (fetch)
        
    def parse_config(self, line):
        """Parse configuration lines"""
        parts = line.split()
        if len(parts) < 3:
            return
        
        category, key, value = parts[0], parts[1], parts[2]
        
        if category == 'microcode':
            if key == 'length':
                self.microcode_length = int(value)
            elif key == 'width':
                self.microcode_width = int(value)
            elif key == 'output':
                self.microcode_output = value
        elif category == 'mapping':
            if key == 'length':
                self.mapping_length = int(value)
                self.mapping_entries = [0] * self.mapping_length
            elif key == 'output':
                self.mapping_output = value
    
    def parse_field(self, line):
        """
        Parse field definition: field <name> <width> <default> [enums]
        Example: field rd 3 imm pcp=0,dsp=1,rsp=2,tmp=3,mdr=4,tos=5,acc=6,imm=7
        """
        parts = line.split(None, 3)  # Split into max 4 parts
        if len(parts) < 4:
            raise ValueError(f"Invalid field: {line}")
        
        name = parts[1]
        width = int(parts[2])
        rest = parts[3]  # "default [enums...]"
        
        # Split rest into tokens
        tokens = rest.split()
        default_str = tokens[0]
        
        # Parse enums (everything after first token, comma-separated)
        enums = {}
        if len(tokens) > 1:
            enum_str = ' '.join(tokens[1:])
            for i, item in enumerate(enum_str.split(',')):
                item = item.strip()
                if '=' in item:
                    k, v = item.split('=')
                    enums[k.strip()] = int(v.strip())
                else:
                    enums[item] = i
        
        # Resolve default
        if default_str.isdigit():
            default_value = int(default_str)
        elif default_str in enums:
            default_value = enums[default_str]
        else:
            # Default is a symbolic name (like 'imm') that should be in enums
            # If not in enums yet, we'll handle it during encoding
            default_value = default_str
        
        field = Field(name, width, default_value, enums)
        self.fields.append(field)
        self.field_map[name] = field
    
    def layout_fields(self):
        """Assign bit positions to fields (LSB first)"""
        pos = 0
        for field in self.fields:
            field.position = pos
            pos += field.width
        
        if pos != self.microcode_width:
            print(f"Warning: Fields total {pos} bits, microcode width is {self.microcode_width}")
    
    def parse_subroutine(self, lines, start_idx):
        """Parse a subroutine block, returns (name, address, instructions, next_idx)"""
        # sub <name> [address]
        sub_match = re.match(r'sub\s+(\w+)(?:\s+([0-9A-Fa-fx]+h?))?', lines[start_idx])
        if not sub_match:
            raise ValueError(f"Invalid sub: {lines[start_idx]}")
        
        name = sub_match.group(1)
        address = None
        if sub_match.group(2):
            addr_str = sub_match.group(2).replace('h', '').replace('x', '')
            address = int(addr_str, 16)
        
        # Collect instructions until 'end'
        instructions = []
        idx = start_idx + 1
        while idx < len(lines):
            line = lines[idx].strip()
            if line.startswith('end'):
                break
            if line and not line.startswith(';'):
                instructions.append(line)
            idx += 1
        
        return name, address, instructions, idx + 1
    
    def parse_map(self, line):
        """Parse map statement"""
        # map entry1,entry2,entry3*count,...
        map_content = line[3:].strip()
        entries = [e.strip() for e in map_content.split(',')]
        
        map_idx = 0
        for entry in entries:
            if '*' in entry:
                # Repeat entry
                name, count_str = entry.split('*')
                name = name.strip()
                count = int(count_str.strip())
                
                if name not in self.sub_addresses:
                    raise ValueError(f"Unknown subroutine '{name}' in map")
                addr = self.sub_addresses[name]
                
                for _ in range(count):
                    if map_idx >= self.mapping_length:
                        raise ValueError("Map overflow")
                    self.mapping_entries[map_idx] = addr
                    map_idx += 1
            else:
                # Single entry
                if entry not in self.sub_addresses:
                    raise ValueError(f"Unknown subroutine '{entry}' in map")
                self.mapping_entries[map_idx] = self.sub_addresses[entry]
                map_idx += 1
    
    def assemble_instruction(self, inst_line):
        """Assemble a microcode instruction line into a word"""
        # Initialize with defaults
        field_values = {}
        for field in self.fields:
            if isinstance(field.default, str):
                # Resolve symbolic default
                if field.default in field.enums:
                    field_values[field.name] = field.enums[field.default]
                else:
                    raise ValueError(f"Default '{field.default}' not in enums for field '{field.name}'")
            else:
                field_values[field.name] = field.default
        
        # Parse assignments
        for assign in inst_line.split(','):
            assign = assign.strip()
            
            if '=' in assign:
                field_name, value_str = assign.split('=', 1)
                field_name = field_name.strip()
                value_str = value_str.strip()
                
                if field_name not in self.field_map:
                    raise ValueError(f"Unknown field '{field_name}': {inst_line}")
                
                field = self.field_map[field_name]
                
                # Parse value
                if value_str.endswith('h'):
                    value = int(value_str[:-1], 16)
                elif value_str.startswith('0x'):
                    value = int(value_str, 16)
                elif value_str.isdigit():
                    value = int(value_str)
                else:
                    # Symbolic value
                    value = field.encode(value_str)
                
                field_values[field_name] = value
            else:
                # Just field name -> set to 1
                if assign not in self.field_map:
                    raise ValueError(f"Unknown field '{assign}': {inst_line}")
                field_values[assign] = 1
        
        # Pack into word
        word = 0
        for field in self.fields:
            value = field_values[field.name]
            if value >= (1 << field.width):
                raise ValueError(f"Value {value} too large for field '{field.name}' (width {field.width})")
            word |= (value << field.position)
        
        return word
    
    def assemble(self, source_file):
        """Main assembly"""
        with open(source_file, 'r') as f:
            lines = [line.strip() for line in f.readlines()]
        
        # Pass 1: Config and fields
        idx = 0
        while idx < len(lines):
            line = lines[idx]
            
            if not line or line.startswith(';'):
                idx += 1
                continue
            
            if line.startswith('microcode') or line.startswith('mapping'):
                self.parse_config(line)
            elif line.startswith('field'):
                self.parse_field(line)
            elif line.startswith('sub'):
                break
            
            idx += 1
        
        self.layout_fields()
        
        # Pass 2: Subroutines
        while idx < len(lines):
            line = lines[idx]
            
            if not line or line.startswith(';'):
                idx += 1
                continue
            
            if line.startswith('sub'):
                name, address, instructions, next_idx = self.parse_subroutine(lines, idx)
                self.subroutines.append((name, address, instructions))
                idx = next_idx
            elif line.startswith('map'):
                break
            else:
                idx += 1
        
        # Assign addresses - FIXED ALGORITHM
        # Build a set of all occupied addresses from explicit subroutines
        occupied = set()
        for name, address, instructions in self.subroutines:
            if address is not None:
                self.sub_addresses[name] = address
                # Mark all addresses this subroutine occupies
                for i in range(len(instructions)):
                    occupied.add(address + i)
        
        # Now assign auto addresses to remaining subroutines
        auto_addr = 0
        for name, address, instructions in self.subroutines:
            if address is None:
                # Find next sequence of free addresses long enough for this subroutine
                while any((auto_addr + i) in occupied for i in range(len(instructions))):
                    auto_addr += 1
                
                self.sub_addresses[name] = auto_addr
                
                # Mark these addresses as occupied
                for i in range(len(instructions)):
                    occupied.add(auto_addr + i)
                
                auto_addr += len(instructions)
        
        # Pass 3: Map
        while idx < len(lines):
            line = lines[idx]
            if not line or line.startswith(';'):
                idx += 1
                continue
            if line.startswith('map'):
                self.parse_map(line)
            idx += 1
        
        # Generate microcode ROM
        microcode_rom = [0] * self.microcode_length
        
        for name, _, instructions in self.subroutines:
            addr = self.sub_addresses[name]
            for inst in instructions:
                if addr >= self.microcode_length:
                    raise ValueError(f"Address {addr} overflow")
                microcode_rom[addr] = self.assemble_instruction(inst)
                addr += 1
        
        # Write outputs
        self.write_hex(self.microcode_output, microcode_rom, self.microcode_width)
        self.write_hex(self.mapping_output, self.mapping_entries, 9)
        
        print(f"Assembly complete!")
        print(f"  Microcode: {self.microcode_output}")
        print(f"  Mapping:   {self.mapping_output}")
        print(f"\nSubroutines:")
        for name, addr in sorted(self.sub_addresses.items(), key=lambda x: x[1]):
            # Find instruction count
            inst_count = next(len(inst) for n, a, inst in self.subroutines if n == name)
            print(f"  {name:12s} @ 0x{addr:03X} ({inst_count} instructions)")
    
    def write_hex(self, filename, data, width):
        """Write Verilog hex file"""
        with open(filename, 'w') as f:
            hex_digits = (width + 3) // 4
            for word in data:
                f.write(f"{word:0{hex_digits}X}\n")

def main():
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <source.src>")
        sys.exit(1)
    
    asm = MicrocodeAssembler()
    try:
        asm.assemble(sys.argv[1])
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == '__main__':
    main()