import io
import rand

fn small_amount() int {
    return int(rand.u8()) + 1
}

struct ArrayWriter {
pub mut:
    result []u8
}

// implement io.Writer
fn (mut aw ArrayWriter) write(buf []u8) !int {
    len := buf.len
    mut res := 0
    for i := 0; i < len; i++ {
        aw.result << buf[i]
        res++
    }
    return res
}

struct ChunkedArrayWriter {
    chunk int
pub mut:
    result []u8
}

fn (mut cw ChunkedArrayWriter) write(buf []u8) !int {
    if buf.len == 0 {
        return 0
    }
    n := if buf.len < cw.chunk { buf.len } else { cw.chunk }
    for i := 0; i < n; i++ {
        cw.result << buf[i]
    }
    return n
}

struct FailOnceChunkedWriter {
    chunk int
pub mut:
    result []u8
mut:
    writes int
}

fn (mut fw FailOnceChunkedWriter) write(buf []u8) !int {
    if fw.writes == 1 {
        fw.writes++
        return error('forced write failure')
    }
    fw.writes++
    n := if buf.len < fw.chunk { buf.len } else { fw.chunk }
    for i := 0; i < n; i++ {
        fw.result << buf[i]
    }
    return n
}

struct ZeroProgressWriter {}

fn (mut w ZeroProgressWriter) write(buf []u8) !int {
    return 0
}

// write less than max bytes, returns number of bytes written
// data is written in chunks.
fn write_random_data(mut aw ArrayWriter, mut bw io.BufferedWriter, max int) !int {
    less_than_max := max - 255 // guarantee 1 full u8 less than max
    mut total := 0
    for total < less_than_max {
        r := rand.u8()
        d := rand.bytes(r)!
        w := bw.write(d)!
        total += w
    }
    return total
}

fn test_flush() {
    mut aw := ArrayWriter{}
    max := 65536
    mut bw := io.new_buffered_writer(writer: &aw, cap: max)!

    // write less data than buffer capacity, the underlying writer should receive no data.
    written := write_random_data(mut aw, mut bw, 65536)!
    assert written == bw.buffered()
    assert aw.result.len == 0

    bw.flush()!
    assert aw.result.len == written
    assert bw.buffered() == 0
    assert bw.available() == max
}

fn test_write() {
    mut aw := ArrayWriter{}
    max := 65536
    mut bw := io.new_buffered_writer(writer: &aw, cap: max)!

    // write less data than buffer capacity, the underlying writer should receive no data.
    written := write_random_data(mut aw, mut bw, 65536)!

    // now exceed buffer capacity by a little
    little := small_amount()
    excess := bw.available() + little
    excess_data := rand.bytes(excess)!
    w := bw.write(excess_data)!
    assert bw.buffered() == little
    assert bw.available() == max - little
    assert aw.result.len == max
}

fn test_write_big() {
    mut aw := ArrayWriter{}
    max := 65536
    mut bw := io.new_buffered_writer(writer: &aw, cap: max)!

    more_than_max := max + small_amount()
    big_source := rand.bytes(more_than_max)!
    w := bw.write(big_source)!
    assert w == more_than_max
    assert bw.buffered() == 0
    assert bw.available() == max
    assert aw.result.len == more_than_max
}

// create_data returns an array with `n` elements plus `\n` as last character.
fn create_data(n int) []u8 {
    mut res := []u8{}
    for i := 0; i < n; i++ {
        res << `X`
    }
    res << `\n`
    return res
}

fn test_simple_write() {
    mut aw := ArrayWriter{}
    mut bw := io.new_buffered_writer(writer: &aw, cap: 10)!
    mut data := create_data(6)
    w1 := bw.write(data)!

    // add data to buffer, less than cap
    // data is written to buffer, not to underlying writer.
    assert w1 == 7 // 6*x + \n
    assert bw.buffered() == 7
    assert aw.result.len == 0

    // add more data, exceed cap
    // data is flushed to underlying writer when buffer is full, then more data is written to buffer.
    w2 := bw.write(data)!
    assert w2 == 7
    assert bw.buffered() == 4
    assert aw.result.len == 10

    // exceed cap immediately
    // all data is written without buffering
    aw.result = []u8{}
    data = create_data(33)
    bw.reset()
    w3 := bw.write(data)!
    assert bw.buffered() == 0
    assert aw.result.len == 34 // 33*x + \n
}

fn test_simple_flush() {
    mut aw := ArrayWriter{}
    mut bw := io.new_buffered_writer(writer: &aw, cap: 10)!
    data := create_data(6)
    w := bw.write(data)!

    assert w == 7 // 6*x + \n
    assert bw.buffered() == 7

    bw.flush()!
    assert bw.buffered() == 0
    assert aw.result.len == 7
}

fn test_flush_handles_partial_writes() {
    mut cw := ChunkedArrayWriter{
        chunk: 2
    }
    mut bw := io.new_buffered_writer(writer: &cw, cap: 8)!
    data := 'abcdefg'.bytes()

    written := bw.write(data)!
    assert written == data.len
    assert bw.buffered() == data.len
    assert cw.result.len == 0

    bw.flush()!
    assert bw.buffered() == 0
    assert cw.result == data
}

fn test_flush_preserves_unwritten_data_on_error() {
    mut fw := FailOnceChunkedWriter{
        chunk: 3
    }
    mut bw := io.new_buffered_writer(writer: &fw, cap: 8)!
    data := 'abcdefg'.bytes()
    _ = bw.write(data)!

    if _ := bw.flush() {
        assert false
    } else {
        assert err.msg() == 'forced write failure'
    }
    assert fw.result == 'abc'.bytes()
    assert bw.buffered() == 4
    assert bw.buf[..bw.buffered()] == 'defg'.bytes()

    bw.flush()!
    assert bw.buffered() == 0
    assert fw.result == data
}

fn test_write_returns_error_for_zero_progress_writes() {
    mut zw := ZeroProgressWriter{}
    mut bw := io.new_buffered_writer(writer: &zw, cap: 4)!

    if _ := bw.write('12345'.bytes()) {
        assert false
    } else {
        assert err.msg() == 'writer returned an invalid number of bytes while writing'
    }
    assert bw.buffered() == 0
}