#!/usr/bin/env python3
"""
CBE System Disk Builder - Clean Build
Builds a 10MB disk image with a TinyCPU OS kernel.
"""
import os, shutil, subprocess, struct

OUT = "/root/CBE/CBE"
BUILD = os.path.join(OUT, 'build')
DISK_SIZE = 10 * 1024 * 1024

OPCODES = {
    'nop':0x00,'mova_b':0x01,'movia':0x02,'add':0x03,
    'sub':0x04,'jmpa':0x05,'jza':0x06,'jca':0x07,
    'cmp':0x08,'movib':0x09,'store':0x0A,'load':0x0B,
    'jmpi':0x0C,'jzi':0x0D,'jci':0x0E,'inc':0x0F,
    'call':0x10,'ret':0x11,'hlt':0xFF,
}

GPU=0xC0; KST=0xFA; KDT=0xFB; DSK=0xB0; DRD=0xB1; DDT=0xB2
VC=0x50; VT=0x51; LB=0x60

class Asm:
    def __init__(self):
        self.ops = []; self.labels = {}
    def label(self, n): self.ops.append(('_label', n))
    def op(self, m, a=None): self.ops.append((m, a))
    def db(self, d):
        if isinstance(d, int): d = bytes([d])
        elif isinstance(d, str): d = d.encode()
        self.ops.append(('_db', bytes(d)))

    def assemble(self, origin=0, max_size=1024):
        pos = origin
        for m, a in self.ops:
            if m == '_label': self.labels[a] = pos
            elif m == '_db': pos += len(a)
            else: pos += 2 if a is not None else 1
        total = pos - origin
        if total > max_size: raise RuntimeError(f"Overflow: {total}>{max_size}")
        buf, pos2, err = bytearray(max_size), origin, False
        for m, a in self.ops:
            if m == '_label': continue
            o = pos2 - origin
            if m == '_db':
                for b in a: buf[o]=b; o+=1; pos2+=1
                continue
            if m not in OPCODES: print(f"  Bad op '{m}'"); err=True; pos2+=1; continue
            buf[o] = OPCODES[m]; pos2+=1
            if a is not None:
                o2 = pos2 - origin
                if isinstance(a, str):
                    if a in self.labels: buf[o2] = self.labels[a] & 0xFF
                    else: print(f"  Unresolved '{a}'"); buf[o2]=0; err=True
                else: buf[o2] = a & 0xFF
                pos2+=1
        return bytes(buf[:total]), err


