#![feature(yggdrasil_os)]

use std::os::fd::AsRawFd;
use std::os::yggdrasil::io::poll::PollChannel;
use std::{
    io::{self, stdin, stdout, Read, Write},
    net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, TcpStream, UdpSocket},
    process::ExitCode,
};

use clap::Parser;

#[derive(Debug, Parser)]
struct Args {
    #[clap(short, long)]
    listen: bool,
    #[clap(short, long)]
    udp: bool,
    #[clap(short = 'p', long)]
    listen_port: Option<u16>,
    address: Option<SocketAddr>,
}

fn udp_client(address: SocketAddr) -> io::Result<()> {
    let mut stdin = stdin();
    let mut stdout = stdout();
    let mut poll = PollChannel::new()?;
    // TODO use ephemeral port?
    // TODO use socket.connect()
    let socket = UdpSocket::bind(SocketAddr::new(
        IpAddr::V4(Ipv4Addr::UNSPECIFIED),
        address.port(),
    ))?;
    let mut buffer = [0; 512];

    let socket_fd = socket.as_raw_fd();
    let stdin_fd = stdin.as_raw_fd();

    poll.add(socket_fd)?;
    poll.add(stdin_fd)?;

    loop {
        let (fd, status) = poll.wait(None, true)?.unwrap();
        status?;

        match fd {
            fd if fd == socket_fd => {
                let (len, _) = socket.recv_from(&mut buffer)?;
                stdout.write_all(&buffer[..len])?;
            }
            fd if fd == stdin_fd => {
                let len = stdin.read(&mut buffer)?;
                if len == 0 {
                    break;
                }
                socket.send_to(&buffer[..len], address)?;
            }
            _ => unreachable!(),
        }
    }

    Ok(())
}

fn udp_server(port: u16) -> io::Result<()> {
    let mut stdout = stdout();
    let socket = UdpSocket::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port))?;
    let mut buffer = [0; 2048];

    loop {
        let (len, _) = socket.recv_from(&mut buffer)?;
        stdout.write_all(&buffer[..len])?;
    }
}

fn tcp_client(address: SocketAddr) -> io::Result<()> {
    let mut stdin = stdin();
    let mut stdout = stdout();
    let mut poll = PollChannel::new()?;
    let mut socket = TcpStream::connect(address)?;
    let mut buffer = [0; 512];

    let socket_fd = socket.as_raw_fd();
    let stdin_fd = stdin.as_raw_fd();

    poll.add(socket_fd)?;
    poll.add(stdin_fd)?;

    loop {
        let (fd, status) = poll.wait(None, true)?.unwrap();
        status?;

        match fd {
            fd if fd == socket_fd => {
                let len = socket.read(&mut buffer)?;
                stdout.write_all(&buffer[..len])?;
            }
            fd if fd == stdin_fd => {
                let len = stdin.read(&mut buffer)?;
                if len == 0 {
                    break;
                }
                socket.write_all(&buffer[..len])?;
            }
            _ => unreachable!(),
        }
    }

    Ok(())
}

fn tcp_server(port: u16) -> io::Result<()> {
    let mut stdout = stdout();
    let listener = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::UNSPECIFIED), port))?;
    let mut buffer = [0; 2048];

    let (mut stream, _) = listener.accept()?;

    loop {
        let len = stream.read(&mut buffer)?;
        stdout.write_all(&buffer[..len])?;
    }
}

fn main() -> ExitCode {
    let args = Args::parse();

    if args.listen && args.address.is_some() {
        eprintln!("--listen/-l and remote address cannot be used simultaneously");
        return ExitCode::FAILURE;
    }

    let result = if args.listen {
        let Some(port) = args.listen_port else {
            eprintln!("--listen/-l requires --listen-port/-p PORT");
            return ExitCode::FAILURE;
        };

        if args.udp {
            udp_server(port)
        } else {
            tcp_server(port)
        }
    } else {
        let Some(address) = args.address else {
            eprintln!("Remote address required");
            return ExitCode::FAILURE;
        };

        if args.udp {
            udp_client(address)
        } else {
            tcp_client(address)
        }
    };

    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("Error: {}", error);
            ExitCode::FAILURE
        }
    }
}
