// vtest build: !windows
import encoding.base64
import os
import net
import net.websocket
import time

const proxy_github_job = os.getenv('GITHUB_JOB')

const proxy_should_skip = proxy_github_job != '' && proxy_github_job != 'websocket_tests'

const proxy_ws_target_port = 30105
const proxy_ws_port = 30106

@[heap]
struct ProxyWebsocketTestResults {
pub mut:
    nr_messages      int
    nr_pong_received int
    error_message    string
}

fn start_proxy_ws_server(listen_port int) ! {
    mut s := websocket.new_server(.ip, listen_port, '')
    s.set_ping_interval(1)
    s.on_connect(fn (mut s websocket.ServerClient) !bool {
        return s.resource_name == '/'
    })!
    s.on_message(fn (mut ws websocket.Client, msg &websocket.Message) ! {
        match msg.opcode {
            .pong { ws.write_string('pong')! }
            else { ws.write(msg.payload, msg.opcode)! }
        }
    })
    start_proxy_ws_server_in_thread(mut s)
}

fn start_proxy_ws_server_in_thread(mut ws websocket.Server) {
    spawn fn [mut ws] () {
        ws.listen() or { panic('websocket server could not listen, err: ${err}') }
    }()
    for ws.get_state() != .open {
        time.sleep(10 * time.millisecond)
    }
}

fn start_proxy_server(listen_port int, target string, requests chan string) {
    ready := chan bool{cap: 1}
    spawn fn [listen_port, target, requests, ready] () {
        mut listener := net.listen_tcp(.ip, ':${listen_port}') or { panic(err) }
        ready <- true
        mut client := listener.accept() or {
            listener.close() or {}
            panic(err)
        }
        listener.close() or {}
        request := read_proxy_request(mut client) or {
            requests <- 'ERROR: ${err}'
            client.close() or {}
            return
        }
        requests <- request
        mut upstream := net.dial_tcp(target) or {
            client.write_string('HTTP/1.1 502 Bad Gateway\r\n\r\n') or {}
            client.close() or {}
            return
        }
        client.write_string('HTTP/1.1 200 Connection Established\r\n\r\n') or {
            client.close() or {}
            upstream.close() or {}
            return
        }
        tunnel_proxy_connection(mut client, mut upstream)
    }()
    _ := <-ready
}

fn read_proxy_request(mut conn net.TcpConn) !string {
    mut total_bytes_read := 0
    mut msg := [4096]u8{}
    mut buffer := [1]u8{}
    for total_bytes_read < msg.len {
        bytes_read := conn.read_ptr(&buffer[0], 1)!
        if bytes_read == 0 {
            return error('unexpected EOF while reading proxy request')
        }
        msg[total_bytes_read] = buffer[0]
        total_bytes_read++
        if total_bytes_read > 3 && msg[total_bytes_read - 1] == `\n`
            && msg[total_bytes_read - 2] == `\r` && msg[total_bytes_read - 3] == `\n`
            && msg[total_bytes_read - 4] == `\r` {
            return msg[..total_bytes_read].bytestr()
        }
    }
    return error('proxy request headers exceeded 4096 bytes')
}

fn tunnel_proxy_connection(mut client net.TcpConn, mut upstream net.TcpConn) {
    defer {
        client.close() or {}
        upstream.close() or {}
    }
    client.set_read_timeout(50 * time.millisecond)
    upstream.set_read_timeout(50 * time.millisecond)
    deadline := time.now().add(4 * time.second)
    for time.now() < deadline {
        relay_proxy_connection(mut client, mut upstream)
        relay_proxy_connection(mut upstream, mut client)
    }
}

fn relay_proxy_connection(mut src net.TcpConn, mut dst net.TcpConn) {
    mut buf := []u8{len: 1024}
    bytes_read := src.read(mut buf) or {
        if err.code() == net.err_timed_out_code {
            return
        }
        return
    }
    if bytes_read > 0 {
        dst.write(buf[..bytes_read]) or {}
    }
}

fn proxy_open_cb(mut client websocket.Client) ! {
    client.pong()!
}

fn proxy_error_cb(mut client websocket.Client, err string, mut res ProxyWebsocketTestResults) ! {
    res.error_message = err
}

fn proxy_message_cb(mut client websocket.Client, msg &websocket.Message, mut res ProxyWebsocketTestResults) ! {
    if msg.opcode != .text_frame {
        return
    }
    match msg.payload.bytestr() {
        'pong' {
            res.nr_pong_received++
        }
        'a' {
            res.nr_messages++
        }
        else {
            res.error_message = 'unexpected payload: ${msg.payload.bytestr()}'
        }
    }
}

fn test_ws_connection_through_proxy() ! {
    if proxy_should_skip {
        return
    }
    start_proxy_ws_server(proxy_ws_target_port)!
    requests := chan string{cap: 1}
    start_proxy_server(proxy_ws_port, '127.0.0.1:${proxy_ws_target_port}', requests)
    mut results := ProxyWebsocketTestResults{}
    mut client := websocket.new_client('ws://127.0.0.1:${proxy_ws_target_port}', websocket.ClientOpt{
        proxy_url: 'http://user:[email protected]:${proxy_ws_port}'
    })!
    client.on_open(proxy_open_cb)
    client.on_error_ref(proxy_error_cb, results)
    client.on_message_ref(proxy_message_cb, results)
    client.connect()!
    spawn client.listen()
    for msg in ['a', 'a'] {
        client.write(msg.bytes(), .text_frame)!
        time.sleep(100 * time.millisecond)
    }
    time.sleep(1500 * time.millisecond)
    request := <-requests
    auth_token := base64.encode('user:pass'.bytes())
    assert request.starts_with('CONNECT 127.0.0.1:${proxy_ws_target_port} HTTP/1.1\r\n')
    assert request.contains('Proxy-Connection: Keep-Alive\r\n')
    assert request.contains('Proxy-Authorization: Basic ${auth_token}\r\n')
    client.close(1000, 'done') or {}
    time.sleep(100 * time.millisecond)
    assert results.error_message == ''
    assert results.nr_pong_received >= 2
    assert results.nr_messages == 2
}