//! What the PDF writer has no operators for — optional content (PDF layers) — drawn as form
//! XObjects of their own (without a transparency group, so they draw as if in place), each
//! starting with a mark: an empty clip whose first point says what the form is. Marked forms also
//! stay apart (the writer shares identical forms). [`finish`] completes them in the written file:
//! a layer's form is drawn as the layer's optional content (`/OC … BDC`).

use std::collections::HashMap;
use std::io::Read;

use krilla::geom::PathBuilder;
use krilla::graphic::Graphic;
use krilla::surface::Surface;
use vectorcraft_doc::{Document, NodeKind};

use crate::encrypt::{Lexer, Obj, StreamSpans, Tok, int, stream_spans};
use crate::output::{deflate, text_string};
use crate::patch::{Patch, Xref};
use crate::{PdfError, PdfSettings, Standard};

/// What a marked form is.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Mark {
    /// The art of top-level layer `i` (`Document::layers`).
    Layer(usize),
}

/// The x coordinate of a layer mark's first point (its y is the layer index).
const LAYER_X: i32 = -32000;

impl Mark {
    fn at(self) -> (f32, f32) {
        match self {
            // Exact in `f32` up to 2^24 layers.
            Self::Layer(i) => (LAYER_X as f32, i as f32),
        }
    }
}

/// Draw `draw` as a form of its own, marked `mark`.
pub(crate) fn form(s: &mut Surface, mark: Mark, draw: impl FnOnce(&mut Surface)) {
    let mut sb = s.stream_builder();
    let mut fs = sb.surface();
    let (x, y) = mark.at();
    let mut pb = PathBuilder::new();
    pb.move_to(x, y);
    pb.line_to(x + 1.0, y);
    pb.line_to(x + 1.0, y + 1.0);
    pb.close();
    // A form's drawing space is its own: the mark's points are written as they are.
    if let Some(p) = pb.finish() {
        fs.push_clip_path(&p, &krilla::paint::FillRule::NonZero);
        fs.pop();
    }
    draw(&mut fs);
    fs.finish();
    let stream = sb.finish();
    s.draw_graphic(Graphic::new(stream, false));
}

// ---------- finishing the written file ----------

/// The prefix of the resource names of the layers' optional content groups (then the index).
const LAYER_OC: &str = "VCoc";
/// The most bytes of a form's content read to find its mark.
const HEAD: u64 = 64;

fn bad(why: &str) -> PdfError {
    PdfError::Write(format!("can't write PDF layers: {why}"))
}

/// An object of the written file with a dictionary, and its stream's [`StreamSpans`].
struct Object<'a> {
    dict: Obj<'a>,
    stream: Option<StreamSpans>,
}

impl Object<'_> {
    /// The stream's data, decoded (Flate or none), at most `limit` bytes.
    fn data(&self, pdf: &[u8], limit: u64) -> Option<Vec<u8>> {
        let (_, (start, end)) = self.stream?;
        let data = pdf.get(start..end)?;
        let mut out = vec![];
        match self.dict.get(b"Filter") {
            None => out.extend_from_slice(data.get(..usize::try_from(limit).unwrap_or(usize::MAX).min(data.len()))?),
            Some(Obj::Name(b"FlateDecode")) => {
                flate2::read::ZlibDecoder::new(data).take(limit).read_to_end(&mut out).ok()?;
            }
            Some(_) => return None,
        }
        Some(out)
    }

    /// Replace the stream's data with `content` (compressed as before).
    fn rewrite(&self, patch: &mut Patch, content: &[u8]) -> Result<(), PdfError> {
        let ((len_start, len_end), (start, end)) = self.stream.ok_or_else(|| bad("a content stream isn't a stream"))?;
        let data = if self.dict.get(b"Filter").is_some() { deflate(content).map_err(|e| bad(&e.to_string()))? } else { content.to_vec() };
        patch.replace(len_start, len_end, data.len().to_string().into_bytes());
        patch.replace(start, end, data);
        Ok(())
    }

    /// The mark of a marked form.
    fn mark(&self, pdf: &[u8]) -> Option<Mark> {
        if self.dict.name(b"Subtype") != Some(b"Form") {
            return None;
        }
        let head = self.data(pdf, HEAD)?;
        let mut lx = Lexer::at(&head, 0);
        let mut word = || match lx.next() {
            Some((Tok::Word(w), ..)) => Some(w),
            _ => None,
        };
        let (Some(b"q"), Some(x), Some(y), Some(b"m")) = (word(), word().and_then(int), word().and_then(int), word()) else { return None };
        match x {
            _ if x == i64::from(LAYER_X) => usize::try_from(y).ok().map(Mark::Layer),
            _ => None,
        }
    }
}

