use std::{
    borrow::Cow,
    collections::{BTreeMap, BTreeSet},
    fs::{File, OpenOptions},
    io::{BufWriter, Write},
    path::Path,
    rc::Rc,
};

use lopdf::{dictionary, Dictionary, Object, ObjectId, StringFormat};

use crate::{
    copy::{CopyOptions, ObjectSource},
    PdfOpsError, Result,
};

const INHERITABLE_PAGE_ATTRS: [&[u8]; 4] = [b"Resources", b"MediaBox", b"CropBox", b"Rotate"];
const MAX_COPY_DEPTH: usize = 256;

pub(crate) struct StreamingPdfWriter {
    output: CountingWriter<BufWriter<File>>,
    offsets: BTreeMap<u32, u32>,
    max_id: u32,
    pages_id: ObjectId,
    catalog_id: ObjectId,
    pages: Vec<ObjectId>,
}

impl StreamingPdfWriter {
    pub(crate) fn create(path: &Path) -> Result<Self> {
        let file = OpenOptions::new().write(true).create_new(true).open(path)?;
        let mut output = CountingWriter::new(BufWriter::new(file));
        output.write_all(b"%PDF-1.7\n%\xC7\xEC\x8F\xA2\n")?;

        let mut writer = Self {
            output,
            offsets: BTreeMap::new(),
            max_id: 0,
            pages_id: (0, 0),
            catalog_id: (0, 0),
            pages: Vec::new(),
        };
        writer.pages_id = writer.new_object_id();
        writer.catalog_id = writer.new_object_id();
        Ok(writer)
    }

    pub(crate) fn new_object_id(&mut self) -> ObjectId {
        self.max_id += 1;
        (self.max_id, 0)
    }

    pub(crate) fn pages_id(&self) -> ObjectId {
        self.pages_id
    }

    pub(crate) fn extend_pages(&mut self, pages: impl IntoIterator<Item = ObjectId>) {
        self.pages.extend(pages);
    }

    pub(crate) fn write_object(&mut self, id: ObjectId, object: &Object) -> Result<()> {
        let offset = u32::try_from(self.output.bytes_written()).map_err(|_| {
            PdfOpsError::InvalidStructure("streaming output offset exceeds PDF xref limit".into())
        })?;
        self.offsets.insert(id.0, offset);
        writeln!(self.output, "{} {} obj", id.0, id.1)?;
        write_object(&mut self.output, object)?;
        self.output.write_all(b"\nendobj\n")?;
        Ok(())
    }

    pub(crate) fn finish(mut self) -> Result<()> {
        let kids = self
            .pages
            .iter()
            .copied()
            .map(Object::Reference)
            .collect::<Vec<_>>();
        let pages = dictionary! {
            "Type" => "Pages",
            "Kids" => Object::Array(kids),
            "Count" => self.pages.len() as i64,
        };
        self.write_object(self.pages_id, &Object::Dictionary(pages))?;

        let catalog = dictionary! {
            "Type" => "Catalog",
            "Pages" => self.pages_id,
        };
        self.write_object(self.catalog_id, &Object::Dictionary(catalog))?;

        let xref_start = self.output.bytes_written();
        self.output.write_all(b"xref\n")?;
        writeln!(self.output, "0 {}", self.max_id + 1)?;
        self.output.write_all(b"0000000000 65535 f \n")?;
        for id in 1..=self.max_id {
            let offset = self.offsets.get(&id).copied().ok_or_else(|| {
                PdfOpsError::InvalidStructure(format!("missing xref offset for object {id}"))
            })?;
            writeln!(self.output, "{offset:010} 00000 n ")?;
        }
        self.output.write_all(b"trailer\n")?;
        let trailer = dictionary! {
            "Size" => (self.max_id + 1) as i64,
            "Root" => self.catalog_id,
        };
        write_dictionary(&mut self.output, &trailer)?;
        write!(self.output, "\nstartxref\n{xref_start}\n%%EOF")?;
        self.output.flush()?;
        Ok(())
    }
}

