package org.kne.cloud.network.tcp;

import java.io.EOFException;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.WritableByteChannel;
import java.util.ArrayList;
import java.util.List;

import org.kne.cloud.network.NetworkPacket;
import org.kne.cloud.network.ipv6.IPv6Packet.IPv6DefaultPayload;

public class TCPPacket extends IPv6DefaultPayload {
    public static final int TCP_PROTOCOL_NUMBER = 6;
    
    // TCP头部常量
    private static final int TCP_HEADER_MIN_LENGTH = 20;
    private static final int TCP_OPTIONS_MAX_LENGTH = 40;
    
    // TCP标志位掩码
    private static final int FLAG_NS = 0x100;
    private static final int FLAG_CWR = 0x80;
    private static final int FLAG_ECE = 0x40;
    private static final int FLAG_URG = 0x20;
    private static final int FLAG_ACK = 0x10;
    private static final int FLAG_PSH = 0x08;
    private static final int FLAG_RST = 0x04;
    private static final int FLAG_SYN = 0x02;
    private static final int FLAG_FIN = 0x01;
    
    private ByteBuffer header = NetworkPacket.bufferAllocator.allocate(TCP_HEADER_MIN_LENGTH);
    private List<TCPOptionTLV> options = new ArrayList<>();
    
    public TCPPacket() {
        super(TCP_PROTOCOL_NUMBER);
        // 设置默认数据偏移量（头部长度）
        setDataOffset(5); // 5 * 4 = 20字节
    }
    
    // ==================== 端口相关方法 ====================
    
    public int getSrcPort() {
        return header.getShort(0) & 0xFFFF;
    }
    
    public void setSrcPort(int srcPort) {
        if (srcPort < 0 || srcPort > 65535) {
            throw new IllegalArgumentException("Source port must be between 0 and 65535");
        }
        header.putShort(0, (short) srcPort);
    }
    
    public int getDstPort() {
        return header.getShort(2) & 0xFFFF;
    }
    
    public void setDstPort(int dstPort) {
        if (dstPort < 0 || dstPort > 65535) {
            throw new IllegalArgumentException("Destination port must be between 0 and 65535");
        }
        header.putShort(2, (short) dstPort);
    }
    
    // ==================== 序列号相关方法 ====================
    
    public long getSeqNumber() {
        return header.getInt(4) & 0xFFFFFFFFL;
    }
    
    public void setSeqNumber(long seqNumber) {
        if (seqNumber < 0 || seqNumber > 0xFFFFFFFFL) {
            throw new IllegalArgumentException("Sequence number must be between 0 and 2^32-1");
        }
        header.putInt(4, (int) seqNumber);
    }
    
    public long getAckNumber() {
        return header.getInt(8) & 0xFFFFFFFFL;
    }
    
    public void setAckNumber(long ackNumber) {
        if (ackNumber < 0 || ackNumber > 0xFFFFFFFFL) {
            throw new IllegalArgumentException("Acknowledgment number must be between 0 and 2^32-1");
        }
        header.putInt(8, (int) ackNumber);
    }
    
    // ==================== 数据偏移和保留位 ====================
    
    public int getDataOffset() {
        return (header.get(12) >>> 4) & 0x0F;
    }
    
    public void setDataOffset(int dataOffset) {
        if (dataOffset < 5 || dataOffset > 15) {
            throw new IllegalArgumentException("Data offset must be between 5 and 15");
        }
        header.put(12, (byte) ((header.get(12) & 0x0F) | ((dataOffset & 0x0F) << 4)));
    }
    
    public boolean isNS() {
        return (header.get(12) & 0x01) != 0;
    }
    
    public void setNS(boolean ns) {
        if (ns) {
            header.put(12, (byte) (header.get(12) | 0x01));
        } else {
            header.put(12, (byte) (header.get(12) & 0xFE));
        }
    }
    
    // ==================== 标志位方法 ====================
    
    public int getFlags() {
        return header.get(13) & 0xFF;
    }
    
    public void setFlags(int flags) {
        header.put(13, (byte) flags);
    }
    
    public boolean isCWR() {
        return (getFlags() & FLAG_CWR) != 0;
    }
    
    public void setCWR(boolean cwr) {
        setFlag(FLAG_CWR, cwr);
    }
    
    public boolean isECE() {
        return (getFlags() & FLAG_ECE) != 0;
    }
    
    public void setECE(boolean ece) {
        setFlag(FLAG_ECE, ece);
    }
    
    public boolean isURG() {
        return (getFlags() & FLAG_URG) != 0;
    }
    
    public void setURG(boolean urg) {
        setFlag(FLAG_URG, urg);
    }
    
    public boolean isACK() {
        return (getFlags() & FLAG_ACK) != 0;
    }
    
    public void setACK(boolean ack) {
        setFlag(FLAG_ACK, ack);
    }
    
