use std::collections::HashMap;
use std::io::{self, Write};
use std::panic::{AssertUnwindSafe, catch_unwind};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};

use avril_core::serde_v8::{JsBuffer, ToJsBuffer};
use avril_core::{Extension, Op};
use brotli::{CompressorWriter, DecompressorWriter};
use flate2::{Compress, Compression, Decompress, FlushCompress, FlushDecompress, Status};
use futures::channel::oneshot;

const MAX_ID: u64 = 9_007_199_254_740_991;
const OUTPUT_BUFFER_SIZE: usize = 16 * 1024;

struct CodecOutput {
    bytes: Vec<u8>,
    terminal_error: bool,
}

impl CodecOutput {
    fn success(bytes: Vec<u8>) -> Self {
        Self { bytes, terminal_error: false }
    }
}

#[derive(Default)]
struct OutputSink(Vec<u8>);

impl Write for OutputSink {
    fn write(&mut self, bytes: &[u8]) -> io::Result<usize> {
        self.0.extend_from_slice(bytes);
        Ok(bytes.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

enum Codec {
    Compress {
        codec: Compress,
        finished: bool,
    },
    Decompress {
        codec: Decompress,
        finished: bool,
    },
    BrotliCompress {
        codec: Option<Box<CompressorWriter<OutputSink>>>,
        finished: bool,
    },
    BrotliDecompress {
        codec: Option<Box<DecompressorWriter<OutputSink>>>,
        finished: bool,
    },
}

// Brotli writer destructors can finish buffered input after a future is canceled.
impl Drop for Codec {
    fn drop(&mut self) {
        match self {
            Self::BrotliCompress { codec, .. } => {
                if let Some(writer) = codec.take() {
                    spawn_compression(move || drop(writer));
                }
            }
            Self::BrotliDecompress { codec, .. } => {
                if let Some(writer) = codec.take() {
                    spawn_compression(move || drop(writer));
                }
            }
            _ => {}
        }
    }
}

impl Codec {
    fn new(format: &str, decompress: bool) -> io::Result<Self> {
        // Brotli state is created by worker operations, not on the isolate thread.
        if format == "brotli" {
            return Ok(if decompress {
                Self::BrotliDecompress { codec: None, finished: false }
            } else {
                Self::BrotliCompress { codec: None, finished: false }
            });
        }
        if decompress {
            let codec = match format {
                "deflate" => Decompress::new(true),
                "deflate-raw" => Decompress::new(false),
                "gzip" => Decompress::new_gzip(15),
                _ => return Err(io::Error::other("Unsupported compression format")),
            };
            Ok(Self::Decompress { codec, finished: false })
        } else {
            let codec = match format {
                "deflate" => Compress::new(Compression::default(), true),
                "deflate-raw" => Compress::new(Compression::default(), false),
                "gzip" => Compress::new_gzip(Compression::default(), 15),
                _ => return Err(io::Error::other("Unsupported compression format")),
            };
            Ok(Self::Compress { codec, finished: false })
        }
    }

    fn transform(&mut self, input: &[u8]) -> io::Result<CodecOutput> {
        if input.is_empty() {
            return Ok(CodecOutput::success(Vec::new()));
        }
        match self {
            Self::Compress { codec, finished } => {
                if *finished {
                    return Err(io::Error::other("Compression stream is closed"));
                }
                run_compress(codec, input, FlushCompress::None, false).map(CodecOutput::success)
            }
            Self::Decompress { codec, finished } => {
                if *finished {
                    return Ok(CodecOutput {
                        bytes: Vec::new(),
                        terminal_error: true,
                    });
                }
                let (bytes, ended, trailing) = run_decompress(codec, input, FlushDecompress::None)?;
                if ended || trailing {
                    *finished = true;
                }
                Ok(CodecOutput { bytes, terminal_error: trailing })
            }
            Self::BrotliCompress { codec, finished } => {
                if *finished {
                    return Err(io::Error::other("Compression stream is closed"));
                }
                let writer = codec.get_or_insert_with(|| {
                    Box::new(CompressorWriter::new(
                        OutputSink::default(),
                        OUTPUT_BUFFER_SIZE,
                        5,
                        22,
                    ))
                });
                writer.write_all(input)?;
                Ok(CodecOutput::success(std::mem::take(
                    &mut writer.get_mut().0,
                )))
            }
            Self::BrotliDecompress { codec, finished } => {
                if *finished {
                    return Ok(CodecOutput {
                        bytes: Vec::new(),
                        terminal_error: true,
                    });
                }
                let writer = codec.get_or_insert_with(|| {
                    Box::new(DecompressorWriter::new(
                        OutputSink::default(),
                        OUTPUT_BUFFER_SIZE,
                    ))
                });
                let consumed = writer.write(input)?;
                let bytes = std::mem::take(&mut writer.get_mut().0);
                let trailing = consumed < input.len();
                if trailing {
                    *finished = true;
                }
                Ok(CodecOutput { bytes, terminal_error: trailing })
            }
        }
    }

    fn finish(&mut self) -> io::Result<CodecOutput> {
        match self {
            Self::Compress { codec, finished } => {
                if *finished {
                    return Ok(CodecOutput::success(Vec::new()));
                }
                let output = run_compress(codec, &[], FlushCompress::Finish, true)?;
                *finished = true;
                Ok(CodecOutput::success(output))
            }
            Self::Decompress { codec, finished } => {
                if *finished {
                    return Ok(CodecOutput::success(Vec::new()));
                }
                let (output, ended, _) = run_decompress(codec, &[], FlushDecompress::Finish)?;
                if !ended {
                    return Err(io::Error::other("Truncated compressed stream"));
                }
                *finished = true;
                Ok(CodecOutput::success(output))
            }
            Self::BrotliCompress { codec, finished } => {
                if *finished {
                    return Ok(CodecOutput::success(Vec::new()));
                }
                let writer = codec.take().unwrap_or_else(|| {
                    Box::new(CompressorWriter::new(
                        OutputSink::default(),
                        OUTPUT_BUFFER_SIZE,
                        5,
                        22,
                    ))
                });
                *finished = true;
                Ok(CodecOutput::success((*writer).into_inner().0))
            }
            Self::BrotliDecompress { codec, finished } => {
                if *finished {
                    return Ok(CodecOutput::success(Vec::new()));
                }
                let writer = codec.take().unwrap_or_else(|| {
                    Box::new(DecompressorWriter::new(
                        OutputSink::default(),
                        OUTPUT_BUFFER_SIZE,
                    ))
                });
                let output = (*writer)
                    .into_inner()
                    .map_err(|_| io::Error::other("Truncated Brotli stream"))?;
                *finished = true;
                Ok(CodecOutput::success(output.0))
            }
        }
    }
}

fn run_compress(
    codec: &mut Compress,
    input: &[u8],
    flush: FlushCompress,
    finishing: bool,
) -> io::Result<Vec<u8>> {
    let mut output = Vec::new();
    let mut consumed = 0;
    loop {
        let before_in = codec.total_in();
        let before_out = codec.total_out();
        let mut buffer = [0; OUTPUT_BUFFER_SIZE];
        let status = codec
            .compress(&input[consumed..], &mut buffer, flush)
            .map_err(io::Error::other)?;
        consumed += (codec.total_in() - before_in) as usize;
        let produced = (codec.total_out() - before_out) as usize;
        output.try_reserve(produced).map_err(io::Error::other)?;
        output.extend_from_slice(&buffer[..produced]);
        if status == Status::StreamEnd {
            if consumed != input.len() {
                return Err(io::Error::other(
                    "Compression ended before input was consumed",
                ));
            }
            return Ok(output);
        }
        if consumed == input.len() && produced < buffer.len() && !finishing {
            return Ok(output);
        }
        if codec.total_in() == before_in && codec.total_out() == before_out {
            return Err(io::Error::other("Compression made no progress"));
        }
    }
}

fn run_decompress(
    codec: &mut Decompress,
    input: &[u8],
    flush: FlushDecompress,
) -> io::Result<(Vec<u8>, bool, bool)> {
    let mut output = Vec::new();
    let mut consumed = 0;
    loop {
        let before_in = codec.total_in();
        let before_out = codec.total_out();
        let mut buffer = [0; OUTPUT_BUFFER_SIZE];
        let status = codec
            .decompress(&input[consumed..], &mut buffer, flush)
            .map_err(io::Error::other)?;
        consumed += (codec.total_in() - before_in) as usize;
        let produced = (codec.total_out() - before_out) as usize;
        output.try_reserve(produced).map_err(io::Error::other)?;
        output.extend_from_slice(&buffer[..produced]);
        if status == Status::StreamEnd {
            return Ok((output, true, consumed != input.len()));
        }
        if consumed == input.len() && produced < buffer.len() && flush != FlushDecompress::Finish {
            return Ok((output, false, false));
        }
        if codec.total_in() == before_in && codec.total_out() == before_out {
            if flush == FlushDecompress::Finish || consumed == input.len() {
                return Ok((output, false, false));
            }
            return Err(io::Error::other("Decompression made no progress"));
        }
    }
}

#[derive(Default)]
struct CodecStore {
    next_id: AtomicU64,
    entries: Mutex<HashMap<u64, Arc<Mutex<Codec>>>>,
}

impl CodecStore {
    fn insert(&self, codec: Codec) -> io::Result<u64> {
        let id = self
            .next_id
            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| {
                (id < MAX_ID).then_some(id + 1)
            })
            .map(|id| id + 1)
            .map_err(|_| io::Error::other("Compression resource IDs exhausted"))?;
        self.entries
            .lock()
            .map_err(|_| io::Error::other("Compression registry poisoned"))?
            .insert(id, Arc::new(Mutex::new(codec)));
        Ok(id)
    }

    fn get(&self, id: u64) -> io::Result<Arc<Mutex<Codec>>> {
        self.entries
            .lock()
            .map_err(|_| io::Error::other("Compression registry poisoned"))?
            .get(&id)
            .cloned()
            .ok_or_else(|| io::Error::other("Compression stream is closed"))
    }

    fn remove(&self, id: u64) {
        if let Ok(mut entries) = self.entries.lock() {
            entries.remove(&id);
        }
    }
}

fn spawn_compression(job: impl FnOnce() + Send + 'static) {
    compression_pool().spawn(move || {
        let _ = catch_unwind(AssertUnwindSafe(job));
    });
}

fn compression_pool() -> &'static rayon::ThreadPool {
    static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
    POOL.get_or_init(|| {
        rayon::ThreadPoolBuilder::new()
            .num_threads(4)
            .thread_name(|index| format!("avril:compression:{index}"))
            .build()
            .expect("failed to initialize compression worker pool")
    })
}

async fn run_operation(
    codec: Arc<Mutex<Codec>>,
    input: Option<Vec<u8>>,
    finish: bool,
    id: u64,
    store: Arc<CodecStore>,
) -> io::Result<(ToJsBuffer, bool)> {
    let (sender, receiver) = oneshot::channel();
    spawn_compression(move || {
        let result = catch_unwind(AssertUnwindSafe(|| {
            let mut codec = codec
                .lock()
                .map_err(|_| io::Error::other("Compression codec poisoned"))?;
            match input {
                Some(input) => codec.transform(&input),
                None => codec.finish(),
            }
        }))
        .unwrap_or_else(|_| Err(io::Error::other("Compression worker panicked")));
        if finish || result.is_err() || result.as_ref().is_ok_and(|output| output.terminal_error) {
            store.remove(id);
        }
        let _ = sender.send(result);
    });
    let output = receiver.await.map_err(io::Error::other)??;
    Ok((ToJsBuffer::from(output.bytes), output.terminal_error))
}

pub fn init() -> Extension {
    let store = Arc::new(CodecStore::default());
    let new_store = store.clone();
    let operation_store = store.clone();
    let finish_store = store.clone();
    let dispose_store = store.clone();
    Extension {
        name: "ext:compression",
        dependencies: &["ext:web", "ext:webstreams"],
        ops: vec![
            Op::new(
                "op_compression_new",
                move |(format, decompress): (String, bool)| {
                    new_store.insert(Codec::new(&format, decompress)?)
                },
            ),
            Op::asynchronous(
                "op_compression_transform",
                move |(id, input): (u64, JsBuffer)| {
                    let input = input.as_ref().to_vec();
                    let codec = operation_store.get(id);
                    let store = operation_store.clone();
                    async move {
                        let codec = codec?;
                        run_operation(codec, Some(input), false, id, store).await
                    }
                },
            ),
            Op::asynchronous("op_compression_finish", move |id: u64| {
                let codec = finish_store.get(id);
                let store = finish_store.clone();
                async move {
                    let codec = codec?;
                    run_operation(codec, None, true, id, store).await
                }
            }),
            Op::new("op_compression_dispose", move |id: u64| {
                dispose_store.remove(id);
                Ok::<_, io::Error>(())
            }),
        ],
        modules: &[(
            "ext:compression/main.js",
            include_str!(concat!(env!("OUT_DIR"), "/main.js")),
        )],
        entry_point: Some("ext:compression/main.js"),
    }
}