use "collections"
use "files"
use "format"
use "itertools"

interface Value
  """
  Values referred to by opcodes in instructions should all be things
  that implement this interface
  """
  fun value(): I64


class Position
  """
  A positional operand (i.e. one whose value is fetched from a memory cell
  """
  let _addr: I64
  let _value: I64

  new create(__addr: I64, __value: I64) =>
    _addr = __addr
    _value = __value

  fun addr(): I64 => _addr
  fun value(): I64 => _value

class Immediate
  """
  An immediate operand (i.e. one whose value is a verbatim argument)
  """
  let _value: I64

  new create(__value: I64) =>
    _value = __value

  fun value(): I64 => _value

class Operands
  """
  The operands to a typical binary operator: two values and a target
  that will always be in positional (i.e. non-immediate) mode
  """
  let _lhs: Value
  let _rhs: Value
  let _tgt: I64

  new create(__lhs: Value, __rhs: Value, __tgt: I64) =>
    _lhs = __lhs
    _rhs = __rhs
    _tgt = __tgt

  fun val1(): I64 => _lhs.value()
  fun val2(): I64 => _rhs.value()
  fun tgt(): USize => _tgt.usize()

trait Error
  """
  The errors we can raise will all implement this so we can give
  nice informative error messages
  """
  fun message(): String

class val OutOfBounds is Error
  let _loc: USize
  new val create(loc: USize) => _loc = loc
  fun message(): String =>
    "Memory access out of bounds: " + Format.int[USize](_loc)

class val UnknownOperand is Error
  let _op: I64
  new val create(op: I64) => _op = op
  fun message(): String =>
    "Unknown operation: " + Format.int[I64](_op)

class val NotEnoughInput is Error
  let _pc: USize
  new val create(pc: USize) => _pc = pc
  fun message(): String =>
    "Not enough input when running instruction at " + Format.int[USize](_pc)


// I should wrap this but I haven't, so, uh, yeah
type Opcode is I64
// a Mode is true if immediate, false if positional
type Modes is (Bool, Bool, Bool)


