import test from 'node:test'
import assert from 'node:assert/strict'
import * as safe from '../src/safeio.js'
import * as io from '../src/io.js'
import { freeze as deepFreeze } from './helpers.js'

const { pure, effect, map, flatMap, sequence, step, drive, traceStep, trace } = safe

const r1 = deepFreeze({ tag: 'Rect', x: 0, y: 0, w: 1, h: 1, color: '#000000' })
const t1 = deepFreeze({ tag: 'Text', x: 1, y: 2, text: 'hi', color: '#ffffff', size: 16, align: 'left' })
const drawOp = (shape) => ({ tag: 'Draw', shape })
const seedOp = { tag: 'RandomSeed' }
const boom = () => { throw new Error('boom') }
const N = 100000

// --- building -------------------------------------------------------------------------------------------------

test('m07b pure builds a Pure', () => assert.deepStrictEqual(pure(3), { tag: 'Pure', value: 3 }))

test('m07b effect(op) is an Effect whose continuation is pure', () => {
  const p = effect(seedOp)
  assert.equal(p.tag, 'Effect')
  assert.deepStrictEqual(p.op, seedOp)
  assert.deepStrictEqual(p.k(5), pure(5))
})

test('m07b flatMap only builds a FlatMap node', () => {
  const m = pure(1)
  const f = (n) => pure(n + 1)
  const p = flatMap(m, f)
  assert.equal(p.tag, 'FlatMap')
  assert.equal(p.io, m)
  assert.equal(p.f, f)
})

test('m07b building flatMap(pure(1), throwingF) does not throw', () =>
  assert.doesNotThrow(() => flatMap(pure(1), boom)))

test('m07b building map(io, throwingF) does not throw', () =>
  assert.doesNotThrow(() => map(effect(seedOp), boom)))

test('m07b tracing map(io, throwingF) runs f and throws', () =>
  assert.throws(() => trace(map(effect(seedOp), boom), 0), /boom/))

// --- step: the three rewrite rules ----------------------------------------------------------------------------

test('m07b step leaves a Pure unchanged', () => {
  const p = pure(1)
  assert.equal(step(p), p)
})

test('m07b step leaves an Effect unchanged', () => {
  const p = effect(seedOp)
  assert.equal(step(p), p)
})

test('m07b step of FlatMap(Pure(a), f) is f(a)', () =>
  assert.deepStrictEqual(step(flatMap(pure(3), (n) => pure(n * 2))), pure(6)))

test('m07b step of FlatMap(Effect(op, k), f) is an Effect with the same op that continues with FlatMap(k(r), f)', () => {
  const f = (n) => pure(n * 2)
  const s = step(flatMap(effect(seedOp), f))
  assert.equal(s.tag, 'Effect')
  assert.deepStrictEqual(s.op, seedOp)
  const next = s.k(5)
  assert.equal(next.tag, 'FlatMap')
  assert.deepStrictEqual(next.io, pure(5))
  assert.equal(next.f, f)
})

test('m07b step re-associates FlatMap(FlatMap(m, g), f) into FlatMap(m, x => FlatMap(g(x), f))', () => {
  const m = effect(seedOp)
  const g = (n) => pure(n + 1)
  const f = (n) => pure(n * 2)
  const s = step(flatMap(flatMap(m, g), f))
  assert.equal(s.tag, 'FlatMap')
  assert.equal(s.io, m)
  const inner = s.f(4)
  assert.equal(inner.tag, 'FlatMap')
  assert.deepStrictEqual(inner.io, pure(5))
  assert.equal(inner.f, f)
})

test('m07b step does not call the functions it re-associates', () =>
  assert.doesNotThrow(() => step(flatMap(flatMap(effect(seedOp), boom), boom))))

test('m07b step of a left-nested chain is one rewrite, not a whole walk', () => {
  // After one step, the top is still a FlatMap: step did constant work, the rest is left for later steps.
  const chain = Array.from({ length: N }).reduce((acc) => flatMap(acc, (n) => pure(n + 1)), pure(0))
  assert.equal(step(chain).tag, 'FlatMap')
})

// --- the interpreter ------------------------------------------------------------------------------------------

test('m07b drive collects every output in order over 200 000 moves', () => {
  const [last, outs] = drive(0, (n) => n === 200000, (n) => [n + 1, n % 50000 === 0 ? [n] : []])
  assert.equal(last, 200000)
  assert.deepStrictEqual(outs, [0, 50000, 100000, 150000])
})

test('m07b drive of a state that is already done makes no move', () =>
  assert.deepStrictEqual(drive(7, () => true, boom), [7, []]))

test('m07b traceStep answers an Effect and records its op', () => {
  const [next, ops] = traceStep(effect(seedOp), 9)
  assert.deepStrictEqual(next, pure(9))
  assert.deepStrictEqual(ops, [seedOp])
})

test('m07b traceStep rewrites a FlatMap and records nothing', () => {
  const [next, ops] = traceStep(flatMap(pure(3), (n) => pure(n * 2)), 0)
  assert.deepStrictEqual(next, pure(6))
  assert.deepStrictEqual(ops, [])
})

test('m07b traceStep leaves a Pure unchanged', () => {
  const p = pure(1)
  assert.deepStrictEqual(traceStep(p, 0), [p, []])
})

test('m07b trace of pure(3) has result 3 and no ops', () =>
  assert.deepStrictEqual(trace(pure(3), 0), { result: 3, ops: [] }))