/// The objects of the written file that have a dictionary, by number.
fn objects<'a>(pdf: &'a [u8], xref: &Xref) -> HashMap<u32, Object<'a>> {
    (1..xref.count())
        .filter_map(|n| {
            let n = u32::try_from(n).ok()?;
            let mut lx = Lexer::at(pdf, xref.offset(pdf, n)?);
            (lx.header()?.0 == n).then_some(())?;
            let dict = lx.object(0)?;
            let Obj::Dict { .. } = dict else { return None };
            let stream = stream_spans(pdf, &mut lx, &dict).ok().flatten();
            Some((n, Object { dict, stream }))
        })
        .collect()
}

/// The resource dictionary of the content of `o` (a page or a form).
fn resources<'o>(objects: &'o HashMap<u32, Object<'_>>, o: &'o Object<'_>) -> Option<&'o Obj<'o>> {
    match o.dict.get(b"Resources")? {
        Obj::Ref(n, _) => objects.get(n).map(|r| &r.dict),
        d => Some(d),
    }
}

/// Add `entries` (`/VCoc0 12 0 R`) to the `category` (`Properties`) of resource dictionary `res`.
fn add_resources(patch: &mut Patch, res: &Obj<'_>, category: &str, entries: &str) -> Result<(), PdfError> {
    match (res.get(category.as_bytes()), res) {
        (Some(Obj::Dict { start, .. }), _) => patch.replace(start + 2, start + 2, entries.as_bytes().to_vec()),
        (None, Obj::Dict { start, .. }) => patch.replace(start + 2, start + 2, format!("/{category}<<{entries}>>").into_bytes()),
        _ => return Err(bad("unexpected resources")),
    }
    Ok(())
}

/// `content` with each `Do` of a layer's form (`forms`: its resource name → the layer) marked as
/// that layer's optional content, and the layers marked; `None` when there is none.
fn mark_layers(content: &[u8], forms: &HashMap<&[u8], usize>) -> Option<(Vec<u8>, Vec<usize>)> {
    let mut lx = Lexer::at(content, 0);
    let mut name = None;
    let mut spans = vec![];
    while let Some((tok, start, end)) = lx.next() {
        name = match tok {
            Tok::Name(n) => forms.get(n).map(|layer| (start, *layer)),
            Tok::Word(b"Do") => {
                spans.extend(name.map(|(name, layer)| (name, end, layer)));
                None
            }
            // Inline image data isn't PDF syntax (the writer writes none).
            Tok::Word(b"ID") => break,
            _ => None,
        };
    }
    if spans.is_empty() {
        return None;
    }
    let mut out = Vec::with_capacity(content.len() + spans.len() * 24);
    let mut at = 0;
    for (start, end, layer) in &spans {
        out.extend_from_slice(content.get(at..*start)?);
        out.extend_from_slice(format!("/OC /{LAYER_OC}{layer} BDC ").as_bytes());
        out.extend_from_slice(content.get(*start..*end)?);
        out.extend_from_slice(b" EMC");
        at = *end;
    }
    out.extend_from_slice(content.get(at..)?);
    let mut layers: Vec<usize> = spans.iter().map(|s| s.2).collect();
    layers.sort_unstable();
    layers.dedup();
    Some((out, layers))
}

/// Finish the marked forms of `pdf` (as [`crate::export`] writes `doc` with `set`): with
/// `layers`, each top-level layer (but template layers) becomes an optional content group —
/// named as the layer, off when it is hidden, not printed (`/PrintState /OFF`) when its Print
/// option is off, locked when it is locked — and the uses of its forms its optional content.
pub(crate) fn finish(pdf: Vec<u8>, doc: &Document, set: &PdfSettings, layers: bool) -> Result<Vec<u8>, PdfError> {
    if !layers {
        return Ok(pdf);
    }
    let xref = Xref::read(&pdf).ok_or_else(|| bad("the written PDF has no cross-reference table"))?;
    let objects = objects(&pdf, &xref);
    let marks: HashMap<u32, Mark> = objects.iter().filter_map(|(n, o)| Some((*n, o.mark(&pdf)?))).collect();
    let mut patch = Patch::new(&xref);
    write_layers(&pdf, &xref, &objects, &marks, &mut patch, doc, set)?;
    let pdf = patch.apply(&pdf, &xref)?;
    // The file still is what its standard says.
    crate::pdfx::verify(&pdf, set.standard)?;
    Ok(pdf)
}