    public boolean isPSH() {
        return (getFlags() & FLAG_PSH) != 0;
    }
    
    public void setPSH(boolean psh) {
        setFlag(FLAG_PSH, psh);
    }
    
    public boolean isRST() {
        return (getFlags() & FLAG_RST) != 0;
    }
    
    public void setRST(boolean rst) {
        setFlag(FLAG_RST, rst);
    }
    
    public boolean isSYN() {
        return (getFlags() & FLAG_SYN) != 0;
    }
    
    public void setSYN(boolean syn) {
        setFlag(FLAG_SYN, syn);
    }
    
    public boolean isFIN() {
        return (getFlags() & FLAG_FIN) != 0;
    }
    
    public void setFIN(boolean fin) {
        setFlag(FLAG_FIN, fin);
    }
    
    private void setFlag(int flag, boolean set) {
        int flags = getFlags();
        if (set) {
            flags |= flag;
        } else {
            flags &= ~flag;
        }
        setFlags(flags);
    }
    
    // ==================== 窗口大小 ====================
    
    public int getWindowSize() {
        return header.getShort(14) & 0xFFFF;
    }
    
    public void setWindowSize(int windowSize) {
        if (windowSize < 0 || windowSize > 65535) {
            throw new IllegalArgumentException("Window size must be between 0 and 65535");
        }
        header.putShort(14, (short) windowSize);
    }
    
    // ==================== 校验和 ====================
    
    public int getChecksum() {
        return header.getShort(16) & 0xFFFF;
    }
    
    public void setChecksum(int checksum) {
        header.putShort(16, (short) checksum);
    }
    
    // ==================== 紧急指针 ====================
    
    public int getUrgentPointer() {
        return header.getShort(18) & 0xFFFF;
    }
    
    public void setUrgentPointer(int urgentPointer) {
        if (urgentPointer < 0 || urgentPointer > 65535) {
            throw new IllegalArgumentException("Urgent pointer must be between 0 and 65535");
        }
        header.putShort(18, (short) urgentPointer);
    }
    
    // ==================== 选项处理 ====================
    
    public List<TCPOptionTLV> getOptions() {
        return new ArrayList<>(options);
    }
    
    public void addOption(TCPOptionTLV option) {
        options.add(option);
    }
    
    public void removeOption(TCPOptionTLV option) {
        options.remove(option);
    }
    
    public TCPOptionTLV findOptionByType(int type) {
        for (TCPOptionTLV option : options) {
            if (option.getType() == type) {
                return option;
            }
        }
        return null;
    }
    
    /**
     * 添加MSS选项
     */
    public void addMssOption(int mss) {
        addOption(new TCPMssOption(mss));
    }
    
    /**
     * 获取MSS值
     */
    public int getMssValue() {
        TCPMssOption mssOption = (TCPMssOption) findOptionByType(TCPMssOption.MAXIMUM_SEGMENT_SIZE);
        return mssOption != null ? mssOption.getMssValue() : 0;
    }
    
    /**
     * 添加窗口缩放选项
     */
    public void addWindowScaleOption(int scaleFactor) {
        addOption(new TCPWindowScaleOption(scaleFactor));
    }
    
    /**
     * 获取窗口缩放因子
     */
    public int getWindowScaleValue() {
        TCPWindowScaleOption scaleOption = (TCPWindowScaleOption) findOptionByType(TCPWindowScaleOption.WINDOW_SCALE);
        return scaleOption != null ? scaleOption.getScaleFactor() : 0;
    }
    
    /**
     * 添加SACK允许选项
     */
    public void addSackPermittedOption() {
        addOption(new TCPSackPermittedOption());
    }
    
    /**
     * 添加时间戳选项
     */
    public void addTimestampOption(int tsValue, int tsEchoReply) {
        addOption(new TCPTimestampOption(tsValue, tsEchoReply));
    }
    
    /**
     * 获取时间戳值
     */
    public long getTimestampValue() {
        TCPTimestampOption tsOption = (TCPTimestampOption) findOptionByType(TCPTimestampOption.TIMESTAMP);
        return tsOption != null ? tsOption.getTimestampValue() : 0;
    }
    
    /**
     * 获取时间戳回显值
     */
    public long getTimestampEchoReply() {
        TCPTimestampOption tsOption = (TCPTimestampOption) findOptionByType(TCPTimestampOption.TIMESTAMP);
        return tsOption != null ? tsOption.getTimestampEchoReply() : 0;
    }
    
    // ==================== 数据部分处理 ====================
    
    @Override
    public long getTotalLength() {
        int headerLength = TCP_HEADER_MIN_LENGTH + calculateOptionsLength();
        return headerLength + (getData() != null ? getData().limit() : 0);
    }
    
    public int getPayloadLength() {
        return getData() != null ? getData().limit() : 0;
    }
    
