import { describe, it } from 'mocha'
import { assert, expect } from 'chai'
import { once } from 'events'
import { PassThrough, Writable } from 'stream'
import { Buffer } from 'buffer'
import chalk from 'chalk'

import { quintRepl } from '../../../src/repl'
import { dedent } from '../../textUtils'
import { version } from '../../../src/version'

// A simple implementation of Writable to a string:
// After: https://bensmithgall.com/blog/jest-mock-trick
//
// Note that the writable string is used to wait for the repl's output. The
// reply is considered to be done processing the input if it emits a prompt
// prefix (>>>, ...) or an internal "Error".
class ToStringWritable extends Writable {
  buffer: string = ''
  waiter: () => void = () => {}

  _write(chunk: Buffer, _encoding: string, next: (_error?: Error | null) => void): void {
    this.buffer += chunk
    if (chunk.includes('>>> ') || chunk.includes('... ') || chunk.includes('Error')) {
      this.waiter()
    }
    next()
  }

  reset() {
    this.buffer = ''
  }

  async isReady() {
    if (this.buffer.endsWith('>>> ') || this.buffer.endsWith('... ')) {
      return
    }
    await new Promise((resolve, _reject) => {
      this.waiter = () => resolve(null)
    })
  }
}

// run a test with mocked input/output and return the input + output
const withIO = async (inputText: string): Promise<string> => {
  // save the current chalk level and reset chalk to no color
  const savedChalkLevel = chalk.level
  chalk.level = 0
  // setup:
  //  - the output that writes to a string
  //  - the input that consumes events
  const output = new ToStringWritable()
  // an input mock designed for testing
  const input = new PassThrough()
  // Pipe input to output to simulate terminal echo behavior
  // (readline doesn't echo input from non-TTY streams)
  // Use { end: false } to prevent ending output when input ends
  input.pipe(output, { end: false })

  const rl = quintRepl(input, output, { verbosity: 2, backend: 'rust' }, () => {})
  await output.isReady()

  // Send input line-by-line to the REPL. We emit 'data' events for each line,
  // with a wait between each one to give the REPL's async processing time to
  // handle each line before the next one arrives. This prevents race conditions
  // where lines are queued faster than the REPL can process them.
  const lines = inputText.split(/(?<=\n)/)
  for (const line of lines) {
    input.emit('data', line)
    await output.isReady()
  }
  input.end()
  input.unpipe(output)
  input.destroy()

  // readline is asynchronous, wait till it terminates
  await once(rl, 'close')

  chalk.level = savedChalkLevel
  // Remove trailing newline that gets added when the REPL closes
  return output.buffer.replace(/\n$/, '')
}

// the standard banner, which gets repeated
const banner = `Quint REPL ${version} (using the Rust backend)
Type ".exit" to exit, or ".help" for more information`

async function assertRepl(input: string, output: string) {
  const expected = `${banner}
${output}`

  const result = await withIO(input)
  assert(typeof result === 'string', 'expected result to be a string')
  expect(result).to.equal(expected)
}

