use crate::register_test;
use libc::{AF_INET, AF_UNIX, SOCK_DGRAM, SOCK_STREAM};
use libc::{accept, bind, connect, listen, shutdown, socket};

use std::io::{Read, Write};

pub fn test_tcp_socket_creation() {
    unsafe {
        let sockfd = socket(AF_INET, SOCK_STREAM, 0);
        if sockfd < 0 {
            panic!("Failed to create TCP socket");
        }
    }
}

register_test!(test_tcp_socket_creation);

pub fn test_unix_socket_creation() {
    unsafe {
        let sockfd = socket(AF_UNIX, SOCK_STREAM, 0);
        if sockfd < 0 {
            panic!("Failed to create UNIX stream socket");
        }
    }
    unsafe {
        let sockfd = socket(AF_UNIX, SOCK_DGRAM, 0);
        if sockfd < 0 {
            panic!("Failed to create UNIX datagram socket");
        }
    }
}

register_test!(test_unix_socket_creation);

pub fn test_unix_socket_basic_functions() {
    let sockfd = unsafe { socket(AF_UNIX, SOCK_STREAM, 0) };
    if sockfd < 0 {
        panic!("Failed to create UNIX stream socket for function tests");
    }
    let path = "/tmp/test_socket";
    let sockaddr = libc::sockaddr_un {
        sun_family: AF_UNIX as u16,
        sun_path: {
            let mut path_array = [0u8; 108];
            for (i, &b) in path.as_bytes().iter().enumerate() {
                path_array[i] = b;
            }
            path_array
        },
    };
    let bind_result = unsafe {
        bind(
            sockfd,
            &sockaddr as *const libc::sockaddr_un as *const libc::sockaddr,
            std::mem::size_of::<libc::sockaddr_un>() as u32,
        )
    };
    if bind_result < 0 {
        panic!("Failed to bind UNIX socket");
    }
    let listen_result = unsafe { listen(sockfd, 5) };
    if listen_result < 0 {
        panic!("Failed to listen on UNIX socket");
    }
    let shutdown_result = unsafe { shutdown(sockfd, 2) };
    if shutdown_result < 0 {
        panic!("Failed to shutdown UNIX socket");
    }
}

register_test!(test_unix_socket_basic_functions);

pub fn test_unix_socket_fork_msg_passing() {
    use std::ptr;

    // Create server socket, bind and listen before fork
    let server_fd = unsafe { socket(AF_UNIX, SOCK_STREAM, 0) };
    if server_fd < 0 {
        panic!("Failed to create server UNIX socket");
    }

    let path = "/tmp/uds_fork_test";
    let sockaddr = libc::sockaddr_un {
        sun_family: AF_UNIX as u16,
        sun_path: {
            let mut path_array = [0u8; 108];
            for (i, &b) in path.as_bytes().iter().enumerate() {
                path_array[i] = b;
            }
            path_array
        },
    };

    let ret = unsafe {
        bind(
            server_fd,
            &sockaddr as *const libc::sockaddr_un as *const libc::sockaddr,
            std::mem::size_of::<libc::sockaddr_un>() as u32,
        )
    };
    if ret < 0 {
        panic!("Server bind failed");
    }
    let ret = unsafe { listen(server_fd, 1) };
    if ret < 0 {
        panic!("Server listen failed");
    }

    let pid = unsafe { libc::fork() };
    if pid < 0 {
        panic!("fork failed");
    }

    if pid == 0 {
        // Child: client
        let client_fd = unsafe { socket(AF_UNIX, SOCK_STREAM, 0) };
        if client_fd < 0 {
            panic!("Client socket creation failed");
        }
        let ret = unsafe {
            connect(
                client_fd,
                &sockaddr as *const libc::sockaddr_un as *const libc::sockaddr,
                std::mem::size_of::<libc::sockaddr_un>() as u32,
            )
        };
        if ret < 0 {
            panic!("Client connect failed");
        }

        // Send request
        let req = b"hello";
        let wr = unsafe { libc::write(client_fd, req.as_ptr() as *const _, req.len()) };
        if wr != req.len() as isize {
            panic!("Client write failed");
        }

        // Receive response
        let mut resp = [0u8; 5];
        let rd = unsafe { libc::read(client_fd, resp.as_mut_ptr() as *mut _, resp.len()) };
        if rd != resp.len() as isize || &resp != b"world" {
            panic!("Client read failed");
        }

        unsafe { libc::close(client_fd) };
        unsafe { libc::_exit(0) };
    } else {
        // Parent: server
        let conn_fd = unsafe { accept(server_fd, ptr::null_mut(), ptr::null_mut()) };
        if conn_fd < 0 {
            panic!("Server accept failed");
        }

        // Receive request
        let mut buf = [0u8; 5];
        let rd = unsafe { libc::read(conn_fd, buf.as_mut_ptr() as *mut _, buf.len()) };
        if rd != buf.len() as isize || &buf != b"hello" {
            panic!("Server read failed");
        }

        // Send response
        let resp = b"world";
        let wr = unsafe { libc::write(conn_fd, resp.as_ptr() as *const _, resp.len()) };
        if wr != resp.len() as isize {
            panic!("Server write failed");
        }

        // Wait for child
        let mut status = 0;
        unsafe { libc::waitpid(pid, &mut status, 0) };
        if !libc::WIFEXITED(status) || libc::WEXITSTATUS(status) != 0 {
            panic!("Client process did not exit cleanly");
        }

        unsafe { libc::close(conn_fd) };
        unsafe { libc::close(server_fd) };
    }
}

register_test!(test_unix_socket_fork_msg_passing);

pub fn test_rust_unix_socket() {
    use std::os::unix::net::{UnixListener, UnixStream};
    use std::thread;

    let path = "/tmp/rust_uds_test";
    let listener = UnixListener::bind(path).expect("Failed to bind UNIX socket");

    thread::spawn(move || {
        let (mut stream, _) = listener.accept().expect("Failed to accept connection");
        let mut buf = [0u8; 5];
        stream
            .read_exact(&mut buf)
            .expect("Failed to read from stream");
        if &buf != b"hello" {
            panic!("Server read incorrect data");
        }
        //     stream
        //         .write_all(b"world")
        //         .expect("Failed to write to stream");
    });

    let mut stream = UnixStream::connect(path).expect("Failed to connect to UNIX socket");
    stream
        .write_all(b"hello")
        .expect("Failed to write to stream");
    // let mut buf = [0u8; 5];
    // stream
    //     .read_exact(&mut buf)
    //     .expect("Failed to read from stream");
    // if &buf != b"world" {
    //     panic!("Client read incorrect data");
    // }
}

register_test!(test_rust_unix_socket);
