use fastly::{Error, Request};

// Chunk size of each origin response read
const CHUNK_SIZE: usize = 4096;

// PNG magic bytes
const PNG_MARKER: &[u8] = &[0x89, 0x50, 0x4e, 0x47];
// IEND marker bytes, which identifies the end of the image
const PNG_IEND: &[u8] = &[0x00, 0x00, 0x00, 0x00, 0x49, 0x45, 0x4e, 0x44];

// Sequence search result
#[derive(PartialEq, Debug)]
enum SeqSearch {
    // Sequence found, parameter is index to carry over for next-round search
    Found(usize),

    // Sequence not found, parameter is index to carry over for next-round search
    NotFound(usize),
}

fn main() -> Result<(), Error> {
    let req = Request::from_client();
    let mut backend_resp = req.send("origin_0")?;

    // Take the body so we can iterate through its lines later
    let mut backend_resp_body = backend_resp.take_body();

    // Start sending the backend response to the client with a now-empty body
    let mut client_body = backend_resp.stream_to_client();

    // Search PNG magic first, and then search IEND marker in the backend stream
    let mut found_png_magic = false;
    let mut found_png_iend = false;
    let mut buffer = Vec::new();
    for chunk_result in backend_resp_body.read_chunks(CHUNK_SIZE) {
        let Ok(mut chunk) = chunk_result else {
            // Abort operation if encountered any read error
            break;
        };

        // Merge previous carry-over buffer and current chunk for PNG marker,
        // which may span between two chunks.
        if !buffer.is_empty() {
            buffer.append(&mut chunk);
        } else {
            buffer = chunk;
        }

        if buffer.is_empty() {
            continue;
        }

        if !found_png_magic {
            // If we have not yet found a PNG marker, see if there's one
            let index = match seq_search(&buffer, PNG_MARKER) {
                SeqSearch::Found(index) => {
                    found_png_magic = true;
                    index
                }
                SeqSearch::NotFound(index) => index,
            };
            // Streaming buffer to client, without carry-over part
            client_body.write_bytes(&buffer[0..index]);

            buffer = buffer[index..].to_vec()
        } else {
            // We found the png magic sequence, looking for PNG_IEND marker
            let index = match seq_search(&buffer, PNG_IEND) {
                SeqSearch::Found(index) => {
                    found_png_iend = true;
                    index
                }
                SeqSearch::NotFound(index) => index,
            };

            // Streaming buffer to client, without carry-over part
            client_body.write_bytes(&buffer[0..index]);

            if found_png_iend {
                // We've past the end of the PNG, any remaining data will not be streamed
                break;
            }

            buffer = buffer[index..].to_vec();
        }
    }

    // If there is no PNG_IEND marker found, we streaming remaining carry-over buffer to client
    if !found_png_iend {
        client_body.write_bytes(&buffer);
    }

    // Finish the streaming body to close the client connection
    client_body.finish().ok();

    Ok(())
}

// Find a sequence of elements in a slice, returns search result enum.
// The enum parameter is the index from where carry-over buffer starts
// If we have below input
// buffer: [0123456789]  seq: [01234]
// we would return index 6, so buffer [6789] will be carried-over for next round of seq search
fn seq_search(buffer: &[u8], seq: &[u8]) -> SeqSearch {
    let seq_len = seq.len();
    let buffer_len = buffer.len();

    assert!(seq_len > 0 && buffer_len > 0);

    if seq_len > buffer_len {
        return SeqSearch::NotFound(buffer_len);
    }

    for ci in 0..=buffer_len - seq_len {
        let mut found = true;
        for si in 0..seq_len {
            if buffer[ci + si] != seq[si] {
                found = false;
                break;
            }
        }

        if found {
            return SeqSearch::Found(ci + seq_len);
        }
    }

    SeqSearch::NotFound(buffer_len - seq_len + 1)
}