pub(crate) struct StreamingCopyContext<'a> {
    writer: &'a mut StreamingPdfWriter,
    object_map: BTreeMap<ObjectId, ObjectId>,
    inherited_attrs_cache: BTreeMap<ObjectId, Rc<InheritedPageAttrs>>,
    selected_pages: BTreeSet<ObjectId>,
    options: CopyOptions,
}

impl<'a> StreamingCopyContext<'a> {
    pub(crate) fn new(writer: &'a mut StreamingPdfWriter, options: CopyOptions) -> Self {
        Self {
            writer,
            object_map: BTreeMap::new(),
            inherited_attrs_cache: BTreeMap::new(),
            selected_pages: BTreeSet::new(),
            options,
        }
    }

    pub(crate) fn copy_pages(
        &mut self,
        source: &impl ObjectSource,
        page_ids: &[ObjectId],
    ) -> Result<Vec<ObjectId>> {
        let mut copied_pages = Vec::with_capacity(page_ids.len());
        for page_id in page_ids {
            copied_pages.push(self.copy_page(source, *page_id)?);
        }
        Ok(copied_pages)
    }

    fn copy_page(&mut self, source: &impl ObjectSource, page_id: ObjectId) -> Result<ObjectId> {
        let new_id = if self.selected_pages.contains(&page_id) {
            self.copy_page_instance(source, page_id)?
        } else {
            let new_id = self.copy_object_at_depth(source, page_id, 0)?;
            self.selected_pages.insert(page_id);
            new_id
        };
        Ok(new_id)
    }

    fn copy_page_instance(
        &mut self,
        source: &impl ObjectSource,
        old_page_id: ObjectId,
    ) -> Result<ObjectId> {
        let object = source
            .get_object_value(old_page_id)
            .map_err(PdfOpsError::Pdf)?;
        let page = object
            .as_dict()
            .map_err(|_| PdfOpsError::InvalidStructure("page is not a dictionary".into()))?;
        if !page.has_type(b"Page") {
            return Err(PdfOpsError::InvalidStructure(
                "copied page does not have /Type /Page".into(),
            ));
        }

        let new_id = self.writer.new_object_id();
        let previous = self.object_map.insert(old_page_id, new_id);
        let copied = match self.copy_page_dictionary(source, old_page_id, page, 0) {
            Ok(copied) => copied,
            Err(err) => {
                restore_object_mapping(&mut self.object_map, old_page_id, previous);
                return Err(err);
            }
        };
        self.writer
            .write_object(new_id, &Object::Dictionary(copied))?;
        restore_object_mapping(&mut self.object_map, old_page_id, previous);
        Ok(new_id)
    }

    fn copy_object_at_depth(
        &mut self,
        source: &impl ObjectSource,
        old_id: ObjectId,
        depth: usize,
    ) -> Result<ObjectId> {
        check_copy_depth(depth)?;
        if let Some(new_id) = self.object_map.get(&old_id) {
            return Ok(*new_id);
        }

        let new_id = self.writer.new_object_id();
        self.object_map.insert(old_id, new_id);

        let object = match source.get_object_value(old_id) {
            Ok(object) => object,
            Err(lopdf::Error::ObjectNotFound(_)) => {
                self.writer.write_object(new_id, &Object::Null)?;
                return Ok(new_id);
            }
            Err(err) => return Err(PdfOpsError::Pdf(err)),
        };

        let copied = match object {
            Cow::Borrowed(object) => match object {
                Object::Dictionary(dict) if dict.has_type(b"Page") => Object::Dictionary(
                    self.copy_page_dictionary(source, old_id, dict, depth + 1)?,
                ),
                Object::Stream(stream) => {
                    Object::Stream(self.copy_stream(source, stream.clone(), depth + 1)?)
                }
                _ => self.copy_value(source, object, depth + 1)?,
            },
            Cow::Owned(object) => match object {
                Object::Dictionary(dict) if dict.has_type(b"Page") => Object::Dictionary(
                    self.copy_page_dictionary(source, old_id, &dict, depth + 1)?,
                ),
                Object::Stream(stream) => {
                    Object::Stream(self.copy_stream(source, stream, depth + 1)?)
                }
                object => self.copy_owned_value(source, object, depth + 1)?,
            },
        };
        self.writer.write_object(new_id, &copied)?;
        Ok(new_id)
    }