    // ==================== 网络读写方法 ====================
    
    @Override
    public void writeToChannel(WritableByteChannel dto) throws IOException {
        // 写入TCP基本头部（20字节）
        header.position(0);
        header.limit(TCP_HEADER_MIN_LENGTH);
        dto.write(header);
        
        // 写入TCP选项
        for (TCPOptionTLV option : options) {
            TCPOptionTLV.writeTCPOptionTLVToChannel(dto, option);
        }
        
        // 写入选项填充（如果需要）
        int optionsLength = calculateOptionsLength();
        int paddingNeeded = (4 - (optionsLength % 4)) % 4;
        for (int i = 0; i < paddingNeeded; i++) {
            ByteBuffer nop = ByteBuffer.allocate(1);
            nop.put((byte) TCPNopOption.NO_OPERATION);
            nop.flip();
            dto.write(nop);
        }
        
        // 写入TCP数据负载
        if (getData() != null && getData().limit() > 0) {
            getData().position(0);
            dto.write(getData());
        }
    }
    
    @Override
    public void readFromChannel(ReadableByteChannel din, long length) throws IOException {
        // 读取TCP基本头部（20字节）
        header.clear();
        header.limit(TCP_HEADER_MIN_LENGTH);
        while (header.hasRemaining()) {
            if (din.read(header) == -1) {
                throw new EOFException();
            }
        }
        header.flip();
        
        // 计算TCP头部总长度和选项长度
        int headerLength = getDataOffset() * 4;
        int optionsLength = headerLength - TCP_HEADER_MIN_LENGTH;
        
        if (headerLength < TCP_HEADER_MIN_LENGTH || headerLength > (TCP_HEADER_MIN_LENGTH + TCP_OPTIONS_MAX_LENGTH)) {
            throw new IOException("Invalid TCP header length: " + headerLength);
        }
        
        // 读取TCP选项
        options.clear();
        if (optionsLength > 0) {
            readOptionsFromChannel(din, optionsLength);
        }
        
        // 读取TCP数据负载
        int payloadLength = (int) length - headerLength;
        if (payloadLength > 0) {
            if (getData() == null) {
            //    setData(NetworkPacket.bufferAllocator.allocate(payloadLength));
            }
            getData().clear();
            getData().limit(payloadLength);
            while (getData().hasRemaining()) {
                if (din.read(getData()) == -1) {
                    throw new EOFException();
                }
            }
            getData().flip();
        }
    }
    
    /**
     * 从通道读取TCP选项
     */
    private void readOptionsFromChannel(ReadableByteChannel din, int optionsLength) throws IOException {
        int bytesRead = 0;
        
        while (bytesRead < optionsLength) {
            TCPOptionTLV option = TCPOptionTLV.readTCPOptionTLVFromChannel(din);
            if (option == null) {
                break;
            }
            
            bytesRead += option.getTotalLength();
            options.add(option);
            
            // 如果遇到End选项，停止读取
            if (option.getType() == TCPEndOption.END_OF_OPTION_LIST) {
                break;
            }
        }
        
        // 跳过剩余的选项字节（如果有）
        int remaining = optionsLength - bytesRead;
        if (remaining > 0) {
            ByteBuffer skipBuffer = NetworkPacket.bufferAllocator.allocate(remaining);
            skipBuffer.limit(remaining);
            while (skipBuffer.hasRemaining()) {
                if (din.read(skipBuffer) == -1) {
                    throw new EOFException();
                }
            }
        }
    }
    
    // ==================== 辅助方法 ====================
    
    private int calculateOptionsLength() {
        int total = 0;
        for (TCPOptionTLV option : options) {
            total += (int) option.getTotalLength();
        }
        return total;
    }
    
    // ==================== 工具方法 ====================
    
    @Override
    public String toString() {
        return String.format(
            "TCPPacket[srcPort=%d, dstPort=%d, seq=%d, ack=%d, flags=%s, window=%d, options=%d, payload=%d bytes]",
            getSrcPort(), getDstPort(), getSeqNumber(), getAckNumber(), getFlagsString(), 
            getWindowSize(), options.size(), getPayloadLength()
        );
    }
    
    private String getFlagsString() {
        StringBuilder sb = new StringBuilder();
        if (isFIN()) sb.append("FIN ");
        if (isSYN()) sb.append("SYN ");
        if (isRST()) sb.append("RST ");
        if (isPSH()) sb.append("PSH ");
        if (isACK()) sb.append("ACK ");
        if (isURG()) sb.append("URG ");
        if (isECE()) sb.append("ECE ");
        if (isCWR()) sb.append("CWR ");
        if (isNS()) sb.append("NS ");
        return sb.toString().trim();
    }
    
    @Override
    protected boolean needEndPosition() {
        return true;
    }
}