pub use crate::decoder::*;
pub use crate::error::*;
pub use crate::marker::*;
pub use crate::parser::CodingProcess;
use std::fmt::Write;

mod byteorder;
mod decoder;
mod error;
mod marker;
mod parser;

#[cfg(target_arch = "wasm32")]
use wasm_bindgen::prelude::*;

#[derive(Debug, Clone, Default, PartialEq)]
pub struct Scans {
    pub metadata_end: Option<usize>,
    pub dc_end: Option<usize>,
    pub good_scan_end: Option<usize>,
}

pub fn scans(input_file: &[u8]) -> Result<Scans> {
    let mut decoder = Decoder::new(input_file);
    let markers = decoder.decode()?;

    let mut found = Scans::default();
    let mut seen_sof = false;
    let mut seen_dc = false;
    let mut number_of_components = 3;
    let mut is_progressive = false;
    let mut dc_components_seen = [false; 3];
    let mut luma_coeff_bits_missing = [255u8; 16]; // there are 64 coefs, but only the first 16 are important for us

    for m in markers {
        // markers specify only their starting position, so to know the end position,
        // we need to wait until the next marker.
        if seen_sof {
            if found.metadata_end.is_none() {
                found.metadata_end = Some(m.position);
                if !is_progressive {
                    break;
                }
            } else if seen_dc {
                if found.dc_end.is_none() {
                    found.dc_end = Some(m.position);
                }
                // mark scan as reaching good-enough quality when either enough bits were sent (the cuttoff value is empirical),
                // or we've crossed half of the file, so we have to stop searching anyway.
                else if found.good_scan_end.is_none() && (m.position >= input_file.len()/2 || coeff_fill_factor(&luma_coeff_bits_missing) >= 91) {
                    found.good_scan_end = Some(m.position);
                    break;
                }
            }
        }

        match m.marker {
            // Image dimensions
            MarkerData::Frame(frame) => {
                number_of_components = frame.components.len();
                is_progressive = frame.coding_process == CodingProcess::DctProgressive;
                seen_sof = true;
            },
            // Progressive scan
            MarkerData::Scan(scan) => {
                // DC is 0th component
                if *scan.spectral_selection.start() == 0 {
                    for comp in scan.component_indices.iter().copied() {
                        if let Some(seen) = dc_components_seen.get_mut(comp as usize) {
                            *seen = true;
                        }
                    }
                }
                // DC may be split into multiple scans. Don't show incomplete DC to avoid green faces problem.
                if dc_components_seen.iter().take(number_of_components).all(|&v| v) {
                    // We don't care how many bits of DC were sent. As soon as there is something in full color, we're good.
                    seen_dc = true;
                }

                // Measure only quality of luma, assuming encoders know how to sensibly interleave chroma
                if scan.component_indices.contains(&0) {
                    for bit in scan.spectral_selection {
                        if (bit as usize) < luma_coeff_bits_missing.len() {
                            luma_coeff_bits_missing[bit as usize] = luma_coeff_bits_missing[bit as usize].min(scan.successive_approximation_low);
                        }
                    }
                }
            },
            MarkerData::Other(_) => {},
        }
    }
    Ok(found)
}

/// Checks if there are enough coefficient bits sent for the image to look good enough.
/// This is entirely speculative. Allows either few precise ACs or many OK-ish ACs.
/// It is *not* JPEG's quality scale!
/// FIXME: this should be scaled by quantization table
fn coeff_fill_factor(bits_missing: &[u8; 16]) -> u16 {
    let mut coeff_fill_factor = 0;
    // includes DC, because posterization is also an issue
    for missing in bits_missing {
        coeff_fill_factor += match missing {
            0 => 12,
            1 => 8,
            2 => 5,
            3 => 1,
            _ => 0,
        }
    }
    coeff_fill_factor
}

/// Assumes the priorities are:
///
///  * 50 = critical js/css, fonts + image metadata
///  * 30 = regular js/css, followed by other images + JPEG DC
///  * 20 = low-priority image bodies
///  * 10 = idle
///
/// Returns `cf-priority` and `cf-priority-change`
pub fn cf_priority_change_header(image: &[u8]) -> Result<(String, String)> {
    cf_priority_change_header_from_scans(scans(image)?, image.len())
}