    fn copy_page_dictionary(
        &mut self,
        source: &impl ObjectSource,
        old_page_id: ObjectId,
        page: &Dictionary,
        depth: usize,
    ) -> Result<Dictionary> {
        check_copy_depth(depth)?;
        let mut copied = Dictionary::new();
        for (key, value) in page.iter() {
            if key.as_slice() == b"Parent" {
                continue;
            }
            if key.as_slice() == b"Annots" && !self.options.copy_annotations {
                continue;
            }
            copied.set(key.clone(), self.copy_value(source, value, depth + 1)?);
        }

        for key in INHERITABLE_PAGE_ATTRS {
            if copied.has(key) {
                continue;
            }
            if let Some(value) = self.inherited_attr(source, old_page_id, page, key)? {
                copied.set(key.to_vec(), self.copy_value(source, &value, depth + 1)?);
            }
        }
        copied.set("Parent", self.writer.pages_id());

        Ok(copied)
    }

    fn copy_value(
        &mut self,
        source: &impl ObjectSource,
        value: &Object,
        depth: usize,
    ) -> Result<Object> {
        check_copy_depth(depth)?;
        match value {
            Object::Reference(id) => Ok(Object::Reference(self.copy_object_at_depth(
                source,
                *id,
                depth + 1,
            )?)),
            Object::Array(items) => {
                let mut copied = Vec::with_capacity(items.len());
                for item in items {
                    copied.push(self.copy_value(source, item, depth + 1)?);
                }
                Ok(Object::Array(copied))
            }
            Object::Dictionary(dict) => {
                let mut copied = Dictionary::new();
                for (key, value) in dict.iter() {
                    copied.set(key.clone(), self.copy_value(source, value, depth + 1)?);
                }
                Ok(Object::Dictionary(copied))
            }
            Object::Stream(stream) => Ok(Object::Stream(self.copy_stream(
                source,
                stream.clone(),
                depth + 1,
            )?)),
            _ => Ok(value.clone()),
        }
    }

    fn copy_owned_value(
        &mut self,
        source: &impl ObjectSource,
        value: Object,
        depth: usize,
    ) -> Result<Object> {
        check_copy_depth(depth)?;
        match value {
            Object::Reference(id) => Ok(Object::Reference(self.copy_object_at_depth(
                source,
                id,
                depth + 1,
            )?)),
            Object::Array(items) => {
                let mut copied = Vec::with_capacity(items.len());
                for item in items {
                    copied.push(self.copy_owned_value(source, item, depth + 1)?);
                }
                Ok(Object::Array(copied))
            }
            Object::Dictionary(dict) => {
                let mut copied = Dictionary::new();
                for (key, value) in dict {
                    copied.set(key, self.copy_owned_value(source, value, depth + 1)?);
                }
                Ok(Object::Dictionary(copied))
            }
            Object::Stream(stream) => Ok(Object::Stream(self.copy_stream(
                source,
                stream,
                depth + 1,
            )?)),
            value => Ok(value),
        }
    }

    fn copy_stream(
        &mut self,
        source: &impl ObjectSource,
        mut stream: lopdf::Stream,
        depth: usize,
    ) -> Result<lopdf::Stream> {
        stream.dict = self.copy_dictionary(source, &stream.dict, depth + 1)?;
        Ok(stream)
    }

    fn copy_dictionary(
        &mut self,
        source: &impl ObjectSource,
        dict: &Dictionary,
        depth: usize,
    ) -> Result<Dictionary> {
        check_copy_depth(depth)?;
        let mut copied = Dictionary::new();
        for (key, value) in dict.iter() {
            copied.set(key.clone(), self.copy_value(source, value, depth + 1)?);
        }
        Ok(copied)
    }