/// The optional content groups of `doc`'s top-level layers, the uses of their forms marked as
/// their content, and the catalog's `/OCProperties`.
fn write_layers(
    pdf: &[u8],
    xref: &Xref,
    objects: &HashMap<u32, Object<'_>>,
    marks: &HashMap<u32, Mark>,
    patch: &mut Patch,
    doc: &Document,
    set: &PdfSettings,
) -> Result<(), PdfError> {
    // One group per layer, in paint order (bottom first).
    let mut groups: HashMap<usize, u32> = HashMap::new();
    let (mut all, mut off, mut locked) = (vec![], vec![], vec![]);
    for (i, l) in doc.layers.iter().enumerate() {
        let NodeKind::Layer { template: false, printable, .. } = l.kind else { continue };
        let name = l.name.clone().unwrap_or_else(|| format!("Layer {}", i + 1));
        let on = |b: bool| if b { "ON" } else { "OFF" };
        let g = patch.add_object(
            format!("<</Type/OCG/Name{}/Usage<</View<</ViewState/{}>>/Print<</PrintState/{}>>>>>>", text_string(&name), on(l.visible), on(printable))
                .into_bytes(),
        );
        groups.insert(i, g);
        all.push(g);
        if !l.visible {
            off.push(g);
        }
        if l.locked {
            locked.push(g);
        }
    }
    // Every content stream drawing layers' forms: the pages' and the forms'.
    for o in objects.values() {
        let content = match o.dict.get(b"Contents") {
            Some(Obj::Ref(c, _)) if o.dict.name(b"Type") == Some(b"Page") => objects.get(c),
            _ if o.dict.name(b"Subtype") == Some(b"Form") => Some(o),
            _ => None,
        };
        let (Some(content), Some(res)) = (content, resources(objects, o)) else { continue };
        let Some(Obj::Dict { entries, .. }) = res.get(b"XObject") else { continue };
        let forms: HashMap<&[u8], usize> = entries
            .iter()
            .filter_map(|(name, x)| {
                let Obj::Ref(n, _) = x else { return None };
                match marks.get(n) {
                    Some(Mark::Layer(i)) if groups.contains_key(i) => Some((*name, *i)),
                    _ => None,
                }
            })
            .collect();
        if forms.is_empty() {
            continue;
        }
        let data = content.data(pdf, u64::MAX).ok_or_else(|| bad("a content stream can't be read"))?;
        let Some((new, marked)) = mark_layers(&data, &forms) else { continue };
        content.rewrite(patch, &new)?;
        let entries: String = marked.iter().filter_map(|i| Some(format!("/{LAYER_OC}{i} {} 0 R", groups.get(i)?))).collect();
        add_resources(patch, res, "Properties", &entries)?;
    }
    let refs = |v: &[u32]| v.iter().map(|g| format!("{g} 0 R")).collect::<Vec<_>>().join(" ");
    // Layers panels list the top layer first.
    let order: Vec<u32> = all.iter().rev().copied().collect();
    let mut config = format!("/Name(Layers)/Order[{}]", refs(&order));
    if !off.is_empty() {
        config.push_str(&format!("/OFF[{}]", refs(&off)));
    }
    if !locked.is_empty() {
        config.push_str(&format!("/Locked[{}]", refs(&locked)));
    }
    // Viewers apply the print states when printing (the standards allow no automatic states).
    if set.standard == Standard::None {
        config.push_str(&format!("/AS[<</Event/Print/OCGs[{}]/Category[/Print]>>]", refs(&all)));
    }
    let root = xref.trailer_ref(pdf, b"/Root").and_then(|n| objects.get(&n));
    let Some(Object { dict: Obj::Dict { start, .. }, .. }) = root else { return Err(bad("the written PDF has no catalog")) };
    patch.replace(start + 2, start + 2, format!("/OCProperties<</OCGs[{}]/D<<{config}>>>>", refs(&all)).into_bytes());
    Ok(())
}