describe('repl ok', () => {
  it('empty input', async () => {
    await assertRepl('', '>>> ')
  })

  it('Set(2 + 3)', async () => {
    const input = 'Set(2 + 3)\n'
    const output = dedent(
      `>>> Set(2 + 3)
      |Set(5)
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('Map(1 -> 2, 3 -> 4)', async () => {
    const input = 'Map(1 -> 2, 3 -> 4)\n'
    const output = dedent(
      `>>> Map(1 -> 2, 3 -> 4)
      |Map(1 -> 2, 3 -> 4)
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('basic expressions', async () => {
    const input = dedent(
      `1 + 1
      |3 > 1
      |1.to(3).map(x => 2 * x)
      |1.to(4).filter(x => x > 2)
      |Set(1, 3).union(Set(5, 6))
      |1.to(4).forall(x => x > 1)
      |(5 - 1, 5, 6)
      |[5 - 1, 5, 6]
      |`
    )
    const output = dedent(
      `>>> 1 + 1
      |2
      |>>> 3 > 1
      |true
      |>>> 1.to(3).map(x => 2 * x)
      |Set(2, 4, 6)
      |>>> 1.to(4).filter(x => x > 2)
      |Set(3, 4)
      |>>> Set(1, 3).union(Set(5, 6))
      |Set(1, 3, 5, 6)
      |>>> 1.to(4).forall(x => x > 1)
      |false
      |>>> (5 - 1, 5, 6)
      |(4, 5, 6)
      |>>> [5 - 1, 5, 6]
      |[4, 5, 6]
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('ill-typed expressions', async () => {
    const input = dedent(
      `1 + false
      |`
    )
    const output = dedent(
      `>>> 1 + false
      |static analysis error: 
      | Error [QNT000]: Couldn't unify int and bool
      |Trying to unify int and bool
      |Trying to unify (int, int) => int and (int, bool) => _t0
      |
      |
      |  at <input-0>:0:1
      |  1 + false
      |  ^^^^^^^^^
      |
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('definitions in expressions', async () => {
    const input = dedent(
      `val x = 3; 2 * x
      |def mult(x, y) = x * y; mult(2, mult(3, 4))
      |`
    )
    const output = dedent(
      `>>> val x = 3; 2 * x
      |6
      |>>> def mult(x, y) = x * y; mult(2, mult(3, 4))
      |24
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('top-level definitions', async () => {
    const input = dedent(
      `val n = 4
      |def mult(x, y) = x * y
      |mult(100, n)
      |def powpow(x, y) = x^y
      |mult(100, powpow(2, 3))
      |`
    )
    const output = dedent(
      `>>> val n = 4
      |
      |>>> def mult(x, y) = x * y
      |
      |>>> mult(100, n)
      |400
      |>>> def powpow(x, y) = x^y
      |
      |>>> mult(100, powpow(2, 3))
      |800
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('clear history', async () => {
    const input = dedent(
      `val n = 4
      |n * n
      |.clear
      |n * n
      |`
    )
    const output = dedent(
      `>>> val n = 4
      |
      |>>> n * n
      |16
      |>>> .clear
      |
      |>>> n * n
      |static analysis error: 
      | Error [QNT404]: Name 'n' not found
      |
      |  at <input-0>:0:1
      |  n * n
      |  ^
      |
      |static analysis error: 
      | Error [QNT404]: Name 'n' not found
      |
      |  at <input-0>:0:5
      |  n * n
      |      ^
      |
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('clear resets evaluator state', async () => {
    const input = dedent(
      `var x: int
      |x' = 0
      |x' = x + 1
      |x
      |.clear
      |var x: int
      |x' = 42
      |x
      |`
    )
    const output = dedent(
      `>>> var x: int
      |
      |>>> x' = 0
      |true
      |{ x: 0 }
      |>>> x' = x + 1
      |true
      |{ x: 0 => 1 }
      |>>> x
      |1
      |>>> .clear
      |
      |>>> var x: int
      |
      |>>> x' = 42
      |true
      |{ x: 42 }
      |>>> x
      |42
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('state diffs suppressed at verbosity 1', async () => {
    const input = dedent(
      `var x: int
      |action Init = x' = 0
      |action Next = x' = x + 1
      |Init
      |.verbosity=1
      |Next
      |x
      |`
    )
    const output = dedent(
      `>>> var x: int
      |
      |>>> action Init = x' = 0
      |
      |>>> action Next = x' = x + 1
      |
      |>>> Init
      |true
      |{ x: 0 }
      |>>> .verbosity=1
      |.verbosity=1
      |>>> Next
      |true
      |>>> x
      |1
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug output suppressed below verbosity 3', async () => {
    const input = dedent(
      `var x: int
      |x' = 0
      |q::debug("x value", x)
      |`
    )
    const output = dedent(
      `>>> var x: int
      |
      |>>> x' = 0
      |true
      |{ x: 0 }
      |>>> q::debug("x value", x)
      |0
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug output with verbosity=3', async () => {
    const input = dedent(
      `pure def plus(x, y) = x + y
      |.verbosity=3
      |q::debug("result", plus(2, 3))
      |`
    )
    const output = dedent(
      `>>> pure def plus(x, y) = x + y
      |
      |>>> .verbosity=3
      |.verbosity=3
      |>>> q::debug("result", plus(2, 3))
      |[DEBUG] result 5
      |5
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug output on failure', async () => {
    const input = dedent(
      `pure def div(x, y) = x / y
      |.verbosity=3
      |q::debug("dividing", div(2, 0))
      |`
    )
    const output = dedent(
      `>>> pure def div(x, y) = x / y
      |
      |>>> .verbosity=3
      |.verbosity=3
      |>>> q::debug("dividing", div(2, 0))
      |runtime error: 
      | Error [QNT503]: Division by zero
      |
      |  at <input-0>:0:22
      |  pure def div(x, y) = x / y
      |                       ^^^^^
      |    at div(2, 0) (<input-1>:0:22)
      |
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug with stateful actions', async () => {
    const input = dedent(
      `var x: int
      |.verbosity=3
      |x' = q::debug("init x", 0)
      |action step = x' = q::debug("stepping", x + 1)
      |step
      |`
    )
    const output = dedent(
      `>>> var x: int
      |
      |>>> .verbosity=3
      |.verbosity=3
      |>>> x' = q::debug("init x", 0)
      |[DEBUG] init x 0
      |true
      |{ x: 0 }
      |>>> action step = x' = q::debug("stepping", x + 1)
      |
      |>>> step
      |[DEBUG] stepping 1
      |true
      |{ x: 0 => 1 }
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug output (initial verbosity 3)', async () => {
    const input = dedent(
      `.verbosity=3
      |pure def plus(x, y) = x + y
      |q::debug("result", plus(2, 3))
      |`
    )
    const output = dedent(
      `>>> .verbosity=3
      |.verbosity=3
      |>>> pure def plus(x, y) = x + y
      |
      |>>> q::debug("result", plus(2, 3))
      |[DEBUG] result 5
      |5
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug returns its value', async () => {
    const input = dedent(
      `.verbosity=3
      |q::debug("check", 1 == 1)
      |q::debug("is even", 4 % 2 == 0)
      |`
    )
    const output = dedent(
      `>>> .verbosity=3
      |.verbosity=3
      |>>> q::debug("check", 1 == 1)
      |[DEBUG] check true
      |true
      |>>> q::debug("is even", 4 % 2 == 0)
      |[DEBUG] is even true
      |true
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug with stateful actions (initial verbosity 3)', async () => {
    const input = dedent(
      `.verbosity=3
      |var x: int
      |x' = q::debug("init x", 0)
      |action step = x' = q::debug("stepping", x + 1)
      |step
      |`
    )
    const output = dedent(
      `>>> .verbosity=3
      |.verbosity=3
      |>>> var x: int
      |
      |>>> x' = q::debug("init x", 0)
      |[DEBUG] init x 0
      |true
      |{ x: 0 }
      |>>> action step = x' = q::debug("stepping", x + 1)
      |
      |>>> step
      |[DEBUG] stepping 1
      |true
      |{ x: 0 => 1 }
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('set and get the seed', async () => {
    const input = dedent(
      `.seed=4
      |.seed
      |.seed=0x1abc
      |`
    )
    const output = dedent(
      `>>> .seed=4
      |.seed=4
      |>>> .seed
      |.seed=4
      |>>> .seed=0x1abc
      |.seed=6844
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('handle exceptions', async () => {
    const input = dedent(
      `Set(Int)
      |`
    )
    const output = dedent(
      `>>> Set(Int)
      |runtime error: 
      | Error [QNT501]: Infinite set Int is non-enumerable
      |
      |  at <input-0>:0:1
      |  Set(Int)
      |  ^^^^^^^^
      |
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('assignments', async () => {
    const input = dedent(
      `var x: int
      |action Init = x' = 0
      |action Next = x' = x + 1
      |Init
      |x
      |Next
      |x
      |Next
      |x
      |`
    )
    const output = dedent(
      `>>> var x: int
      |
      |>>> action Init = x' = 0
      |
      |>>> action Next = x' = x + 1
      |
      |>>> Init
      |true
      |{ x: 0 }
      |>>> x
      |0
      |>>> Next
      |true
      |{ x: 0 => 1 }
      |>>> x
      |1
      |>>> Next
      |true
      |{ x: 1 => 2 }
      |>>> x
      |2
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('action-level disjunctions and conjunctions', async () => {
    const input = dedent(
      `
      |var x: int
      |action Init = x' = 0
      |action Next = any {
      |  all {
      |    x == 0,
      |    x' = 1,
      |  },
      |  all {
      |    x == 1,
      |    x' = 0,
      |  },
      |}
      |
      |Init
      |x
      |Next
      |x
      |Next
      |x
      |Next
      |x
      |`
    )
    const output = dedent(
      `>>> 
      |>>> var x: int
      |
      |>>> action Init = x' = 0
      |
      |>>> action Next = any {
      |...   all {
      |...     x == 0,
      |...     x' = 1,
      |...   },
      |...   all {
      |...     x == 1,
      |...     x' = 0,
      |...   },
      |... }
      |... 
      |
      |>>> Init
      |true
      |{ x: 0 }
      |>>> x
      |0
      |>>> Next
      |true
      |{ x: 0 => 1 }
      |>>> x
      |1
      |>>> Next
      |true
      |{ x: 1 => 0 }
      |>>> x
      |0
      |>>> Next
      |true
      |{ x: 0 => 1 }
      |>>> x
      |1
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('action-level disjunctions and non-determinism', async () => {
    const input = dedent(
      `
      |var x: int
      |action Init = x' = 0
      |action Next = any {
      |  x' = x + 1,
      |  x' = x - 1,
      |}
      |
      |.seed=42
      |Init
      |-1 <= x and x <= 1
      |Next
      |-2 <= x and x <= 2
      |Next
      |-3 <= x and x <= 3
      |Next
      |-4 <= x and x <= 4
      |`
    )
    const output = dedent(
      `>>> 
      |>>> var x: int
      |
      |>>> action Init = x' = 0
      |
      |>>> action Next = any {
      |...   x' = x + 1,
      |...   x' = x - 1,
      |... }
      |... 
      |
      |>>> .seed=42
      |.seed=42
      |>>> Init
      |true
      |{ x: 0 }
      |>>> -1 <= x and x <= 1
      |true
      |>>> Next
      |true
      |{ x: 0 => -1 }
      |>>> -2 <= x and x <= 2
      |true
      |>>> Next
      |true
      |{ x: -1 => -2 }
      |>>> -3 <= x and x <= 3
      |true
      |>>> Next
      |true
      |{ x: -2 => -1 }
      |>>> -4 <= x and x <= 4
      |true
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('nondet and oneOf', async () => {
    const input = dedent(
      `
      |var x: int
      |.seed=42
      |
      |x' = 0
      |x == 0
      |{ nondet y = oneOf(Set(1, 2, 3))
      |  x' = y }
      |
      |1 <= x and x <= 3
      |nondet y = oneOf(2.to(5)); x' = y
      |2 <= x and x <= 5
      |nondet t = oneOf(tuples(2.to(5), 3.to(4))); x' = t._1 + t._2
      |5 <= x and x <= 9
      |nondet i = oneOf(1.to(100)); x' = i
      |x >= 0
      |nondet i = oneOf(0.to(100)); x' = i
      |Int.contains(x)
      |nondet m = 1.to(5).setOfMaps(1.to(10)).oneOf(); x' = m.get(3)
      |x.in(Int)
      |nondet m = Set().oneOf(); x' = m
      |`
    )
    const output = dedent(
      `>>> 
      |>>> var x: int
      |
      |>>> .seed=42
      |.seed=42
      |>>> 
      |>>> x' = 0
      |true
      |{ x: 0 }
      |>>> x == 0
      |true
      |>>> { nondet y = oneOf(Set(1, 2, 3))
      |...   x' = y }
      |... 
      |true
      |{ x: 0 => 2 }
      |>>> 1 <= x and x <= 3
      |true
      |>>> nondet y = oneOf(2.to(5)); x' = y
      |true
      |{ x: 2 => 5 }
      |>>> 2 <= x and x <= 5
      |true
      |>>> nondet t = oneOf(tuples(2.to(5), 3.to(4))); x' = t._1 + t._2
      |true
      |{ x: 5 => 6 }
      |>>> 5 <= x and x <= 9
      |true
      |>>> nondet i = oneOf(1.to(100)); x' = i
      |true
      |{ x: 6 => 62 }
      |>>> x >= 0
      |true
      |>>> nondet i = oneOf(0.to(100)); x' = i
      |true
      |{ x: 62 => 45 }
      |>>> Int.contains(x)
      |true
      |>>> nondet m = 1.to(5).setOfMaps(1.to(10)).oneOf(); x' = m.get(3)
      |true
      |{ x: 45 => 8 }
      |>>> x.in(Int)
      |true
      |>>> nondet m = Set().oneOf(); x' = m
      |false
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('nondet and oneOf over sets of sets', async () => {
    const input = dedent(
      `var S: Set[int]
      |.seed=42
      |nondet y = oneOf(powerset(1.to(3))); S' = y
      |S.subseteq(1.to(3))
      |`
    )
    const output = dedent(
      `>>> var S: Set[int]
      |
      |>>> .seed=42
      |.seed=42
      |>>> nondet y = oneOf(powerset(1.to(3))); S' = y
      |true
      |{ S: Set(1) }
      |>>> S.subseteq(1.to(3))
      |true
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug with chained actions', async () => {
    const input = dedent(
      `var n: int
      |action init = n' = q::debug("init", 0)
      |action step = n' = q::debug("step", n + 1)
      |.verbosity=3
      |init.then(step).then(step)
      |`
    )
    const output = dedent(
      `>>> var n: int
      |
      |>>> action init = n' = q::debug("init", 0)
      |
      |>>> action step = n' = q::debug("step", n + 1)
      |
      |>>> .verbosity=3
      |.verbosity=3
      |>>> init.then(step).then(step)
      |[DEBUG] init 0
      |[DEBUG] step 1
      |[DEBUG] step 2
      |true
      |{ n: 1 => 2 }
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug with expect (condition passes)', async () => {
    const input = dedent(
      `var n: int
      |action init = n' = q::debug("init", 0)
      |action step = n' = q::debug("step", n + 1)
      |.verbosity=3
      |init.then(step).expect(q::debug("check", n) == 1)
      |`
    )
    const output = dedent(
      `>>> var n: int
      |
      |>>> action init = n' = q::debug("init", 0)
      |
      |>>> action step = n' = q::debug("step", n + 1)
      |
      |>>> .verbosity=3
      |.verbosity=3
      |>>> init.then(step).expect(q::debug("check", n) == 1)
      |[DEBUG] init 0
      |[DEBUG] step 1
      |[DEBUG] check 1
      |true
      |{ n: 1 }
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('diagnostics: debug with multiple calls in single expression (initial verbosity 3)', async () => {
    // Uses arithmetic to trigger two q::debug calls in one evaluation.
    // Both diagnostics are preserved since no state shift happens between them.
    const input = dedent(
      `.verbosity=3
      |q::debug("a", 1) + q::debug("b", 2)
      |`
    )
    const output = dedent(
      `>>> .verbosity=3
      |.verbosity=3
      |>>> q::debug("a", 1) + q::debug("b", 2)
      |[DEBUG] a 1
      |[DEBUG] b 2
      |3
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('run q::test, q::testOnce, and q::lastTrace', async () => {
    const input = dedent(
      `
      |var n: int
      |action Init = n' = 0
      |action Next = n' = n + 1
      |val Inv = n < 10
      |q::testOnce(5, 1, Init, Next, Inv)
      |q::testOnce(10, 1, Init, Next, Inv)
      |q::test(5, 5, 1, Init, Next, Inv)
      |q::test(5, 10, 1, Init, Next, Inv)
      |q::lastTrace.length()
      |q::lastTrace.nth(q::lastTrace.length() - 1)
      |`
    )
    const output = dedent(
      `>>> 
      |>>> var n: int
      |
      |>>> action Init = n' = 0
      |
      |>>> action Next = n' = n + 1
      |
      |>>> val Inv = n < 10
      |
      |>>> q::testOnce(5, 1, Init, Next, Inv)
      |"ok"
      |>>> q::testOnce(10, 1, Init, Next, Inv)
      |"violation"
      |>>> q::test(5, 5, 1, Init, Next, Inv)
      |"ok"
      |>>> q::test(5, 10, 1, Init, Next, Inv)
      |"violation"
      |>>> q::lastTrace.length()
      |11
      |>>> q::lastTrace.nth(q::lastTrace.length() - 1)
      |{ n: 10 }
      |>>> `
    )
    await assertRepl(input, output)
  })

  it('REPL consumes its output', async () => {
    const input = dedent(
      `>>> 1 + 1
      |
      |>>> all {
      |...   true,
      |...   true
      |... }
      |
      | >>> if (true) {
      | ...   3
      | ... } else {
      | ...   4
      | ... }
      |
      |>>> Set(2 + 3)
      |Set(5)
      |   // a multiline comment
      | /// a doc comment
      | /* a multiline
      |    comment
      |  */
      |`
    )
    const output = dedent(
      `>>> >>> 1 + 1
      |... 
      |2
      |>>> >>> all {
      |... ...   true,
      |... ...   true
      |... ... }
      |... 
      |true
      |>>>  >>> if (true) {
      |...  ...   3
      |...  ... } else {
      |...  ...   4
      |...  ... }
      |... 
      |3
      |>>> >>> Set(2 + 3)
      |... Set(5)
      |Set(5)
      |>>>    // a multiline comment
      |>>>  /// a doc comment
      |>>>  /* a multiline
      |...     comment
      |...   */
      |... `
    )
    await assertRepl(input, output)
  })
})