    fn inherited_attr(
        &mut self,
        source: &impl ObjectSource,
        page_id: ObjectId,
        page: &Dictionary,
        key: &[u8],
    ) -> Result<Option<Object>> {
        if let Ok(value) = page.get(key) {
            return Ok(Some(value.clone()));
        }

        let mut current = match page.get(b"Parent") {
            Ok(Object::Reference(parent)) => *parent,
            Ok(_) => {
                return Err(PdfOpsError::InvalidStructure(
                    "page tree parent is not a reference".into(),
                ));
            }
            Err(_) => return Ok(None),
        };
        let mut visited = BTreeSet::from([page_id]);

        loop {
            if !visited.insert(current) {
                return Err(PdfOpsError::InvalidStructure(
                    "cycle detected while resolving inherited page attributes".into(),
                ));
            }

            let Some(attrs) = self.inherited_attrs(source, current)? else {
                return Ok(None);
            };
            if let Some(value) = attrs.get(key) {
                return Ok(Some(value.clone()));
            }
            match attrs.parent {
                ParentRef::Reference(parent) => current = parent,
                ParentRef::Missing => return Ok(None),
                ParentRef::Invalid => {
                    return Err(PdfOpsError::InvalidStructure(
                        "page tree parent is not a reference".into(),
                    ));
                }
            }
        }
    }

    fn inherited_attrs(
        &mut self,
        source: &impl ObjectSource,
        id: ObjectId,
    ) -> Result<Option<Rc<InheritedPageAttrs>>> {
        if let Some(attrs) = self.inherited_attrs_cache.get(&id) {
            return Ok(Some(Rc::clone(attrs)));
        }

        let object = match source.get_object_value(id) {
            Ok(object) => object,
            Err(lopdf::Error::ObjectNotFound(_)) => return Ok(None),
            Err(err) => return Err(PdfOpsError::Pdf(err)),
        };
        let dict = object.as_dict().map_err(|_| {
            PdfOpsError::InvalidStructure("page tree node is not a dictionary".into())
        })?;
        let attrs = Rc::new(InheritedPageAttrs::from_dict(dict));
        self.inherited_attrs_cache.insert(id, Rc::clone(&attrs));
        Ok(Some(attrs))
    }
}

#[derive(Debug, Default)]
struct InheritedPageAttrs {
    resources: Option<Object>,
    media_box: Option<Object>,
    crop_box: Option<Object>,
    rotate: Option<Object>,
    parent: ParentRef,
}

impl InheritedPageAttrs {
    fn from_dict(dict: &Dictionary) -> Self {
        Self {
            resources: dict.get(b"Resources").ok().cloned(),
            media_box: dict.get(b"MediaBox").ok().cloned(),
            crop_box: dict.get(b"CropBox").ok().cloned(),
            rotate: dict.get(b"Rotate").ok().cloned(),
            parent: match dict.get(b"Parent") {
                Ok(Object::Reference(parent)) => ParentRef::Reference(*parent),
                Ok(_) => ParentRef::Invalid,
                Err(_) => ParentRef::Missing,
            },
        }
    }

    fn get(&self, key: &[u8]) -> Option<&Object> {
        match key {
            b"Resources" => self.resources.as_ref(),
            b"MediaBox" => self.media_box.as_ref(),
            b"CropBox" => self.crop_box.as_ref(),
            b"Rotate" => self.rotate.as_ref(),
            _ => None,
        }
    }
}

#[derive(Debug, Clone, Copy, Default)]
enum ParentRef {
    Reference(ObjectId),
    #[default]
    Missing,
    Invalid,
}

fn restore_object_mapping(
    object_map: &mut BTreeMap<ObjectId, ObjectId>,
    old_id: ObjectId,
    previous: Option<ObjectId>,
) {
    if let Some(previous) = previous {
        object_map.insert(old_id, previous);
    } else {
        object_map.remove(&old_id);
    }
}

fn check_copy_depth(depth: usize) -> Result<()> {
    if depth > MAX_COPY_DEPTH {
        return Err(PdfOpsError::InvalidStructure(format!(
            "PDF object nesting exceeds maximum copy depth of {MAX_COPY_DEPTH}"
        )));
    }
    Ok(())
}

