const REGISTERS = ['R0','R1','R2','R3','R4','R5','R6','R7']
const OPCODES = [
  'MOV','ADD','SUB','MUL','XDIV','XMOD',
  'AND','ORR','EOR','MVN','LSL','LSR',
  'CMP','B','BEQ','BNE','BGT','BLT','BGE','BLE',
  'LDR','STR','PUSH','POP','XOUT','XHLT','NOP',
]
const BRANCH_OPS = ['B','BEQ','BNE','BGT','BLT','BGE','BLE']
const MAX_STEPS = 10000

function parseNumber(s) {
  s = s.trim()
  if (/^0x[0-9a-fA-F]+$/i.test(s)) return parseInt(s, 16)
  if (/^0b[01]+$/i.test(s)) return parseInt(s.slice(2), 2)
  if (/^-?\d+$/.test(s)) return parseInt(s, 10)
  return null
}

function splitOperands(str) {
  const result = []; let cur = ''; let depth = 0
  for (const ch of str) {
    if (ch === '[' || ch === '{') depth++
    if (ch === ']' || ch === '}') depth--
    if (ch === ',' && depth === 0) { result.push(cur.trim()); cur = '' }
    else cur += ch
  }
  if (cur.trim()) result.push(cur.trim())
  return result
}

function parseOperand(s) {
  s = s.trim()
  const upper = s.toUpperCase()
  if (REGISTERS.includes(upper)) return { type: 'reg', value: upper }

  // Immediate: #42, #0xFF, #0b101
  if (s.startsWith('#')) {
    const n = parseNumber(s.slice(1))
    if (n !== null) return { type: 'imm', value: n }
  }

  // Memory: [R0], [R0, #4], [#addr]
  const memMatch = s.match(/^\[([^\],]+)(?:,\s*(.+))?\]$/)
  if (memMatch) {
    const base = memMatch[1].trim()
    const baseUp = base.toUpperCase()
    if (REGISTERS.includes(baseUp)) {
      let offset = 0
      if (memMatch[2]) {
        let off = memMatch[2].trim()
        if (off.startsWith('#')) off = off.slice(1)
        offset = parseNumber(off) || 0
      }
      return { type: 'mem', base: baseUp, offset }
    }
    // [#addr] direct
    let addr = base.startsWith('#') ? base.slice(1) : base
    const n = parseNumber(addr)
    if (n !== null) return { type: 'mem_direct', addr: n & 0xFF }
  }

  // {R0} for PUSH/POP
  const braceMatch = s.match(/^\{(.+)\}$/)
  if (braceMatch) {
    const inner = braceMatch[1].trim().toUpperCase()
    if (REGISTERS.includes(inner)) return { type: 'reg', value: inner }
  }

  // Label
  if (/^[a-zA-Z_]\w*$/.test(s)) return { type: 'label', value: s.toLowerCase() }

  // Bare number (lenient)
  const n = parseNumber(s)
  if (n !== null) return { type: 'imm', value: n }

  return { type: 'unknown', raw: s }
}

export function parse(code) {
  const lines = code.split('\n')
  const instructions = []
  const labels = {}
  const errors = []

  // Pass 1: find labels
  let idx = 0
  for (let i = 0; i < lines.length; i++) {
    let line = lines[i]
    const ci = line.indexOf(';'); if (ci >= 0) line = line.slice(0, ci)
    line = line.trim(); if (!line) continue
    const lm = line.match(/^([a-zA-Z_]\w*)\s*:(.*)$/)
    if (lm) { labels[lm[1].toLowerCase()] = idx; line = lm[2].trim(); if (!line) continue }
    idx++
  }

  // Pass 2: parse instructions
  for (let i = 0; i < lines.length; i++) {
    let line = lines[i]
    const ci = line.indexOf(';'); if (ci >= 0) line = line.slice(0, ci)
    line = line.trim(); if (!line) continue
    const lm = line.match(/^([a-zA-Z_]\w*)\s*:(.*)$/)
    if (lm) { line = lm[2].trim(); if (!line) continue }

    const parts = line.match(/^(\w+)(?:\s+(.*))?$/)
    if (!parts) { errors.push({ line: i, msg: `Syntax error: ${lines[i].trim()}` }); continue }

    const opcode = parts[1].toUpperCase()
    if (!OPCODES.includes(opcode)) { errors.push({ line: i, msg: `Unknown instruction: ${parts[1]}` }); continue }

    let operands = []
    if (parts[2]) {
      operands = splitOperands(parts[2]).map(parseOperand)
      const bad = operands.find(o => o.type === 'unknown')
      if (bad) { errors.push({ line: i, msg: `Unrecognized operand: ${bad.raw}` }); continue }
    }
    instructions.push({ opcode, operands, srcLine: i })
  }
  return { instructions, labels, errors }
}

