import java.util.ArrayList;
import java.util.List;
import java.util.Scanner;

class IPFragment {
    int id;                 // Unique Identifier for the IP packet (same for all fragments)
    int offset;             // Fragment offset in 8-byte units
    byte[] fragmentData;    // Data of the fragment
    boolean moreFragments;  // Flag to indicate if more fragments follow
    int fragmentSize;       // Size of the fragment including the header

    public IPFragment(int id, int offset, byte[] fragmentData, boolean moreFragments, int fragmentSize) {
        this.id = id;
        this.offset = offset;
        this.fragmentData = fragmentData;
        this.moreFragments = moreFragments;
        this.fragmentSize = fragmentSize;
    }

    @Override
    public String toString() {
        return String.format("Fragment ID: %d, Offset: %d, More: %b, Fragment Size: %d",
                id, offset, moreFragments, fragmentSize);
    }
}

public class IPFragmenter {

    // Fragmentation algorithm
    public static List<IPFragment> fragmentPacket(byte[] originalPacket, int mtu, boolean doNotFragment) {
        List<IPFragment> fragments = new ArrayList<>();
        
        int headerSize = 20;  // Size of the IP header
        int dataSize = originalPacket.length - headerSize;
        
        // If the packet is smaller than or equal to the MTU, no fragmentation occurs
        if (doNotFragment || originalPacket.length <= mtu) {
            IPFragment singleFragment = new IPFragment(1, 0, originalPacket, false, originalPacket.length);
            fragments.add(singleFragment);
            return fragments;
        }

        int maxDataSize = mtu - headerSize;  // Max data size per fragment
        
        int fragmentId = (int) (Math.random() * 65536);  // Generate a random ID for the packet
        int offset = 0;
        boolean moreFragments = true;

        // Fragment the packet
        for (int i = 0; i < dataSize; i += maxDataSize) {
            int length = Math.min(maxDataSize, dataSize - i);
            byte[] fragmentData = new byte[length];
            System.arraycopy(originalPacket, headerSize + i, fragmentData, 0, length);

            // If it's the last fragment, set "moreFragments" to false
            moreFragments = (i + length < dataSize);

            // Create a fragment and add it to the list
            int fragmentSize = headerSize + fragmentData.length;
            IPFragment fragment = new IPFragment(fragmentId, offset, fragmentData, moreFragments, fragmentSize);
            fragments.add(fragment);

            offset += length / 8;  // Increase the offset by the number of 8-byte blocks
        }

        return fragments;
    }

    // Method to reassemble fragmented packets
    public static byte[] reassembleFragments(List<IPFragment> fragments) {
        // Sort fragments by offset to reconstruct the original order
        fragments.sort((f1, f2) -> Integer.compare(f1.offset, f2.offset));

        // Calculate total size of the reassembled packet
        int totalSize = 0;
        for (IPFragment fragment : fragments) {
            totalSize += fragment.fragmentData.length;
        }

        // Allocate space for the original packet (including the header)
        byte[] reassembledPacket = new byte[totalSize + 20];
        int currentOffset = 20;  // Start after the header

        // Copy data from fragments into the reassembled packet
        for (IPFragment fragment : fragments) {
            System.arraycopy(fragment.fragmentData, 0, reassembledPacket, currentOffset, fragment.fragmentData.length);
            currentOffset += fragment.fragmentData.length;
        }

        return reassembledPacket;
    }

    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);

        // Input the original packet size (for testing purposes)
        byte[] originalPacket = new byte[1500];  // Example large IP packet
        for (int i = 0; i < originalPacket.length; i++) {
            originalPacket[i] = (byte) (i % 256);  // Fill with some data
        }

        // Get MTU from the user
        System.out.print("Enter MTU (Maximum Transmission Unit): ");
        int mtu = scanner.nextInt();

        // Get the DoNotFragment flag from the user
        System.out.print("Do you want to allow fragmentation (yes/no)? ");
        boolean doNotFragment = !scanner.next().equalsIgnoreCase("yes");

        // Fragment the original packet
        List<IPFragment> fragments = fragmentPacket(originalPacket, mtu, doNotFragment);
        System.out.println("\nFragments:");
        for (IPFragment fragment : fragments) {
            System.out.println(fragment);
        }

        // Reassemble the packet
        byte[] reassembledPacket = reassembleFragments(fragments);
        System.out.println("\nReassembled Packet Size: " + reassembledPacket.length);
        
        scanner.close();
    }
}