import flag
import net.http
import os
import time

const probe_url = 'https://httpbin.org/post'
const vschannel_16kb_boundary = 16 * 1024
const default_sizes_csv = '12000,15000,17000,24000,32768'

fn main() {
    if os.args.len >= 2 && os.args[1] == '--worker' {
        run_worker_mode()
        return
    }
    run_parent_mode()
}

fn run_worker_mode() {
    if os.args.len < 3 {
        eprintln('missing payload size for --worker')
        exit(2)
    }
    size := os.args[2].int()
    if size <= 0 {
        eprintln('invalid payload size `${os.args[2]}`')
        exit(2)
    }
    payload := '{"size":${size},"payload":"' + 'x'.repeat(size) + '"}'
    mut headers := http.new_header()
    headers.add(.content_type, 'application/json')
    mut req := http.Request{
        method:        .post
        url:           probe_url
        header:        headers
        data:          payload
        read_timeout:  60 * time.second
        write_timeout: 30 * time.second
    }
    resp := req.do() or {
        eprintln('request failed at size=${size}: ${err}')
        exit(1)
    }
    if resp.status_code != 200 {
        eprintln('unexpected status at size=${size}: ${resp.status_code}')
        exit(1)
    }
    if !resp.body.contains('httpbin.org/post') {
        eprintln('unexpected response body at size=${size}')
        exit(1)
    }
    println('OK size=${size} payload_bytes=${payload.len}')
}

fn run_parent_mode() {
    mut fp := flag.new_flag_parser(os.args)
    fp.application('vschannel_16kb_httpbin_probe')
    fp.version('0.0.1')
    fp.description('Probe net.http HTTPS POST behavior around 16KB payloads using https://httpbin.org/post.')
    fp.skip_executable()

    expect_mode := fp.string('expect', `e`, 'after',
        'Expected behavior mode: before|after|none. before expects >16KB failures, after expects all succeed.')
    sizes_csv := fp.string('sizes', `s`, default_sizes_csv,
        'Comma separated payload sizes in bytes.')
    verbose := fp.bool('verbose', `v`, false, 'Print child process output for all sizes.')
    show_help := fp.bool('help', `h`, false, 'Show this help screen.')
    free_args := fp.finalize() or {
        eprintln('flag parse failed: ${err}')
        exit(2)
    }
    if free_args.len > 0 {
        eprintln('unexpected positional args: ${free_args}')
        exit(2)
    }
    if show_help {
        println(fp.usage())
        return
    }

    mode := expect_mode.to_lower()
    if mode !in ['before', 'after', 'none'] {
        eprintln('invalid --expect `${expect_mode}` (allowed: before|after|none)')
        exit(2)
    }
    sizes := parse_sizes_csv(sizes_csv) or {
        eprintln(err)
        exit(2)
    }

    self_exe := os.executable()
    println('probe url: ${probe_url}')
    println('expect mode: ${mode}')
    println('sizes: ${sizes}')

    mut mismatches := 0
    for size in sizes {
        // Launch worker directly without going through a shell, and capture stdout/stderr separately.
        mut p := os.new_process(self_exe)
        p.set_args(['--worker', size.str()])
        p.set_redirect_stdio()
        p.run()
        p.wait()
        exit_code := p.code
        stdout := p.stdout_read().trim_space()
        stderr := p.stderr_read().trim_space()
        p.close()
        success := exit_code == 0
        expected_success := expected_success_for_mode(mode, size)
        status := if success { 'OK' } else { 'FAIL' }
        expect_text := if expected_success { 'expect=OK' } else { 'expect=FAIL' }
        println('${status:4} size=${size:6} exit=${exit_code:3} ${expect_text}')
        if verbose || !success {
            if stdout.len > 0 {
                println(stdout)
            }
            if stderr.len > 0 {
                eprintln(stderr)
            }
        }
        if mode != 'none' && success != expected_success {
            mismatches++
        }
    }

    if mismatches > 0 {
        eprintln('mismatch count: ${mismatches}')
        exit(1)
    }
    println('probe completed without mismatches')
}

fn parse_sizes_csv(raw string) ![]int {
    mut sizes := []int{}
    for part in raw.split(',') {
        item := part.trim_space()
        if item.len == 0 {
            continue
        }
        size := item.int()
        if size <= 0 {
            return error('invalid size entry `${item}` in --sizes')
        }
        sizes << size
    }
    if sizes.len == 0 {
        return error('no sizes provided in --sizes')
    }
    return sizes
}

fn expected_success_for_mode(mode string, size int) bool {
    return match mode {
        'before' {
            size <= vschannel_16kb_boundary
        }
        'after' {
            true
        }
        'none' {
            true
        }
        else {
            eprintln('unexpected mode `${mode}` in expected_success_for_mode')
            false
        }
    }
}