typst

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, LazyLock, Mutex};

use rustler::{Binary, Env, NewBinary, NifStruct, Term};
use typst::diag::{
    eco_format, FileError, FileResult, PackageError, PackageResult, SourceDiagnostic,
};
use typst::ecow::EcoVec;
use typst::foundations::{Bytes, Datetime, Duration};
use typst::syntax::package::PackageSpec;
use typst::syntax::{FileId, RootedPath, Source, VirtualPath, VirtualRoot};
use typst::text::{Font, FontBook};
use typst::utils::LazyHash;
use typst::WorldExt;
use typst::{Library, LibraryExt};
use typst_kit::fonts::FontStore;
use typst_layout::PagedDocument;
use typst_pdf::{PdfOptions, PdfStandard, PdfStandards};

struct CachedFonts {
    key: Vec<String>,
    store: Arc<FontStore>,
}

static FONT_CACHE: LazyLock<Mutex<Option<CachedFonts>>> = LazyLock::new(|| Mutex::new(None));

static LIBRARY_CACHE: LazyLock<LazyHash<Library>> =
    LazyLock::new(|| LazyHash::new(Library::default()));

static HTTP_AGENT: LazyLock<ureq::Agent> = LazyLock::new(|| {
    ureq::AgentBuilder::new()
        .timeout(std::time::Duration::from_secs(30))
        .build()
});

static CACHE_DIRECTORY: LazyLock<PathBuf> = LazyLock::new(|| {
    std::env::var_os("CACHE_DIRECTORY")
        .map(|os_path| os_path.into())
        .unwrap_or(std::env::temp_dir())
});

/// Main interface that determines the environment for Typst.
pub struct TypstNifWorld {
    /// Root path to which files will be resolved.
    root: PathBuf,

    /// The content of a source.
    source: Source,

    /// Metadata about all known fonts.
    store: Arc<FontStore>,

    /// Map of all known files.
    files: Arc<Mutex<HashMap<FileId, FileEntry>>>,

    /// Cache directory (e.g. where packages are downloaded to).
    cache_directory: PathBuf,

    /// Datetime.
    time: time::OffsetDateTime,
}

impl TypstNifWorld {
    pub fn new(root: String, source: String, extra_fonts: Vec<String>, cache_fonts: bool) -> Self {
        let root = PathBuf::from(root);

        let store = if cache_fonts {
            let mut sorted_key = extra_fonts.clone();
            sorted_key.sort();
            let mut cache = FONT_CACHE
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());

            let needs_scan = match cache.as_ref() {
                Some(cached) if cached.key == sorted_key => false,
                _ => true,
            };

            if needs_scan {
                *cache = Some(CachedFonts {
                    key: sorted_key,
                    store: Arc::new(build_font_store(&extra_fonts)),
                });
            }

            Arc::clone(&cache.as_ref().unwrap().store)
        } else {
            Arc::new(build_font_store(&extra_fonts))
        };

        Self {
            root,
            store,
            source: Source::detached(source),
            time: time::OffsetDateTime::now_utc(),
            cache_directory: CACHE_DIRECTORY.clone(),
            files: Arc::new(Mutex::new(HashMap::new())),
        }
    }
}

/// A File that will be stored in the HashMap.
#[derive(Clone, Debug)]
struct FileEntry {
    bytes: Bytes,
    source: Option<Source>,
}

impl FileEntry {
    fn new(bytes: Vec<u8>, source: Option<Source>) -> Self {
        Self {
            bytes: Bytes::new(bytes),
            source,
        }
    }

    fn source(&mut self, id: FileId) -> FileResult<Source> {
        let source = if let Some(source) = &self.source {
            source
        } else {
            let contents = std::str::from_utf8(&self.bytes).map_err(|_| FileError::InvalidUtf8)?;
            let contents = contents.trim_start_matches('\u{feff}');
            let source = Source::new(id, contents.into());
            self.source.insert(source)
        };
        Ok(source.clone())
    }
}