export function createVM() {
  const state = {
    registers: { R0:0,R1:0,R2:0,R3:0,R4:0,R5:0,R6:0,R7:0 },
    flags: { zero:false, carry:false, negative:false, overflow:false },
    memory: new Array(256).fill(0),
    stack: [],
    pc: 0, halted: false, error: null,
    output: [], input: [],
    stepCount: 0, instructions: [], labels: {},
    cmpA: 0, cmpB: 0,
  }

  function getVal(op) {
    switch (op.type) {
      case 'reg': return state.registers[op.value]
      case 'imm': return op.value
      case 'mem': return state.memory[(state.registers[op.base] + (op.offset||0)) & 0xFF]
      case 'mem_direct': return state.memory[op.addr]
      default: throw new Error(`Cannot read: ${JSON.stringify(op)}`)
    }
  }

  function setReg(op, val, changes) {
    val = ((val % 65536) + 65536) % 65536
    if (op.type !== 'reg') throw new Error('Destination must be a register')
    const old = state.registers[op.value]
    state.registers[op.value] = val
    changes.push({ type: 'reg', name: op.value, old, val })
  }

  function setMem(addr, val, changes) {
    addr = addr & 0xFF
    val = ((val % 65536) + 65536) % 65536
    const old = state.memory[addr]
    state.memory[addr] = val
    changes.push({ type: 'mem', addr, old, val })
  }

  function updateFlags(val) {
    val = ((val % 65536) + 65536) % 65536
    state.flags.zero = val === 0
    state.flags.negative = (val & 0x8000) !== 0
  }

  function condMet(op) {
    const a = state.cmpA, b = state.cmpB
    switch (op) {
      case 'B': return true
      case 'BEQ': return a === b
      case 'BNE': return a !== b
      case 'BGT': return a > b
      case 'BLT': return a < b
      case 'BGE': return a >= b
      case 'BLE': return a <= b
    }
    return false
  }

  // For 2-or-3 operand arithmetic: if 3 ops → Rd = op(Rn, Rm); if 2 ops → Rd = op(Rd, Rm)
  function arith3(ops, fn) {
    if (ops.length >= 3) return fn(getVal(ops[1]), getVal(ops[2]))
    return fn(getVal(ops[0]), getVal(ops[1]))
  }

  function step() {
    if (state.halted || state.error) return null
    if (state.pc >= state.instructions.length) { state.halted = true; return null }
    if (state.stepCount >= MAX_STEPS) {
      state.error = 'Exceeded maximum step count (10000) — likely an infinite loop!'; return null
    }

    const instr = state.instructions[state.pc]
    const { opcode, operands: ops } = instr
    const changes = []
    let jumped = false
    state.stepCount++

    try {
      switch (opcode) {
        case 'NOP': break
        case 'XHLT': state.halted = true; break

        case 'MOV':
          setReg(ops[0], getVal(ops[1]), changes)
          break

        case 'ADD': { const r = arith3(ops, (a,b) => a+b); state.flags.carry = r > 0xFFFF; setReg(ops[0], r, changes); updateFlags(r); break }
        case 'SUB': { const r = arith3(ops, (a,b) => a-b); state.flags.carry = r < 0; setReg(ops[0], r, changes); updateFlags(r); break }
        case 'MUL': { const r = arith3(ops, (a,b) => a*b); state.flags.carry = r > 0xFFFF; setReg(ops[0], r, changes); updateFlags(r); break }
        case 'XDIV': {
          const r = arith3(ops, (a,b) => { if(b===0) throw new Error('Division by zero!'); return Math.floor(a/b) })
          setReg(ops[0], r, changes); updateFlags(r); break
        }
        case 'XMOD': {
          const r = arith3(ops, (a,b) => { if(b===0) throw new Error('Division by zero!'); return a%b })
          setReg(ops[0], r, changes); updateFlags(r); break
        }

        case 'AND': { const r = arith3(ops, (a,b) => a&b); setReg(ops[0], r, changes); updateFlags(r); break }
        case 'ORR': { const r = arith3(ops, (a,b) => a|b); setReg(ops[0], r, changes); updateFlags(r); break }
        case 'EOR': { const r = arith3(ops, (a,b) => a^b); setReg(ops[0], r, changes); updateFlags(r); break }
        case 'MVN': { const r = (~getVal(ops[1])) & 0xFFFF; setReg(ops[0], r, changes); updateFlags(r); break }
        case 'LSL': { const r = arith3(ops, (a,b) => (a << b) & 0xFFFF); setReg(ops[0], r, changes); updateFlags(r); break }
        case 'LSR': { const r = arith3(ops, (a,b) => a >>> b); setReg(ops[0], r, changes); updateFlags(r); break }

        case 'CMP': {
          state.cmpA = getVal(ops[0]); state.cmpB = getVal(ops[1])
          const d = state.cmpA - state.cmpB
          state.flags.zero = d === 0; state.flags.negative = d < 0; state.flags.carry = state.cmpA < state.cmpB
          break
        }

        case 'B': case 'BEQ': case 'BNE': case 'BGT': case 'BLT': case 'BGE': case 'BLE': {
          if (condMet(opcode)) {
            const lbl = ops[0].value
            if (state.labels[lbl] === undefined) throw new Error(`Unknown label: ${lbl}`)
            state.pc = state.labels[lbl]; jumped = true
          }
          break
        }

        case 'LDR': {
          const memOp = ops[1]
          let addr
          if (memOp.type === 'mem') addr = (state.registers[memOp.base] + (memOp.offset||0)) & 0xFF
          else if (memOp.type === 'mem_direct') addr = memOp.addr
          else throw new Error('LDR needs a memory address, e.g. [R0] or [R0, #4]')
          const v = state.memory[addr]
          setReg(ops[0], v, changes)
          changes.push({ type: 'mem_read', addr })
          break
        }
        case 'STR': {
          const memOp = ops[1]
          let addr
          if (memOp.type === 'mem') addr = (state.registers[memOp.base] + (memOp.offset||0)) & 0xFF
          else if (memOp.type === 'mem_direct') addr = memOp.addr
          else throw new Error('STR needs a memory address, e.g. [R0] or [R0, #4]')
          setMem(addr, getVal(ops[0]), changes)
          break
        }

        case 'PUSH': {
          if (state.stack.length >= 64) throw new Error('Stack overflow!')
          const v = getVal(ops[0])
          state.stack.push(v)
          changes.push({ type: 'stack_push', val: v })
          break
        }
        case 'POP': {
          if (state.stack.length === 0) throw new Error('Stack is empty!')
          setReg(ops[0], state.stack.pop(), changes)
          changes.push({ type: 'stack_pop' })
          break
        }

        case 'XOUT': {
          const v = getVal(ops[0])
          state.output.push(v)
          changes.push({ type: 'output', val: v })
          break
        }
      }
    } catch (e) {
      state.error = `Line ${instr.srcLine + 1}: ${e.message}`
      return { pc: state.pc, instr, changes, error: state.error }
    }

    if (!jumped) state.pc++
    return { pc: state.pc, instr, changes, srcLine: instr.srcLine }
  }

  function loadProgram(code) {
    const result = parse(code)
    if (result.errors.length > 0) { state.error = result.errors[0].msg; return result }
    state.instructions = result.instructions
    state.labels = result.labels
    state.pc = 0; state.halted = false; state.error = null; state.stepCount = 0
    state.output = []; state.stack = []
    state.flags = { zero:false, carry:false, negative:false, overflow:false }
    state.cmpA = 0; state.cmpB = 0
    return result
  }

  function run() {
    while (!state.halted && !state.error && state.stepCount < MAX_STEPS) step()
    if (state.stepCount >= MAX_STEPS && !state.halted && !state.error)
      state.error = 'Exceeded maximum step count (10000) — likely an infinite loop!'
  }

  function reset() {
    state.registers = { R0:0,R1:0,R2:0,R3:0,R4:0,R5:0,R6:0,R7:0 }
    state.flags = { zero:false, carry:false, negative:false, overflow:false }
    state.memory = new Array(256).fill(0)
    state.stack = []; state.pc = 0; state.halted = false; state.error = null
    state.output = []; state.input = []; state.stepCount = 0
    state.cmpA = 0; state.cmpB = 0
  }

  return { state, step, run, loadProgram, reset }
}