def build_kernel():
    a = Asm()

    # === START (entry at 0x00) ===
    # CPU boots at PC=0, jump to main immediately
    a.label('start')
    a.op('jmpi', 'main')

    # === PRINT SUBROUTINE ===
    # Call with: B=0, write string addr to print_oper, CALL print_sub
    a.label('print_sub')
    a.op('load')              # LOAD_A with self-modifying operand
    a.label('print_oper')
    a.db(0x00)                # operand byte = address of current char
    a.op('cmp')               # A vs B (0 = null)
    a.op('jzi', 'print_end')
    a.op('store', GPU)        # output char
    a.op('load', 'print_oper')# load pointer value
    a.op('inc')
    a.op('store', 'print_oper')
    a.op('jmpi', 'print_sub')
    a.label('print_end')
    a.op('ret')

    # === SELF-MODDING STORE ===
    # Jump here with: A=char, store target addr in store_oper
    # After execution, continues at after_store
    a.label('sma_store')
    a.op('store')             # STORE_A with self-modifying operand
    a.label('store_oper')
    a.db(0x00)                # operand = target address
    # Continuation at sma_store+2:
    a.label('after_store')
    a.op('load', VC)
    a.op('inc')
    a.op('store', VC)
    a.op('jmpi', 'kbd')

    # === CMD: help ===
    a.label('cmd_help')
    a.op('movib', 0)
    a.op('movia', 'str_help')
    a.op('store', 'print_oper')
    a.op('call', 'print_sub')
    a.op('jmpi', 'prompt')

    # === CMD: clear ===
    a.label('cmd_clear')
    a.op('movia', 0x0C); a.op('store', GPU)
    a.op('jmpi', 'prompt')

    # === CMD: reboot ===
    a.label('cmd_reboot')
    a.op('jmpi', 0x00)

    # === CMD: cat (read sector 3 and print) ===
    a.label('cmd_cat')
    a.op('movia', 3); a.op('store', DSK)
    a.label('cat_wait')
    a.op('load', DRD); a.op('movib', 1); a.op('cmp')
    a.op('jzi', 'cat_ready'); a.op('jmpi', 'cat_wait')
    a.label('cat_ready')
    a.op('movib', 0); a.op('movia', DDT)
    a.op('store', 'print_oper'); a.op('call', 'print_sub')
    a.op('jmpi', 'prompt')

    # === CMD: ls (read sector 2 directory and print) ===
    a.label('cmd_ls')
    a.op('movia', 2); a.op('store', DSK)
    a.label('ls_wait')
    a.op('load', DRD); a.op('movib', 1); a.op('cmp')
    a.op('jzi', 'ls_ready'); a.op('jmpi', 'ls_wait')
    a.label('ls_ready')
    a.op('movib', 0); a.op('movia', DDT)
    a.op('store', 'print_oper'); a.op('call', 'print_sub')
    a.op('jmpi', 'prompt')

    # === MAIN ===
    a.label('main')
    a.label('prompt')
    a.op('movia', 0x0D); a.op('store', GPU)
    a.op('movia', 0x0A); a.op('store', GPU)
    a.op('movia', ord('>')); a.op('store', GPU)
    a.op('movia', ord(' ')); a.op('store', GPU)
    a.op('movib', 0)
    a.op('store', VC)

    # === KBD LOOP ===
    a.label('kbd')
    a.op('load', KST); a.op('cmp'); a.op('jzi', 'kbd')
    a.op('load', KDT)
    a.op('store', VT); a.op('store', GPU)
    a.op('movia', 0); a.op('store', KST)  # ack

    # Check Enter (VK_ENTER=0x0A)
    a.op('load', VT); a.op('movia', 0x0A); a.op('cmp')
    a.op('jzi', 'process')

    # Check Backspace (VK_BACK_SPACE=0x08)
    a.op('load', VT); a.op('movia', 0x08); a.op('cmp')
    a.op('jzi', 'backspace')

    # Buffer full?
    a.op('load', VC); a.op('movia', 15); a.op('cmp')
    a.op('jzi', 'kbd')

    # Store char: self-modding via sma_store
    a.op('load', VC); a.op('movib', LB); a.op('add')
    a.op('store', 'store_oper')
    a.op('load', VT)
    a.op('jmpi', 'sma_store')
    # after_store continues with cursor++ and jmpi kbd

    # === BACKSPACE ===
    a.label('backspace')
    a.op('load', VC); a.op('cmp'); a.op('jzi', 'kbd')
    a.op('load', VC); a.op('movib', 1); a.op('sub')
    a.op('store', VC)
    a.op('jmpi', 'kbd')

    # === PROCESS LINE ===
    a.label('process')
    a.op('movia', 0x0D); a.op('store', GPU)
    a.op('movia', 0x0A); a.op('store', GPU)

    # Dispatch (VK codes: H=72, C=67, R=82, D=68, L=76)
    a.op('load', LB)
    a.op('movia', 72); a.op('cmp'); a.op('jzi', 'cmd_help')
    a.op('load', LB)
    a.op('movia', 67); a.op('cmp'); a.op('jzi', 'cmd_clear')
    a.op('load', LB)
    a.op('movia', 82); a.op('cmp'); a.op('jzi', 'cmd_reboot')
    a.op('load', LB)
    a.op('movia', 68); a.op('cmp'); a.op('jzi', 'cmd_cat')
    a.op('load', LB)
    a.op('movia', 76); a.op('cmp'); a.op('jzi', 'cmd_ls')

    # Unknown
    a.op('movia', ord('?')); a.op('store', GPU)
    a.op('movia', 0x0A); a.op('store', GPU)
    a.op('jmpi', 'prompt')

    # === STRINGS ===
    a.label('str_help')
    a.db(b'Commands: H=help C=clear R=reboot D=cat L=ls\0')

    return a


