//
//  IPPacket.swift
//  Tun
//
//  Created by 安礼成 on 2024/1/18.
//

import Foundation

enum IPVersion: UInt8 {
    case ipv4 = 4
    case ipv6 = 6
}

enum TransportProtocol: UInt8 {
    case icmp = 1
    case tcp  = 6
    case udp  = 17
}

// MARK: - IP Header

struct IPHeader {
    let version: UInt8
    let headerLength: UInt8
    let typeOfService: UInt8
    let totalLength: UInt16
    let id: UInt16
    let offset: UInt16
    let ttl: UInt8
    let proto: UInt8
    let checksum: UInt16
    let source: UInt32
    let destination: UInt32

    var headerBytes: Int {
        Int(headerLength)
    }
}

// MARK: - Lightweight IP Packet View

struct IPPacketView {
    let header: IPHeader
    let data: Data
    let transportPacket: TransportPacket

    enum TransportPacket {
        case tcp(srcPort: UInt16, dstPort: UInt16)
        case udp(srcPort: UInt16, dstPort: UInt16, payloadOffset: Int)
        case icmp
        case unsupported(UInt8)
        case malformed
    }

    init?(_ data: Data) {
        guard data.count >= 20 else {
            return nil
        }

        let firstByte = Self.byte(at: 0, in: data)
        let version = firstByte >> 4
        let headerLen = (firstByte & 0x0F) * 4

        guard headerLen >= 20, data.count >= headerLen else {
            return nil
        }

        let proto = Self.byte(at: 9, in: data)
        self.header = IPHeader(
            version: version,
            headerLength: headerLen,
            typeOfService: Self.byte(at: 1, in: data),
            totalLength: UInt16(bytes: (Self.byte(at: 2, in: data), Self.byte(at: 3, in: data))),
            id: UInt16(bytes: (Self.byte(at: 4, in: data), Self.byte(at: 5, in: data))),
            offset: UInt16(bytes: (Self.byte(at: 6, in: data), Self.byte(at: 7, in: data))),
            ttl: Self.byte(at: 8, in: data),
            proto: proto,
            checksum: UInt16(bytes: (Self.byte(at: 10, in: data), Self.byte(at: 11, in: data))),
            source: UInt32(bytes: (Self.byte(at: 12, in: data), Self.byte(at: 13, in: data), Self.byte(at: 14, in: data), Self.byte(at: 15, in: data))),
            destination: UInt32(bytes: (Self.byte(at: 16, in: data), Self.byte(at: 17, in: data), Self.byte(at: 18, in: data), Self.byte(at: 19, in: data)))
        )

        self.data = data
        self.transportPacket = Self.parseTransportPacket(proto: proto, offset: Int(headerLen), data: data)
    }

    private static func parseTransportPacket(proto: UInt8, offset: Int, data: Data) -> TransportPacket {
        guard let proto = TransportProtocol(rawValue: proto) else {
            return .unsupported(proto)
        }

        switch proto {
        case .tcp:
            guard data.count >= offset + 20 else {
                return .malformed
            }
            return .tcp(
                srcPort: UInt16(bytes: (Self.byte(at: offset, in: data), Self.byte(at: offset + 1, in: data))),
                dstPort: UInt16(bytes: (Self.byte(at: offset + 2, in: data), Self.byte(at: offset + 3, in: data)))
            )

        case .udp:
            guard data.count >= offset + 8 else {
                return .malformed
            }
            return .udp(
                srcPort: UInt16(bytes: (Self.byte(at: offset, in: data), Self.byte(at: offset + 1, in: data))),
                dstPort: UInt16(bytes: (Self.byte(at: offset + 2, in: data), Self.byte(at: offset + 3, in: data))),
                payloadOffset: offset + 8
            )

        case .icmp:
            guard data.count >= offset + 4 else {
                return .malformed
            }
            return .icmp
        }
    }

    private static func byte(at offset: Int, in data: Data) -> UInt8 {
        data[data.index(data.startIndex, offsetBy: offset)]
    }
}