package org.kne.cloud.network.tcp;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.util.Arrays;

import org.kne.cloud.network.NetworkPacket;
import org.kne.cloud.network.ipv6.IPv6Packet.IPv6Payload;
import org.kne.cloud.network.klalb.PortPacket;
import org.kne.io.KNEChannels;

public class UDPPacket extends IPv6Payload implements PortPacket {
    public static final int UDP_PROTOCOL_NUMBER = 17;
    public static final int UDP_HEADER_LENGTH = 8; // UDP头部固定8字节
    
    private ByteBuffer header = NetworkPacket.bufferAllocator.allocate(UDP_HEADER_LENGTH);
    private ByteBuffer data;

    public UDPPacket() {
        super(UDP_PROTOCOL_NUMBER);
    }

    public UDPPacket(int srcPort,int dstPort) {
        super(UDP_PROTOCOL_NUMBER);
        header.putChar((char) srcPort);
        header.putChar((char) dstPort);
        data = NetworkPacket.bufferAllocator.allocate(65535);
        
    }
    
    @Override
	public long getTotalLength() {
		return data.limit()+UDP_HEADER_LENGTH;
	}

	
    @Override
    public void writeToChannel(WritableByteChannel dto) throws IOException {
        setUDPLength(getDataLength() + UDP_HEADER_LENGTH); // 设置总长度（头+数据）
        dto.write(header.slice(0, UDP_HEADER_LENGTH));
        dto.write(data.slice(0, data.limit()));
    }

    @Override
    public void readFromChannel(ReadableByteChannel din, long length) throws IOException {
        header.clear().limit(UDP_HEADER_LENGTH);
        KNEChannels.readFully(din, header);
        
        // 读取数据部分（总长度减去头部长度）
        int totalLength = getUDPLength();
        int dataLength = totalLength - UDP_HEADER_LENGTH;
        
        if (dataLength < 0) {
            throw new IOException("Invalid UDP packet: total length < header length");
        }
        
        data = NetworkPacket.bufferAllocator.allocate(dataLength);
        KNEChannels.readFully(din, data);
        data.flip();
    }

    @Override
    protected boolean needEndPosition() {
        return false;
    }

    public int getSrcPort() {
        return header.getChar(0);
    }

    public void setSrcPort(int srcPort) {
        header.putChar(0, (char) srcPort);
    }

    public int getDstPort() {
        return header.getChar(2);
    }

    public void setDstPort(int dstPort) {
        header.putChar(2, (char) dstPort);
    }

    /**
     * 获取UDP数据包总长度（头部+数据）
     */
    public int getUDPLength() {
        return header.getChar(4);
    }

    /**
     * 设置UDP数据包总长度（头部+数据）
     */
    public void setUDPLength(int totalLength) {
        header.putChar(4, (char) totalLength);
    }

    /**
     * 获取数据部分长度（不包含头部）
     */
    public int getDataLength() {
        return data != null ? data.limit() : 0;
    }

    public int getChecksum() {
        return header.getChar(6);
    }

    public void setChecksum(int checksum) {
        header.putChar(6, (char) checksum);
    }

    public ByteBuffer getData() {
        return data;
    }

    public int calculateChecksum() {
        if (getParent() == null) {
            throw new IllegalStateException("Parent IPv6 packet required for checksum calculation");
        }

        // 保存原始校验和值
        int originalChecksum = getChecksum();
        
        try {
            // 临时将校验和字段设为0
            setChecksum(0); 


            // IPv6伪首部

            long sum = getParent().calculateAddressChecksum();

            // 使用总长度（头部+数据）
            int totalLength = getUDPLength();
            sum += (totalLength >>> 16) + totalLength;
            sum += UDP_PROTOCOL_NUMBER;

            // UDP首部（此时校验和字段为0）
            ByteBuffer headerCopy = header.duplicate();
            headerCopy.position(0).limit(UDP_HEADER_LENGTH);
            while (headerCopy.remaining() >= 2) {
                sum += headerCopy.getChar();
            }

            // UDP数据
            if (data != null && data.limit() > 0) {
                ByteBuffer dataCopy = data.duplicate();
                dataCopy.position(0);
                
                while (dataCopy.remaining() >= 2) {
                    sum += dataCopy.getChar();
                }
                
                if (dataCopy.remaining() == 1) {
                    sum += (dataCopy.get() & 0xFF) << 8;
                }
            }

            // 折叠进位并取反
            while ((sum >> 16) != 0) {
                sum = (sum & 0xFFFF) + (sum >> 16);
            }

            int checksum = (int) (~sum & 0xFFFF);
            return checksum == 0 ? 0xFFFF : checksum;
            
        } finally {
            // 恢复原始校验和值
            setChecksum(originalChecksum);
        }
    }

    public void updateChecksum() {
        setChecksum(calculateChecksum());
    }

    public boolean verifyChecksum() {
        int storedChecksum = getChecksum();
        return storedChecksum == 0 || storedChecksum == calculateChecksum();
    }

    @Override
    public String toString() {
        byte[] b = new byte[Math.min(data.limit(), 10)];
        data.get(0, b);
        return "UDP " + getSrcPort() + "->" + getDstPort() + " [Total:" + getUDPLength() 
               + ", Data:" + getDataLength() + "] " + Arrays.toString(b);
    }
  
	
}