def build_disk(kernel):
    if kernel is None: return None
    disk = bytearray(DISK_SIZE)
    disk[510] = 0x55; disk[511] = 0xAA
    ks = min(len(kernel), 512)
    disk[512:512+ks] = kernel[:ks]
    # Sector 2: directory text
    dir_txt = b"FILES: README  HELLO\0"
    disk[1024:1024+len(dir_txt)] = dir_txt
    # Sector 3: README
    readme = b"CBE/OS v2.0 for TinyCPU\r\nCommands: H C R D L\r\n\0"
    disk[1536:1536+len(readme)] = readme
    # Sector 4: HELLO
    hello = b"Hello from CBE/OS!\r\n\0"
    disk[2048:2048+len(hello)] = hello
    return bytes(disk)


def main():
    os.makedirs(BUILD, exist_ok=True)
    print("CBE Disk Builder - Clean Build")
    print("=" * 40)
    
    a = build_kernel()
    print(f"Ops: {len(a.ops)}")
    code, err = a.assemble(0, 1024)
    if err:
        for n, addr in sorted(a.labels.items(), key=lambda x: x[1] if x[1] is not None else 0):
            print(f"  {n}: 0x{addr:02X}" if addr is not None else f"  {n}: ???")
        return 1
    
    print(f"Kernel: {len(code)} bytes")
    for n, addr in sorted(a.labels.items(), key=lambda x: x[1]):
        print(f"  0x{addr:02X}: {n}")
    
    # Check for label positions
    if a.labels.get('start') != 0:
        print(f"ERROR: start should be 0x00, got 0x{a.labels.get('start', -1):02X}")
        return 1
    if a.labels.get('print_oper', -1) != a.labels.get('print_sub', -1) + 1:
        print(f"ERROR: print_oper not at print_sub+1")
        return 1
    if a.labels.get('store_oper', -1) != a.labels.get('sma_store', -1) + 1:
        print(f"ERROR: store_oper not at sma_store+1")
        return 1

    disk = build_disk(code)
    disk_path = os.path.join(BUILD, 'data.bin')
    with open(disk_path, 'wb') as f: f.write(disk)
    
    module_path = os.path.join(BUILD, 'big-disk.disk', 'banks', 'data.bin')
    os.makedirs(os.path.dirname(module_path), exist_ok=True)
    with open(module_path, 'wb') as f: f.write(disk)
    
    jar = os.path.join(OUT, 'modules', 'gui', 'build', 'stage', 'cbe-emu.jar')
    if os.path.exists(jar):
        src_dir = os.path.join(BUILD, 'big-disk.disk')
        out_name = os.path.join(BUILD, 'big-disk.cbeplugin')
        print(f"Compiling...")
        r = subprocess.run(['java', '-cp', jar, 'com.cbe.cbecc.Main', 'build', src_dir, '-o', out_name],
            capture_output=True, text=True, timeout=30)
        if r.returncode == 0: print("  OK")
        else: print(f"  FAIL: {r.stderr[:200]}")
        
        android_dir = os.path.join(OUT, 'android', 'app', 'src', 'main', 'assets', 'plugins')
        os.makedirs(android_dir, exist_ok=True)
        src = os.path.join(BUILD, 'big-disk.cbeplugin')
        if os.path.exists(src): shutil.copy(src, os.path.join(android_dir, 'big-disk.cbeplugin'))
    else:
        print(f"  JAR not found at {jar}")
    
    print("Done!")
    return 0

if __name__ == '__main__':
    exit(main())