export function countInstructions(code) {
  let c = 0
  for (const line of code.split('\n')) {
    let l = line; const ci = l.indexOf(';'); if (ci >= 0) l = l.slice(0, ci); l = l.trim()
    if (!l) continue
    const lm = l.match(/^[a-zA-Z_]\w*\s*:(.*)$/); if (lm) { l = lm[1].trim(); if (!l) continue }
    c++
  }
  return c
}

export function validateLevel(level, vm) {
  const s = vm.state
  if (s.error) return { passed: false, msg: s.error }

  for (const tc of level.testCases) {
    vm.reset()
    if (level.initialState) {
      if (level.initialState.registers) Object.assign(vm.state.registers, level.initialState.registers)
      if (level.initialState.memory) for (const [a,v] of Object.entries(level.initialState.memory)) vm.state.memory[+a] = v
      if (level.initialState.input) vm.state.input = [...level.initialState.input]
    }
    if (tc.init) {
      if (tc.init.registers) Object.assign(vm.state.registers, tc.init.registers)
      if (tc.init.memory) for (const [a,v] of Object.entries(tc.init.memory)) vm.state.memory[+a] = v
      if (tc.init.input) vm.state.input = [...tc.init.input]
    }

    vm.run()
    if (vm.state.error) return { passed: false, msg: vm.state.error }

    const exp = tc.expected
    if (exp.registers) {
      for (const [r,v] of Object.entries(exp.registers)) {
        if (vm.state.registers[r] !== v)
          return { passed: false, msg: `${r} should be ${v}, but got ${vm.state.registers[r]}` }
      }
    }
    if (exp.memory) {
      for (const [a,v] of Object.entries(exp.memory)) {
        if (vm.state.memory[+a] !== v)
          return { passed: false, msg: `memory[${a}] should be ${v}, but got ${vm.state.memory[+a]}` }
      }
    }
    if (exp.output) {
      for (let i = 0; i < exp.output.length; i++) {
        if (vm.state.output[i] !== exp.output[i])
          return { passed: false, msg: `Output #${i+1} should be ${exp.output[i]}, but got ${vm.state.output[i] ?? 'none'}` }
      }
    }
  }
  return { passed: true, msg: 'Passed!' }
}