class val OutputRequest
  let _o: I64
  new val create(_o': I64) => _o = _o'
  fun output(): I64 => _o
  fun string(): String =>
    "OutputRequest(" + _o.string() + ")"
primitive InputRequest
  fun string(): String => "InputRequest"
primitive Halted
  fun string(): String => "Halted"

type ProgramResult is (OutputRequest | InputRequest | Halted)


class Program
  """
  All the logic we need for running programs
  """

  // an array of memory cells
  var storage: Array[I64]
  // the current instruction we're running
  var pc: USize = 0
  //
  var input: Array[I64] = Array[I64]()

  // a list of errors that may or may not have happened. (right now,
  // we'll only ever have one, but it's worth planning ahead, eh?)
  var errors: Array[Error val] = Array[Error val]()

  // a queue of debug messages
  var debug: Array[String] = Array[String]()

  new create(s: Array[I64]) =>
    """
    Create a Program from a verbatim storage and input array
    """
    storage = s

  new from_file(file: File) =>
    """
    Create a Program by walking input from a file. This just skips any
    errors if they happen (e.g. if there are non-numbers), which
    might not be what we want!
    """
    storage = Array[I64](0)
    for line in file.lines() do
      for chunk in line.split(",").values() do
        try
          storage.push(chunk.i64()?)
        end
      end
    end

  fun ref send_input(i: I64) =>
    input.push(i)

  fun print(env: Env) =>
    """
    Print out the current state of memory
    """
    env.out.write("[")
    for x in storage.values() do
      env.out.write(Format.int[I64](x))
      env.out.write(" ")
    end
    env.out.print("]")

  fun ref dbg(msg: String) =>
    """
    Log a debug message
    """
    debug.push(msg)


  fun ref read_storage(addr: USize): I64 ? =>
    """
    Attempt to read a word from storage, logging an error message
    if we fail to do so
    """
    try
      storage(addr)?
    else
      errors.push(OutOfBounds(addr))
      error
    end

  fun ref write_storage(addr: USize, value: I64) ? =>
    """
    Attempt to write a word to storage, logging an error message
    if we fail to do so
    """
    try
      storage(addr)? = value
    else
      errors.push(OutOfBounds(addr))
      error
    end


  fun ref get_immediate(addr_loc: USize): Immediate ? =>
    """
    Read a single immediate value at the provided loc
    """
    let value = read_storage(addr_loc)?
    Immediate(value)

  fun ref get_position(addr_loc: USize): Position ? =>
    """
    Read a positional value whose address is stored at the provided loc
    """
    let addr = read_storage(addr_loc)?
    let value = read_storage(addr.usize())?
    Position(addr, value)

  fun ref get_value(addr_loc: USize, mode: Bool): Value ? =>
    """
    Read a value at the provided loc based on the mode: if true, then
    treat it as an immediate, otherwise treat it as a positional
    """
    if mode then
      get_immediate(addr_loc)?
    else
      get_position(addr_loc)?
    end


  fun ref get_op_data(modes: Modes): Operands ? =>
    """
    Get the data that a binary op needs based on the provided modes
    """
    (let lm, let rm, _) = modes
    Operands(
      get_value(pc+1, lm)?,
      get_value(pc+2, rm)?,
      read_storage(pc+3)?
    )

  fun ref parse_op(n: I64): (Opcode, Modes) =>
    """
    Parse out an opcode and set of modes. This hard-codes modes
    for three operands; it might be worth it in the future to generalize
    to any N operands
    """
    let opcode = n % 100
    let mode1 =   (n / 100) % 10
    let mode2 =  (n / 1000) % 10
    let mode3 = (n / 10000) % 10
    (opcode, (mode1 == 1, mode2 == 1, mode3 == 1))


  fun ref do_binop(modes: Modes, op: {(I64, I64): I64}, op_name: String) ? =>
    """
    Get the values of the first two operands by the provided modes, run
    the function `op` on them, and store the result of that at the
    provided location
    """
    let data = get_op_data(modes)?
    dbg("    writing " + data.val1().string() + op_name + data.val2().string()
      + " = " + op(data.val1(), data.val2()).string() + " to *" + data.tgt().string())
    write_storage(data.tgt(), op(data.val1(), data.val2()))?
    pc = pc + 4

  fun ref do_jump(modes: Modes, jump_cond: {(I64): Bool}) ? =>
    """
    Get the value indicated by the modes at [pc+1], and the function
    `jump_cond` returns true on that value, then jump to the location
    indicated by [pc+2], otherwise move ahead as usual
    """
    (let cond, let tgt, _) = modes
    if jump_cond(get_value(pc+1, cond)?.value()) then
      pc = get_value(pc+2, tgt)?.value().usize()
    else
      pc = pc + 3
    end

  fun ref do_input(modes: Modes) ? =>
    """
    Read a value from the input and put it into storage. This logs an
    error if there are not enough values in the input queue.
    """
    let i = try
      input.shift()?
    else
      errors.push(NotEnoughInput(pc))
      error
    end
    let addr = read_storage(pc+1)?.usize()
    dbg("    writing " + i.string() + " to *" + addr.string())
    write_storage(addr.usize(), i)?
    pc = pc + 2

  fun ref do_output(modes: Modes): OutputRequest ? =>
    """
    Write the value indicated according to the mode to the output queue.
    """
    (let m, _, _) = modes
    let r = get_value(pc+1, m)?.value()
    pc = pc + 2
    OutputRequest(r)

  fun ref at_pc(): USize => pc

  fun ref fmt_op(opcode: I64, modes: Modes): String =>
    (let a, let b, let c) = modes
    let op1 = try read_storage(pc+1)? else -99 end
    let op2 = try read_storage(pc+2)? else -99 end
    let op3 = try read_storage(pc+3)? else -99 end
    let arg1 = if not a then "*" else "" end + op1.string()
    let arg2 = if not b then "*" else "" end + op2.string()
    let arg3 = if not c then "*" else "" end + op3.string()
    match opcode
      | 1 => "add " + arg1 + " " + arg2 + " -> " + arg3
      | 2 => "sub " + arg1 + " " + arg2 + " -> " + arg3
      | 3 => "input " + arg1
      | 4 => "output " + arg1
      | 5 => "jnz " + arg1 + " to " + arg2
      | 6 => "jez " + arg1 + " to " + arg2
      | 7 => "lt " + arg1 + " " + arg2 + " -> " + arg3
      | 8 => "eq " + arg1 + " " + arg2 + " -> " + arg3
      | 99 => "halt"
    else
      "unknown op(" + opcode.string() + ")"
    end

  fun ref run_loop(): ProgramResult ? =>
    """
    This continues running the VM until it encounters a halt instruction
    or it encounters some other error
    """

    while true do
      (let opcode, let modes) = parse_op(read_storage(pc)?)
      dbg("  " + pc.string() + " -> " + fmt_op(opcode, modes))
      match opcode
        | 1 => do_binop(modes, {(x, y) => x + y}, "+")?
        | 2 => do_binop(modes, {(x, y) => x * y}, "*")?
        | 3 => try do_input(modes)? else return InputRequest end
        | 4 => return do_output(modes)?
        | 5 => do_jump(modes, {(x) => x != 0})?
        | 6 => do_jump(modes, {(x) => x == 0})?
        | 7 => do_binop(modes, {(x, y) => if x < y then 1 else 0 end}, "<")?
        | 8 => do_binop(modes, {(x, y) => if x == y then 1 else 0 end}, "==")?
        | 99 => return Halted
        | let x: I64 =>
          errors.push(UnknownOperand(x))
          error
      end
    end
    Halted

  fun ref run_with(env: Env, i: Array[I64]): Array[I64] =>
    """
    This runs the VM and prints output: if an error occurs, it prints
    the error message(s) and any debug messages that happened during
    program execution, and then it prints everything printed via an
    output instruction.
    """
    input.append(i)

    var output = Array[I64]()
    var is_done = false
    try
      repeat
        match run_loop()?
          | Halted => is_done = true
          | InputRequest => env.err.print("Not enough input for program"); is_done = true
          | let o: OutputRequest => output.push(o.output())
        end
      until is_done end
    else
      for err in errors.values() do
        env.out.print(err.message())
      end
      for n in debug.values() do
        env.out.print("[debug] " + n)
      end
    end
    for n in output.values() do
      env.out.print("> " + Format.int[I64](n))
    end
    output

  fun ref run_step(env: Env): ProgramResult =>
    try
      env.out.print("restarting from " + pc.string())
      let r = run_loop()?
      for msg in debug.values() do
        env.out.print(msg)
      end
      debug.clear()
      r
    else
      for err in errors.values() do
        env.out.print(err.message())
      end
      Halted
    end

  fun ref clone(): Program =>
    Program(storage.clone())

actor Main
  new create(env: Env) =>
    let caps = recover val FileCaps.>set(FileRead).>set(FileStat) end

    let target_file = try env.args(1)? else "input.txt" end
    try
      with
        file = OpenFile(FilePath(env.root as AmbientAuth, target_file, caps)?) as File
      do
        let program = Program.from_file(file)
        var max: (I64 | None) = None
        var choice = Array[I64]()
        for p in permutations([5;6;7;8;9]).values() do
          env.out.print(".")
          let res = test_phases_pt2(env, program, p)
          if try (res as I64) > (max as I64) else true end then
            max = res
            choice = p
          end
        end
        env.out.print("maximum is " + max.string())
        env.out.write("permutation was")
        for i in choice.values() do
          env.out.write(" " + i.string())
        end
        env.out.print("")
      end
    else
      // if something failed, then print an error message of some kind and exit
      env.err.print("Couldn't read expected file `input.txt'")
      env.exitcode(99)
    end

  fun test_phases_pt1(env: Env, program: Program, phases: Array[I64]): (I64 | None) =>
    var input: I64 = 0
    for i in phases.values() do
      let output = program.clone().run_with(env, [i;input])
      input = try output(0)? else
        env.out.print("Not enough output returned")
        return None
      end
      None
    end
    input

  fun test_phases_pt2(env: Env, program: Program, phases: Array[I64]): (I64 | None) =>
    var input: (I64 | None) = 0
    var programs = Array[(I64, Program)]()
    var n: I64 = 0
    for i in phases.values() do
      let p = program.clone()
      env.out.print("first step on " + n.string() + " produces " + p.run_step(env).string())
      env.out.print("  (pc = " + p.at_pc().string() + ")")
      env.out.print("sending input " + i.string())
      p.send_input(i)
      env.out.print("second step on " + n.string() + " produces " + p.run_step(env).string())
      env.out.print("  (pc = " + p.at_pc().string() + ")")
      programs.push((n, p))
      n = n + 1
    end

    var programs_left = phases.size()
    while programs_left > 0 do
      for (n', p) in programs.values() do
        env.out.print("running program " + n'.string() + " at " + p.at_pc().string() + " with " + input.string() + ": ")
        try
          p.send_input((input = None) as I64)
        else
          env.out.print("used input more than once!")
          return None
        end
        var running = true
        repeat
          let r = p.run_step(env)
          match r
            | Halted => programs_left = programs_left - 1; running = false
            | InputRequest => running = false
            | let o: OutputRequest =>
            input = o.output()
            env.out.print("produced " + input.string() + ", ")
          end
        until not running end
        env.out.print("")
      end
    end
    input

  fun permutations(array: Array[I64]): Array[Array[I64]] =>
    if array.size() == 0 then
      [[]]
    else
      let res = Array[Array[I64]]()
      for a in array.values() do
        let sm = Array[I64]()
        Iter[I64](array.values()).filter({(x: I64) => x != a}).collect(sm)
        let ps = permutations(sm)
        for p in ps.values() do
          p.unshift(a)
          res.push(p)
        end
      end
      res
    end