use std::alloc::{alloc, dealloc, Layout};
use std::slice;

const HEADER_LEN: usize = 5;

#[unsafe(no_mangle)]
pub extern "C" fn wasm_alloc(size: u32) -> u32 {
    if size == 0 {
        return 0;
    }
    let layout = Layout::from_size_align(size as usize, 1).unwrap();
    unsafe { alloc(layout) as u32 }
}

#[unsafe(no_mangle)]
pub extern "C" fn wasm_dealloc(ptr: u32, size: u32) {
    if ptr == 0 || size == 0 {
        return;
    }
    let layout = Layout::from_size_align(size as usize, 1).unwrap();
    unsafe { dealloc(ptr as *mut u8, layout) }
}

fn parse_header(state: &[u8]) -> (u16, u16, u8) {
    let rows = u16::from_le_bytes([state[0], state[1]]);
    let cols = u16::from_le_bytes([state[2], state[3]]);
    let flags = state[4];
    (rows, cols, flags)
}

fn make_header(rows: u16, cols: u16, flags: u8) -> [u8; HEADER_LEN] {
    let mut hdr = [0u8; HEADER_LEN];
    hdr[0..2].copy_from_slice(&rows.to_le_bytes());
    hdr[2..4].copy_from_slice(&cols.to_le_bytes());
    hdr[4] = flags;
    hdr
}

fn restore_parser(state: &[u8], rows: u16, cols: u16) -> vt100::Parser {
    let mut parser = vt100::Parser::new(rows, cols, 0);
    if state.len() <= HEADER_LEN {
        return parser;
    }

    let flags = state[4];
    if flags & 1 != 0 && state.len() > HEADER_LEN + 4 {
        // Alt screen was active: [header][main_len:4][main_state][alt_state]
        let main_len =
            u32::from_le_bytes([state[5], state[6], state[7], state[8]]) as usize;
        let main_start = HEADER_LEN + 4;
        let main_end = main_start + main_len;

        if main_len > 0 && main_end <= state.len() {
            parser.process(&state[main_start..main_end]);
        }
        parser.process(b"\x1b[?1049h");
        if main_end < state.len() {
            parser.process(&state[main_end..]);
        }
    } else {
        parser.process(&state[HEADER_LEN..]);
    }

    parser
}

fn pack_result(ptr: u32, len: u32) -> u64 {
    ((ptr as u64) << 32) | (len as u64)
}

fn return_bytes(data: &[u8]) -> u64 {
    let len = data.len() as u32;
    if len == 0 {
        return 0;
    }
    let ptr = wasm_alloc(len);
    if ptr != 0 {
        unsafe {
            std::ptr::copy_nonoverlapping(data.as_ptr(), ptr as *mut u8, data.len());
        }
    }
    pack_result(ptr, len)
}

fn serialize_state(parser: &mut vt100::Parser, rows: u16, cols: u16) -> Vec<u8> {
    let is_alt = parser.screen().alternate_screen();
    let flags: u8 = if is_alt { 1 } else { 0 };
    let hdr = make_header(rows, cols, flags);

    if is_alt {
        let alt_formatted = parser.screen().state_formatted();
        parser.process(b"\x1b[?1049l");
        let main_formatted = parser.screen().state_formatted();
        let main_len = (main_formatted.len() as u32).to_le_bytes();

        let mut out =
            Vec::with_capacity(HEADER_LEN + 4 + main_formatted.len() + alt_formatted.len());
        out.extend_from_slice(&hdr);
        out.extend_from_slice(&main_len);
        out.extend_from_slice(&main_formatted);
        out.extend_from_slice(&alt_formatted);
        out
    } else {
        let formatted = parser.screen().state_formatted();
        let mut out = Vec::with_capacity(HEADER_LEN + formatted.len());
        out.extend_from_slice(&hdr);
        out.extend_from_slice(&formatted);
        out
    }
}

#[unsafe(no_mangle)]
pub extern "C" fn apply_delta(
    state_ptr: u32,
    state_len: u32,
    delta_ptr: u32,
    delta_len: u32,
    rows: u32,
    cols: u32,
) -> u64 {
    let rows = rows as u16;
    let cols = cols as u16;

    let state: Option<&[u8]> = if state_ptr == 0 || (state_len as usize) < HEADER_LEN {
        None
    } else {
        Some(unsafe { slice::from_raw_parts(state_ptr as *const u8, state_len as usize) })
    };

    let delta = if delta_ptr == 0 || delta_len == 0 {
        &[]
    } else {
        unsafe { slice::from_raw_parts(delta_ptr as *const u8, delta_len as usize) }
    };

    let mut parser = match state {
        Some(s) => {
            let (prev_rows, prev_cols, _) = parse_header(s);
            let mut p = restore_parser(s, prev_rows, prev_cols);
            if prev_rows != rows || prev_cols != cols {
                p.screen_mut().set_size(rows, cols);
            }
            p
        }
        None => vt100::Parser::new(rows, cols, 0),
    };

    parser.process(delta);

    return_bytes(&serialize_state(&mut parser, rows, cols))
}

#[unsafe(no_mangle)]
pub extern "C" fn screen_contents(state_ptr: u32, state_len: u32) -> u64 {
    if state_ptr == 0 || (state_len as usize) < HEADER_LEN {
        return 0;
    }
    let state = unsafe { slice::from_raw_parts(state_ptr as *const u8, state_len as usize) };
    let (rows, cols, _) = parse_header(state);
    let parser = restore_parser(state, rows, cols);
    let contents = parser.screen().contents();
    return_bytes(contents.as_bytes())
}

#[unsafe(no_mangle)]
pub extern "C" fn screen_contents_formatted(state_ptr: u32, state_len: u32) -> u64 {
    if state_ptr == 0 || (state_len as usize) < HEADER_LEN {
        return 0;
    }
    let state = unsafe { slice::from_raw_parts(state_ptr as *const u8, state_len as usize) };
    let (rows, cols, _) = parse_header(state);
    let parser = restore_parser(state, rows, cols);
    let formatted = parser.screen().contents_formatted();
    return_bytes(&formatted)
}