impl TypstNifWorld {
    /// Helper to handle file requests.
    ///
    /// Requests will be either in packages or a local file.
    fn file(&self, id: FileId) -> FileResult<FileEntry> {
        let mut files = self.files.lock().map_err(|_| FileError::AccessDenied)?;
        if let Some(entry) = files.get(&id) {
            return Ok(entry.clone());
        }
        let path = if let VirtualRoot::Package(package) = id.root() {
            // Fetching file from package
            let package_dir = self.download_package(package)?;
            id.vpath().realize(&package_dir)
        } else {
            // Fetching file from disk
            id.vpath().realize(&self.root)
        }
        .map_err(FileError::Realize)?;

        let content = std::fs::read(&path).map_err(|error| FileError::from_io(error, &path))?;
        Ok(files
            .entry(id)
            .or_insert(FileEntry::new(content, None))
            .clone())
    }

    /// Downloads the package and returns the system path of the unpacked package.
    fn download_package(&self, package: &PackageSpec) -> PackageResult<PathBuf> {
        let package_subdir = format!("{}/{}/{}", package.namespace, package.name, package.version);
        let path = self.cache_directory.join(package_subdir);

        if path.exists() {
            return Ok(path);
        }

        eprintln!("downloading {package}");
        let url = format!(
            "https://packages.typst.org/{}/{}-{}.tar.gz",
            package.namespace, package.name, package.version,
        );

        let response = retry(|| {
            let response = HTTP_AGENT
                .get(&url)
                .call()
                .map_err(|error| eco_format!("{error}"))?;

            let status = response.status();
            if !http_successful(status) {
                return Err(eco_format!(
                    "response returned unsuccessful status code {status}",
                ));
            }

            Ok(response)
        })
        .map_err(|error| PackageError::NetworkFailed(Some(error)))?;

        let mut compressed_archive = Vec::new();
        response
            .into_reader()
            .read_to_end(&mut compressed_archive)
            .map_err(|error| PackageError::NetworkFailed(Some(eco_format!("{error}"))))?;
        let raw_archive = zune_inflate::DeflateDecoder::new(&compressed_archive)
            .decode_gzip()
            .map_err(|error| PackageError::MalformedArchive(Some(eco_format!("{error}"))))?;
        let mut archive = tar::Archive::new(raw_archive.as_slice());
        archive.unpack(&path).map_err(|error| {
            _ = std::fs::remove_dir_all(&path);
            PackageError::MalformedArchive(Some(eco_format!("{error}")))
        })?;

        Ok(path)
    }

    pub fn insert_virtual_file<S: Into<String>>(
        &self,
        vpath: S,
        bytes: Vec<u8>,
    ) -> FileResult<FileId> {
        let vp = VirtualPath::new(vpath.into()).map_err(|_| FileError::Other(None))?;
        let id = FileId::new(RootedPath::new(VirtualRoot::Project, vp));
        let mut files = self.files.lock().map_err(|_| FileError::AccessDenied)?;
        files.insert(id, FileEntry::new(bytes, None)); // Source created lazily if ever needed
        Ok(id)
    }
}

/// This is the interface we have to implement such that `typst` can compile it.
///
/// I have tried to keep it as minimal as possible
impl typst::World for TypstNifWorld {
    /// Standard library.
    fn library(&self) -> &LazyHash<Library> {
        &LIBRARY_CACHE
    }

    /// Metadata about all known Books.
    fn book(&self) -> &LazyHash<FontBook> {
        self.store.book()
    }

    /// Accessing the main source file.
    fn main(&self) -> FileId {
        self.source.id()
    }

    /// Accessing a specified source file (based on `FileId`).
    fn source(&self, id: FileId) -> FileResult<Source> {
        if id == self.source.id() {
            Ok(self.source.clone())
        } else {
            self.file(id)?.source(id)
        }
    }

    /// Accessing a specified file (non-file).
    fn file(&self, id: FileId) -> FileResult<Bytes> {
        self.file(id).map(|file| file.bytes.clone())
    }

