#pragma once
#include "uefi.h"

#define NET_MAC_LEN     6
#define NET_IP_LEN      4
#define NET_ETH_HDR_LEN 14
#define NET_IP_HDR_LEN  20
#define NET_ARP_HDR_LEN 28
#define NET_ICMP_HDR_LEN 8
#define NET_TX_BUFS     32
#define NET_RX_BUFS     32
#define NET_BUF_SIZE    2048

/* Ethernet */
#define ETH_P_IP   0x0800
#define ETH_P_ARP  0x0806

/* ARP */
#define ARP_OP_REQUEST 1
#define ARP_OP_REPLY   2

/* ICMP */
#define ICMP_ECHO_REQUEST  8
#define ICMP_ECHO_REPLY    0

/* IP */
#define IP_PROTO_ICMP 1
#define IP_PROTO_UDP  17
#define IP_PROTO_TCP  6

struct net_eth_hdr {
    uint8_t  dst[6];
    uint8_t  src[6];
    uint16_t type;
} __attribute__((packed));

struct net_arp_hdr {
    uint16_t hw_type;
    uint16_t proto_type;
    uint8_t  hw_len;
    uint8_t  proto_len;
    uint16_t opcode;
    uint8_t  sender_mac[6];
    uint8_t  sender_ip[4];
    uint8_t  target_mac[6];
    uint8_t  target_ip[4];
} __attribute__((packed));

struct net_ip_hdr {
    uint8_t  ver_ihl;
    uint8_t  tos;
    uint16_t total_len;
    uint16_t id;
    uint16_t flags_frag;
    uint8_t  ttl;
    uint8_t  proto;
    uint16_t checksum;
    uint8_t  src_ip[4];
    uint8_t  dst_ip[4];
} __attribute__((packed));

struct net_icmp_hdr {
    uint8_t  type;
    uint8_t  code;
    uint16_t checksum;
    uint16_t id;
    uint16_t seq;
} __attribute__((packed));

struct net_udp_hdr {
    uint16_t src_port;
    uint16_t dst_port;
    uint16_t length;
    uint16_t checksum;
} __attribute__((packed));

struct net_tcp_hdr {
    uint16_t src_port;
    uint16_t dst_port;
    uint32_t seq;
    uint32_t ack;
    uint8_t  doff_flags;
    uint8_t  flags;
    uint16_t window;
    uint16_t checksum;
    uint16_t urgent;
} __attribute__((packed));

#define TCP_FIN  0x01
#define TCP_SYN  0x02
#define TCP_RST  0x04
#define TCP_PSH  0x08
#define TCP_ACK  0x10

/* TCP states */
enum tcp_state {
    TCP_CLOSED = 0,
    TCP_SYN_SENT,
    TCP_ESTABLISHED,
    TCP_FIN_WAIT_1,
    TCP_FIN_WAIT_2,
    TCP_CLOSE_WAIT,
    TCP_LAST_ACK
};

#define TCP_BUF_SIZE 65536

struct tcp_conn {
    enum tcp_state state;
    uint8_t  dst_ip[4];
    uint16_t dst_port;
    uint16_t src_port;
    uint32_t snd_nxt;
    uint32_t rcv_nxt;
    uint8_t  rx_buf[TCP_BUF_SIZE];
    uint32_t rx_head;
    uint32_t rx_tail;
};

/* Init */
int      net_init(void);

/* Low-level send */
int      net_send_raw(uint16_t eth_type, const uint8_t *dst_mac,
                       const void *data, uint32_t len);
int      net_send_ip(const uint8_t *dst_ip, uint8_t proto,
                      const void *data, uint32_t len);

/* UDP */
int      net_send_udp(const uint8_t *dst_ip, uint16_t dst_port,
                       uint16_t src_port, const void *data, uint32_t len);

/* ARP */
int      net_arp_resolve(const uint8_t *ip, uint8_t *mac_out);

/* ICMP */
int      net_ping(const uint8_t *ip, int count);

/* DNS */
int      dns_resolve(const char *hostname, uint8_t *ip_out);

/* TCP */
int      tcp_connect(const uint8_t *ip, uint16_t port);
int      tcp_send(const void *data, uint32_t len);
int      tcp_recv(void *buf, uint32_t max_len, uint32_t timeout_ms);
void     tcp_close(void);

/* Poll */
void     net_poll(void);

/* Utilities */
void     net_get_mac(uint8_t *mac);
void     net_get_ip(uint8_t *ip);
void     net_get_mask(uint8_t *mask);
void     net_get_gw(uint8_t *gw);
void     net_set_ip(const uint8_t *ip);
void     net_set_mask(const uint8_t *mask);
void     net_set_gw(const uint8_t *gw);
uint16_t net_checksum(const void *data, uint32_t len);