#[cfg(target_arch = "wasm32")]
#[global_allocator]
static ALLOC: wee_alloc::WeeAlloc = wee_alloc::WeeAlloc::INIT;

#[cfg(target_arch = "wasm32")]
#[wasm_bindgen]
pub fn cf_priority_change_header_wasm(image: &[u8]) -> Option<String> {
    cf_priority_change_header(image).ok()
}

fn cf_priority_change_header_from_scans(s: Scans, image_len: usize) -> Result<(String, String)> {
    let mut pri_header = "30/1"; // default for images
    let mut change_header = String::with_capacity(127);

    if let Some(metadata_end) = s.metadata_end {
        // if there's a fat color profile or Adobe Garbage,
        // then sending "just" the metadata is not fast, and shouldn't be prioritized.
        let fat_metadata_limit = ((image_len / 8) + 100).min(2000);
        if metadata_end < fat_metadata_limit {
            pri_header = "50/0";
        }

        // if there's DC preview, then send it quickly. If it's a baseline image
        // (or progressive, but huge), it'll have to wait.
        let next_priority = match s.dc_end {
            Some(n) if n < 25000 => "30/0", // small DC can be downloaded in one go, to save on re-rendering
            Some(_) => "30/1",
            None if image_len < 2000 => "30/n", // tiny baseline image
            None => "20/n", // baseline image
        };
        let _ = write!(change_header, "{}:{}", metadata_end, next_priority);
    }
    if let Some(dc_end) = s.dc_end {
        if !change_header.is_empty() {change_header.push(',');}
        let is_next_scan_small = s.good_scan_end.map_or(false, |s| s < 100_000);
        let scan_concurrency = if is_next_scan_small { '1' } else { 'n' };
        let _ = write!(change_header, "{}:20/{}", dc_end, scan_concurrency);
    }
    if let Some(good_scan_end) = s.good_scan_end {
        if !change_header.is_empty() {change_header.push(',');}
        let file_is_small = image_len < 100_000;
        let scan_concurrency = if file_is_small { '1' } else { 'n' };
        let _ = write!(change_header, "{}:10/{}", good_scan_end, scan_concurrency);
    }
    if !change_header.is_empty() {
        Ok((String::from(pri_header), change_header))
    } else {
        Err(Error::Format("Can't find useful scans"))
    }
}

#[cfg(test)]
fn s(a: &str, b: &str) -> (String, String) { (a.into(), b.into()) }

#[test]
fn test_baseline() {
    assert!(cf_priority_change_header_from_scans(Scans {
        metadata_end: None,
        dc_end: None,
        good_scan_end: None,
    }, 100_000).is_err());

    // regular baseline image
    assert_eq!(s("50/0", "101:20/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(101),
        dc_end: None,
        good_scan_end: None,
    }, 100_000).unwrap());

    // tiny image
    assert_eq!(s("50/0", "101:30/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(101),
        dc_end: None,
        good_scan_end: None,
    }, 1000).unwrap());

    // fat metadata
    assert_eq!(s("30/1", "9999:20/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(9999),
        dc_end: None,
        good_scan_end: None,
    }, 100_000).unwrap());

    // relatively fat metadata
    assert_eq!(s("30/1", "1000:20/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(1000),
        dc_end: None,
        good_scan_end: None,
    }, 3000).unwrap());
}

#[test]
fn test_progressive() {
    assert_eq!(s("50/0", "1000:30/0,10000:20/n,100000:10/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(1_000),
        dc_end: Some(10_000),
        good_scan_end: Some(100_000),
    }, 200_000).unwrap());

    // fat metadata
    assert_eq!(s("30/1", "4000:30/0,10000:20/n,100000:10/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(4_000),
        dc_end: Some(10_000),
        good_scan_end: Some(100_000),
    }, 200_000).unwrap());

    // fat DC
    assert_eq!(s("50/0", "1000:30/1,50000:20/n,100000:10/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(1_000),
        dc_end: Some(50_000),
        good_scan_end: Some(100_000),
    }, 200_000).unwrap());

    // small good scan
    assert_eq!(s("50/0", "1000:30/0,10000:20/1,11000:10/n"), cf_priority_change_header_from_scans(Scans {
        metadata_end: Some(1_000),
        dc_end: Some(10_000),
        good_scan_end: Some(11_000),
    }, 200_000).unwrap());
}