pub(crate) fn write_object(output: &mut dyn Write, object: &Object) -> std::io::Result<()> {
    match object {
        Object::Null => output.write_all(b"null"),
        Object::Boolean(value) => {
            if *value {
                output.write_all(b"true")
            } else {
                output.write_all(b"false")
            }
        }
        Object::Integer(value) => write!(output, "{value}"),
        Object::Real(value) => write!(output, "{value}"),
        Object::Name(name) => write_name(output, name),
        Object::String(text, format) => write_string(output, text, *format),
        Object::Array(items) => write_array(output, items),
        Object::Dictionary(dict) => write_dictionary(output, dict),
        Object::Stream(stream) => write_stream(output, stream),
        Object::Reference(id) => write!(output, "{} {} R", id.0, id.1),
    }
}

fn write_name(output: &mut dyn Write, name: &[u8]) -> std::io::Result<()> {
    output.write_all(b"/")?;
    for &byte in name {
        if b" \t\n\r\x0C()<>[]{}/%#".contains(&byte) || !(33..=126).contains(&byte) {
            write!(output, "#{byte:02X}")?;
        } else {
            output.write_all(&[byte])?;
        }
    }
    Ok(())
}

fn write_string(output: &mut dyn Write, text: &[u8], format: StringFormat) -> std::io::Result<()> {
    match format {
        StringFormat::Literal => {
            output.write_all(b"(")?;
            for &byte in text {
                match byte {
                    b'(' => output.write_all(b"\\(")?,
                    b')' => output.write_all(b"\\)")?,
                    b'\\' => output.write_all(b"\\\\")?,
                    b'\r' => output.write_all(b"\\r")?,
                    byte => output.write_all(&[byte])?,
                }
            }
            output.write_all(b")")
        }
        StringFormat::Hexadecimal => {
            output.write_all(b"<")?;
            for &byte in text {
                write!(output, "{byte:02X}")?;
            }
            output.write_all(b">")
        }
    }
}

fn write_array(output: &mut dyn Write, items: &[Object]) -> std::io::Result<()> {
    output.write_all(b"[")?;
    let mut first = true;
    for item in items {
        if first {
            first = false;
        } else if needs_separator(item) {
            output.write_all(b" ")?;
        }
        write_object(output, item)?;
    }
    output.write_all(b"]")
}

pub(crate) fn write_dictionary(
    output: &mut dyn Write,
    dictionary: &Dictionary,
) -> std::io::Result<()> {
    output.write_all(b"<<")?;
    for (key, value) in dictionary {
        write_name(output, key)?;
        if needs_separator(value) {
            output.write_all(b" ")?;
        }
        write_object(output, value)?;
    }
    output.write_all(b">>")
}

fn write_stream(output: &mut dyn Write, stream: &lopdf::Stream) -> std::io::Result<()> {
    write_dictionary(output, &stream.dict)?;
    output.write_all(b"stream\n")?;
    output.write_all(&stream.content)?;
    output.write_all(b"\nendstream")
}

fn needs_separator(object: &Object) -> bool {
    matches!(
        object,
        Object::Null
            | Object::Boolean(_)
            | Object::Integer(_)
            | Object::Real(_)
            | Object::Reference(_)
    )
}

struct CountingWriter<W: Write> {
    inner: W,
    bytes_written: usize,
}

impl<W: Write> CountingWriter<W> {
    fn new(inner: W) -> Self {
        Self {
            inner,
            bytes_written: 0,
        }
    }

    fn bytes_written(&self) -> usize {
        self.bytes_written
    }
}

impl<W: Write> Write for CountingWriter<W> {
    fn write(&mut self, buffer: &[u8]) -> std::io::Result<usize> {
        let bytes = self.inner.write(buffer)?;
        self.bytes_written += bytes;
        Ok(bytes)
    }

    fn write_all(&mut self, buffer: &[u8]) -> std::io::Result<()> {
        self.bytes_written += buffer.len();
        self.inner.write_all(buffer)
    }

    fn flush(&mut self) -> std::io::Result<()> {
        self.inner.flush()
    }
}
