//! Bounded inputs and outputs; parser allocations and CPU are not process isolated.
use crate::SourceLocation;
use anyhow::{bail, Context, Result};
use quick_xml::{events::Event, name::ResolveResult, NsReader, XmlVersion};
use sha2::{Digest, Sha256};
use std::{
    fs::File,
    io::{Cursor, Read},
    path::Path,
};

const MAX_TEXT: usize = 2 * 1024 * 1024;
const MAX_DOCUMENT: usize = 32 * 1024 * 1024;
const MAX_XML: usize = 8 * 1024 * 1024;
const MAX_PDF_STREAM: usize = 4 * 1024 * 1024;
const MAX_PASSAGES: usize = 10_000;

pub(crate) struct Passage {
    pub text: String,
    pub source: SourceLocation,
}

pub(crate) struct SourceInput {
    bytes: Vec<u8>,
    pub hash: String,
}

/// Read once: the fingerprint and parser consume precisely the same bytes.
pub(crate) fn read_source(path: &Path) -> Result<Option<SourceInput>> {
    let extension = path
        .extension()
        .and_then(|value| value.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();
    let document = matches!(extension.as_str(), "pdf" | "docx");
    if !document
        && !crate::is_code_path(path)
        && !matches!(
            extension.as_str(),
            "txt"
                | "md"
                | "json"
                | "toml"
                | "yaml"
                | "yml"
                | "csv"
                | "html"
                | "css"
                | "xml"
                | "log"
                | "sh"
        )
    {
        return Ok(None);
    }
    if !document && File::open(path)?.metadata()?.len() > MAX_TEXT as u64 {
        return Ok(None);
    }
    let bytes = read_bounded(path, if document { MAX_DOCUMENT } else { MAX_TEXT })?;
    let hash = format!("{:x}", Sha256::digest(&bytes));
    Ok(Some(SourceInput { bytes, hash }))
}

pub(crate) fn extract(path: &Path, input: SourceInput) -> Result<Option<Vec<Passage>>> {
    let extension = path
        .extension()
        .and_then(|value| value.to_str())
        .unwrap_or("")
        .to_ascii_lowercase();
    if extension == "pdf" {
        return pdf(&input.bytes).map(Some);
    }
    if extension == "docx" {
        return docx(&input.bytes).map(Some);
    }
    let code = crate::is_code_path(path);
    if input.bytes.contains(&0) {
        return Ok(None);
    }
    let Ok(text) = String::from_utf8(input.bytes) else {
        return Ok(None);
    };
    let lines = text.split_inclusive('\n').collect::<Vec<_>>();
    if code {
        let mut passages = Vec::new();
        for start in (0..lines.len()).step_by(96) {
            if passages.len() >= MAX_PASSAGES {
                return Ok(None);
            }
            let end = (start + 128).min(lines.len());
            passages.push(Passage {
                text: lines[start..end].concat(),
                source: SourceLocation {
                    kind: "line".into(),
                    number: (start + 1) as u32,
                },
            });
            if end == lines.len() {
                break;
            }
        }
        return Ok(Some(passages));
    }
    Ok(Some(
        lines
            .chunks(64)
            .enumerate()
            .map(|(index, lines)| Passage {
                text: lines.concat(),
                source: SourceLocation {
                    kind: "line".into(),
                    number: (index * 64 + 1) as u32,
                },
            })
            .collect(),
    ))
}

fn read_bounded(path: &Path, maximum: usize) -> Result<Vec<u8>> {
    let file = File::open(path)?;
    if file.metadata()?.len() > maximum as u64 {
        bail!("document input exceeds {maximum} bytes");
    }
    let mut bytes = Vec::new();
    file.take(maximum as u64 + 1).read_to_end(&mut bytes)?;
    if bytes.len() > maximum {
        bail!("document input exceeds {maximum} bytes");
    }
    Ok(bytes)
}

fn pdf(bytes: &[u8]) -> Result<Vec<Passage>> {
    let document = lopdf::Document::load_mem_with_options(
        bytes,
        lopdf::LoadOptions {
            strict: true,
            max_decompressed_size: Some(MAX_PDF_STREAM),
            ..Default::default()
        },
    )
    .context("invalid, encrypted, or over-limit PDF")?;
    if document.was_encrypted() || document.is_encrypted() {
        bail!("encrypted PDF is unsupported");
    }
    let pages = document.get_pages();
    if pages.is_empty() {
        bail!("PDF has no pages");
    }
    if pages.len() > 1000 {
        bail!("PDF exceeds 1000 pages");
    }
    let mut passages = Vec::new();
    let mut total = 0;
    for number in pages.keys() {
        let text = document
            .extract_text_with_limit(&[*number], MAX_PDF_STREAM)
            .with_context(|| format!("PDF text extraction failed on page {number}"))?;
        total += text.len();
        if total > MAX_TEXT {
            bail!("PDF text exceeds {MAX_TEXT} bytes");
        }
        if !text.trim().is_empty() {
            passages.push(Passage {
                text,
                source: SourceLocation {
                    kind: "page".into(),
                    number: *number,
                },
            });
        }
    }
    Ok(passages)
}

fn docx(bytes: &[u8]) -> Result<Vec<Passage>> {
    let mut archive =
        zip::ZipArchive::new(Cursor::new(bytes)).context("invalid DOCX ZIP container")?;
    if archive.len() > 10_000 {
        bail!("DOCX exceeds 10000 archive entries");
    }
    let mut xml = String::new();
    let part = archive
        .by_name("word/document.xml")
        .context("DOCX lacks word/document.xml")?;
    if part.size() > MAX_XML as u64 {
        bail!("DOCX document XML exceeds {MAX_XML} bytes");
    }
    part.take(MAX_XML as u64 + 1)
        .read_to_string(&mut xml)
        .context("DOCX XML must be valid UTF-8")?;
    if xml.len() > MAX_XML {
        bail!("DOCX document XML exceeds {MAX_XML} bytes");
    }
    let mut reader = NsReader::from_str(&xml);
    let mut passages = Vec::new();
    let mut paragraph = String::new();
    let mut number = 0;
    let mut in_paragraph = false;
    let mut in_text = false;
    let mut total = 0;
    let mut depth = 0usize;
    let mut saw_document = false;
    loop {
        let (namespace, event) = reader.read_resolved_event().context("invalid DOCX XML")?;
        let word_namespace = matches!(namespace, ResolveResult::Bound(value) if matches!(value.as_ref(),
            "http://schemas.openxmlformats.org/wordprocessingml/2006/main" | "http://purl.oclc.org/ooxml/wordprocessingml/main"));
        match &event {
            Event::Start(start) => {
                if depth == 0 {
                    if saw_document || !word_namespace || start.local_name().as_ref() != "document"
                    {
                        bail!("invalid DOCX document root");
                    }
                    saw_document = true;
                }
                depth += 1;
                if depth > 128 {
                    bail!("DOCX XML exceeds 128 nesting levels");
                }
            }
            Event::End(_) => depth = depth.checked_sub(1).context("invalid DOCX XML depth")?,
            Event::Empty(_) if depth == 0 => bail!("invalid DOCX document root"),
            _ => {}
        }
        match event {
            Event::Start(event) if word_namespace => match event.local_name().as_ref() {
                "p" => {
                    if in_paragraph {
                        bail!("nested DOCX paragraphs are unsupported");
                    }
                    number += 1;
                    if number > MAX_PASSAGES {
                        bail!("DOCX exceeds 10000 paragraphs");
                    }
                    in_paragraph = true;
                    paragraph.clear();
                }
                "t" => in_text = in_paragraph,
                _ => {}
            },
            Event::End(event) if word_namespace => match event.local_name().as_ref() {
                "t" => in_text = false,
                "p" => {
                    in_paragraph = false;
                    in_text = false;
                    total += paragraph.len();
                    if total > MAX_TEXT {
                        bail!("DOCX text exceeds {MAX_TEXT} bytes");
                    }
                    if !paragraph.trim().is_empty() {
                        passages.push(Passage {
                            text: std::mem::take(&mut paragraph),
                            source: SourceLocation {
                                kind: "paragraph".into(),
                                number: number as u32,
                            },
                        });
                    }
                }
                _ => {}
            },
            Event::Empty(event) if word_namespace => match event.local_name().as_ref() {
                "tab" if in_paragraph => paragraph.push('\t'),
                "br" | "cr" if in_paragraph => paragraph.push('\n'),
                "p" => {
                    number += 1;
                    if number > MAX_PASSAGES {
                        bail!("DOCX exceeds 10000 paragraphs");
                    }
                }
                _ => {}
            },
            Event::Text(event) if in_text => {
                paragraph.push_str(&event.xml_content(XmlVersion::Implicit1_0))
            }
            Event::CData(event) if in_text => {
                paragraph.push_str(&event.xml_content(XmlVersion::Implicit1_0))
            }
            Event::GeneralRef(event) if in_text => {
                let reference = format!("&{};", event.as_ref());
                paragraph.push_str(
                    &quick_xml::escape::unescape(&reference).context("unsupported DOCX entity")?,
                );
            }
            Event::DocType(_) => bail!("DOCX XML document types are unsupported"),
            Event::Eof => {
                if in_paragraph || depth != 0 || !saw_document {
                    bail!("incomplete DOCX document");
                }
                break;
            }
            _ => {}
        }
        if total + paragraph.len() > MAX_TEXT {
            bail!("DOCX text exceeds {MAX_TEXT} bytes");
        }
    }
    Ok(passages)
}

#[cfg(test)]
mod provenance_tests {
    use super::*;

    #[test]
    fn parsing_and_fingerprint_share_one_snapshot_even_when_path_changes() -> Result<()> {
        let directory = tempfile::tempdir()?;
        let path = directory.path().join("notes.txt");
        std::fs::write(&path, "original parser input")?;
        let input = read_source(&path)?.unwrap();
        assert_eq!(
            input.hash,
            format!("{:x}", Sha256::digest(b"original parser input"))
        );
        std::fs::write(&path, "replacement on disk")?;
        let passages = extract(&path, input)?.unwrap();
        assert_eq!(passages[0].text, "original parser input");
        Ok(())
    }
}