    /// Accessing a specified font per index of font book.
    fn font(&self, id: usize) -> Option<Font> {
        self.store.font(id)
    }

    /// Get the current date.
    ///
    /// Optionally, an offset is given.
    fn today(&self, offset: Option<Duration>) -> Option<Datetime> {
        let offset = offset.map(time::Duration::from).unwrap_or_default();
        let shifted = self.time.checked_add(offset)?;
        Some(Datetime::Date(shifted.date()))
    }
}

fn build_font_store(extra_fonts: &[String]) -> FontStore {
    let mut store = FontStore::new();
    store.extend(typst_kit::fonts::embedded());
    store.extend(typst_kit::fonts::system());
    for dir in extra_fonts {
        store.extend(typst_kit::fonts::scan(Path::new(dir)));
    }
    store
}

fn retry<T, E>(mut f: impl FnMut() -> Result<T, E>) -> Result<T, E> {
    if let Ok(ok) = f() {
        Ok(ok)
    } else {
        f()
    }
}

fn http_successful(status: u16) -> bool {
    // 2XX
    status / 100 == 2
}

#[derive(NifStruct)]
#[module = "Typst.NIF.PdfOptions"]
struct PdfOpts {
    standards: Vec<String>,
}

fn parse_pdf_standard(s: &str) -> Result<PdfStandard, String> {
    match s {
        "a-1a" | "1a" => Ok(PdfStandard::A_1a),
        "a-1b" | "1b" => Ok(PdfStandard::A_1b),
        "a-2a" | "2a" => Ok(PdfStandard::A_2a),
        "a-2b" | "2b" => Ok(PdfStandard::A_2b),
        "a-2u" | "2u" => Ok(PdfStandard::A_2u),
        "a-3a" | "3a" => Ok(PdfStandard::A_3a),
        "a-3b" | "3b" => Ok(PdfStandard::A_3b),
        "a-3u" | "3u" => Ok(PdfStandard::A_3u),
        "a-4" | "4" => Ok(PdfStandard::A_4),
        "a-4e" | "4e" => Ok(PdfStandard::A_4e),
        "a-4f" | "4f" => Ok(PdfStandard::A_4f),
        other => Err(format!("unknown PDF standard: {}", other)),
    }
}

#[rustler::nif(schedule = "DirtyCpu")]
fn compile_pdf<'a>(
    env: Env<'a>,
    markup: String,
    root_dir: String,
    extra_fonts: Vec<String>,
    assets: Vec<(String, Binary<'a>)>,
    cache_fonts: bool,
    pdf_opts: PdfOpts,
) -> Result<Term<'a>, String> {
    let world = TypstNifWorld::new(root_dir, markup, extra_fonts, cache_fonts);

    for (vpath, bin) in assets {
        world
            .insert_virtual_file(vpath, bin.as_slice().to_vec())
            .map_err(|e| format!("{:#?}", e))?;
    }

    let document: PagedDocument = typst::compile(&world)
        .output
        .map_err(|e| collect_typst_errors(e, &world))?;

    comemo::evict(0);

    let options = if pdf_opts.standards.is_empty() {
        PdfOptions::default()
    } else {
        let standards: Vec<PdfStandard> = pdf_opts
            .standards
            .iter()
            .map(|s| parse_pdf_standard(s))
            .collect::<Result<Vec<_>, _>>()?;
        PdfOptions {
            standards: PdfStandards::new(&standards).map_err(|e| e.message().to_string())?,
            ..PdfOptions::default()
        }
    };

    let pdf_bytes =
        typst_pdf::pdf(&document, &options).map_err(|e| collect_typst_errors(e, &world))?;

    let mut binary = NewBinary::new(env, pdf_bytes.len());
    binary.copy_from_slice(pdf_bytes.as_slice());

    Ok(binary.into())
}