test('m07b trace answers RandomSeed with the seed and Draw with undefined', () =>
  assert.deepStrictEqual(
    trace(flatMap(effect(seedOp), (s) => map(effect(drawOp(r1)), () => s)), 9),
    { result: 9, ops: [seedOp, drawOp(r1)] }))

test('m07b trace refuses operations of later modules', () =>
  assert.throws(() => trace(effect({ tag: 'LoadBest' }), 0), { message: "SafeIO.trace answers only Module 7's operations" }))

test('m07b sequence of nothing has no ops', () =>
  assert.deepStrictEqual(trace(sequence([]), 0), { result: undefined, ops: [] }))

test('m07b sequence runs effects left to right', () =>
  assert.deepStrictEqual(
    trace(sequence([effect(seedOp), effect(drawOp(t1)), effect(drawOp(r1))]), 0).ops,
    [seedOp, drawOp(t1), drawOp(r1)]))

// --- stack safety ---------------------------------------------------------------------------------------------

const leftDraws = (n) =>
  Array.from({ length: n }).reduce((acc, _, i) => flatMap(acc, () => effect(drawOp({ ...r1, x: i }))), pure(undefined))
const rightDraws = (n) =>
  Array.from({ length: n }).reduceRight((rest, _, i) => flatMap(effect(drawOp({ ...r1, x: i })), () => rest), pure(undefined))

test('m07b 100 000 left-nested flatMaps of draws run, in order', () => {
  const { result, ops } = trace(leftDraws(N), 0)
  assert.equal(result, undefined)
  assert.equal(ops.length, N)
  assert.deepStrictEqual([ops[0], ops[N - 1]], [drawOp({ ...r1, x: 0 }), drawOp({ ...r1, x: N - 1 })])
})

test('m07b 100 000 right-nested flatMaps of draws run, in order', () => {
  const { ops } = trace(rightDraws(N), 0)
  assert.equal(ops.length, N)
  assert.deepStrictEqual([ops[0], ops[N - 1]], [drawOp({ ...r1, x: 0 }), drawOp({ ...r1, x: N - 1 })])
})

test('m07b 100 000 left-nested maps run without any op', () =>
  assert.deepStrictEqual(
    trace(Array.from({ length: N }).reduce((acc) => map(acc, (n) => n + 1), pure(0)), 0),
    { result: N, ops: [] }))

test('m07b sequence of 100 000 draws runs', () =>
  assert.equal(trace(sequence(Array.from({ length: N }, (_, i) => effect(drawOp({ ...r1, x: i })))), 0).ops.length, N))

test('m07b a while loop over step is a complete interpreter for 100 000 left-nested flatMaps', () => {
  // What an impure runtime would do (a loop is allowed there): rewrite until the top is Pure or Effect, perform.
  const performed = []
  let current = leftDraws(N)
  while (current.tag !== 'Pure') {
    current = step(current)
    if (current.tag === 'Effect') {
      performed.push(current.op)
      current = current.k(undefined)
    }
  }
  assert.equal(performed.length, N)
})

// --- the monad laws, observed through trace -------------------------------------------------------------------

const programs = {
  'a pure value': () => pure(4),
  'one effect': () => effect(seedOp),
  'a draw followed by reading the seed': () => flatMap(effect(drawOp(r1)), () => effect(seedOp)),
}
const f = (n) => flatMap(effect(drawOp({ ...r1, x: n })), () => pure(n * 2))
const g = (n) => map(effect(seedOp), (s) => s + n)

Object.entries(programs).forEach(([name, program]) => {
  test(`m07b left identity holds for ${name}`, () => {
    const a = trace(program(), 5).result
    assert.deepStrictEqual(trace(flatMap(pure(a), f), 5), trace(f(a), 5))
  })
  test(`m07b right identity holds for ${name}`, () =>
    assert.deepStrictEqual(trace(flatMap(program(), pure), 5), trace(program(), 5)))
  test(`m07b associativity holds for ${name}`, () =>
    assert.deepStrictEqual(
      trace(flatMap(flatMap(program(), f), g), 5),
      trace(flatMap(program(), (x) => flatMap(f(x), g)), 5)))
})

// --- the same programs under Module 7's IO --------------------------------------------------------------------

// Each program is written once against an interface {pure, effect, map, flatMap, sequence} and built with both.
const shared = {
  'seed, then a draw that depends on it': (M) =>
    M.flatMap(M.effect(seedOp), (s) => M.map(M.effect(drawOp({ ...r1, x: s % 20 })), () => s)),
  'a sequence of three effects': (M) => M.sequence([M.effect(drawOp(t1)), M.effect(seedOp), M.effect(drawOp(r1))]),
  '200 left-nested flatMaps': (M) =>
    Array.from({ length: 200 }).reduce((acc, _, i) => M.flatMap(acc, (n) => M.map(M.effect(drawOp({ ...r1, x: i })), () => n + 1)), M.pure(0)),
  'nested on both sides': (M) =>
    M.flatMap(M.flatMap(M.effect(seedOp), (s) => M.flatMap(M.effect(drawOp(r1)), () => M.pure(s + 1))),
      (n) => M.map(M.sequence([M.effect(drawOp(t1)), M.effect(seedOp)]), () => n * 3)),
}

Object.entries(shared).forEach(([name, program]) =>
  test(`m07b trace equals Module 7's trace for ${name}`, () =>
    assert.deepStrictEqual(trace(program(safe), 7), io.trace(program(io), 7))))