#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/dma-mapping.h>
#include <linux/net.h>
#include <linux/in.h>
#include <linux/inet.h>
#include <linux/socket.h>
#include <linux/wait.h>

MODULE_LICENSE("GPLv3");
MODULE_AUTHOR("QuarkTree");
MODULE_DESCRIPTION("Enhanced DMA/Network Module with Parameter Support");

#define DMA_BUF_SIZE 4096

static char *target_ip = "10.0.0.20";
module_param(target_ip, charp, 0644);
MODULE_PARM_DESC(target_ip, "Target IP address (default: 10.0.0.20)");

static int target_port = 1234;
module_param(target_port, int, 0644);
MODULE_PARM_DESC(target_port, "Target port number (default: 1234)");

static bool use_dma = true;
module_param(use_dma, bool, 0644);
MODULE_PARM_DESC(use_dma, "Use DMA buffer (true/false, default: true)");

static void *data_buf;
static dma_addr_t dma_handle;
static struct socket *sock;

static int __init dma_net_module_init(void) {
    struct sockaddr_in sin;
    int ret;
    struct msghdr msg;
    struct kvec iov;
    __be32 ip_addr;

    if (target_port < 1 || target_port > 65535) {
        printk(KERN_ERR "Invalid port number: %d\n", target_port);
        return -EINVAL;
    }

    ip_addr = in_aton(target_ip);
    if (ip_addr == INADDR_NONE) {
        printk(KERN_ERR "Invalid IP address: %s\n", target_ip);
        return -EINVAL;
    }

    if (use_dma) {
        data_buf = dma_alloc_coherent(NULL, DMA_BUF_SIZE, &dma_handle, GFP_KERNEL);
        printk(KERN_INFO "Allocated DMA buffer at %p (phys: %pad)\n",
               data_buf, &dma_handle);
    } else {
        data_buf = kmalloc(DMA_BUF_SIZE, GFP_KERNEL);
        printk(KERN_INFO "Allocated kernel buffer at %p\n", data_buf);
    }

    if (!data_buf) {
        printk(KERN_ERR "Failed to allocate buffer\n");
        return -ENOMEM;
    }

    ret = sock_create_kern(&init_net, AF_INET, SOCK_STREAM, IPPROTO_TCP, &sock);
    if (ret < 0) {
        printk(KERN_ERR "Socket creation failed: %d\n", ret);
        goto free_buffer;
    }

    sock->flags |= O_NONBLOCK;

    memset(&sin, 0, sizeof(sin));
    sin.sin_family = AF_INET;
    sin.sin_port = htons(target_port);
    sin.sin_addr.s_addr = ip_addr;

    printk(KERN_INFO "Connecting to %s:%d...\n", target_ip, target_port);

    ret = sock->ops->connect(sock, (struct sockaddr *)&sin, sizeof(sin), 0);
    if (ret != -EINPROGRESS && ret < 0) {
        printk(KERN_ERR "Connect error: %d\n", ret);
        goto close_sock;
    }

    for (int i = 0; i < 10; i++) {
        if (sock->sk->sk_state == TCP_ESTABLISHED) break;
        msleep(100);
    }

    if (sock->sk->sk_state != TCP_ESTABLISHED) {
        printk(KERN_ERR "Connection timeout (state: %d)\n", 
               sock->sk->sk_state);
        ret = -ETIMEDOUT;
        goto close_sock;
    }

    printk(KERN_INFO "Connected successfully to %s:%d\n", 
           target_ip, target_port);

    memset(data_buf, 0xAA, DMA_BUF_SIZE);

    memset(&msg, 0, sizeof(msg));
    iov.iov_base = data_buf;
    iov.iov_len = DMA_BUF_SIZE;
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;

    ret = sock_sendmsg(sock, &msg, DMA_BUF_SIZE);
    if (ret < 0) {
        printk(KERN_ERR "Send failed: %d (sent %d bytes)\n", ret, ret > 0 ? ret : 0);
        goto close_sock;
    }

    printk(KERN_INFO "Sent %d bytes to %s:%d\n", 
           ret, target_ip, target_port);
    return 0;

close_sock:
    if (sock) {
        sock_release(sock);
        sock = NULL;
    }
free_buffer:
    if (data_buf) {
        if (use_dma) {
            dma_free_coherent(NULL, DMA_BUF_SIZE, data_buf, dma_handle);
            printk(KERN_INFO "Freed DMA buffer\n");
        } else {
            kfree(data_buf);
            printk(KERN_INFO "Freed kernel buffer\n");
        }
        data_buf = NULL;
    }
    return ret;
}

static void __exit dma_net_module_exit(void) {
    if (sock) {
        sock_release(sock);
        printk(KERN_INFO "Socket released\n");
    }
    if (data_buf) {
        if (use_dma) {
            dma_free_coherent(NULL, DMA_BUF_SIZE, data_buf, dma_handle);
            printk(KERN_INFO "Freed DMA buffer\n");
        } else {
            kfree(data_buf);
            printk(KERN_INFO "Freed kernel buffer\n");
        }
        data_buf = NULL;
    }
    printk(KERN_INFO "Module unloaded\n");
}

module_init(dma_net_module_init);
module_exit(dma_net_module_exit);