#[rustler::nif(schedule = "DirtyCpu")]
fn compile_png<'a>(
    env: Env<'a>,
    markup: String,
    root_dir: String,
    extra_fonts: Vec<String>,
    pixels_per_pt: f64,
    assets: Vec<(String, Binary<'a>)>,
    cache_fonts: bool,
) -> Result<Vec<Binary<'a>>, String> {
    let world = TypstNifWorld::new(root_dir, markup, extra_fonts, cache_fonts);

    for (vpath, bin) in assets {
        world
            .insert_virtual_file(vpath, bin.as_slice().to_vec())
            .map_err(|e| format!("{:#?}", e))?;
    }

    let document: PagedDocument = typst::compile(&world)
        .output
        .map_err(|e| collect_typst_errors(e, &world))?;

    comemo::evict(0);

    let options = typst_render::RenderOptions {
        pixel_per_pt: typst::utils::Scalar::new(pixels_per_pt),
        render_bleed: false,
    };
    let pngs: Result<Vec<Binary>, String> = document
        .pages()
        .iter()
        .map(|page| {
            let pixmap = typst_render::render(page, &options);
            let png = pixmap.encode_png().map_err(|e| format!("{:#?}", e))?;

            let mut binary = NewBinary::new(env, png.len());
            binary.copy_from_slice(&png);
            Ok(binary.into())
        })
        .collect();

    pngs
}

#[rustler::nif(schedule = "DirtyCpu")]
fn compile_svg<'a>(
    env: Env<'a>,
    markup: String,
    root_dir: String,
    extra_fonts: Vec<String>,
    assets: Vec<(String, Binary<'a>)>,
    cache_fonts: bool,
) -> Result<Vec<Binary<'a>>, String> {
    let world = TypstNifWorld::new(root_dir, markup, extra_fonts, cache_fonts);

    for (vpath, bin) in assets {
        world
            .insert_virtual_file(vpath, bin.as_slice().to_vec())
            .map_err(|e| format!("{:#?}", e))?;
    }

    let document: PagedDocument = typst::compile(&world)
        .output
        .map_err(|e| collect_typst_errors(e, &world))?;

    comemo::evict(0);

    let options = typst_svg::SvgOptions::default();
    let svgs: Vec<Binary> = document
        .pages()
        .iter()
        .map(|page| {
            let svg_string = typst_svg::svg(page, &options);
            let svg_bytes = svg_string.as_bytes();

            let mut binary = NewBinary::new(env, svg_bytes.len());
            binary.copy_from_slice(svg_bytes);
            binary.into()
        })
        .collect();

    Ok(svgs)
}

fn collect_typst_errors(errors: EcoVec<SourceDiagnostic>, world: &TypstNifWorld) -> String {
    let mut error_messages = Vec::new();
    let source = &world.source;

    for error in errors {
        let span = error.span;

        let mut error_msg = format!("Error: {}", error.message);

        if !span.is_detached() && span.id() == Some(source.id()) {
            if let Some(range) = world.range(span) {
                let lines = source.lines();
                let line = lines.byte_to_line(range.start).unwrap_or(0) + 1;
                let column = lines.byte_to_column(range.start).unwrap_or(0) + 1;

                error_msg = format!("[line {}:{}] {}", line, column, error.message);

                // Try to get the actual source line for context
                if let Some(line_range) = lines.line_to_range(line - 1) {
                    let source_line = &source.text()[line_range];
                    let trimmed_line = source_line.trim_end();

                    // Calculate the position of the error marker
                    let leading_spaces = source_line.len() - source_line.trim_start().len();
                    let marker_pos = column.saturating_sub(leading_spaces);

                    error_msg = format!(
                        "{}\n  Source: {}\n         {}{}",
                        error_msg,
                        trimmed_line.trim(),
                        " ".repeat(marker_pos),
                        "^"
                    );
                }
            }
        }

        // Add hints if any
        for hint in &error.hints {
            error_msg = format!("{}\n  Hint: {}", error_msg, hint.v);
        }

        error_messages.push(error_msg);
    }

    error_messages.join("\n\n")
}

rustler::init!("Elixir.Typst.NIF");