use std::fs;
use std::io::Read as _;
use std::path::{Path, PathBuf};
use std::sync::{Arc, OnceLock};

use anyhow::{Context, Result};
use arborium_highlight::{HtmlFormat, Span, html_escape, spans_to_flat_tokens, spans_to_html};
use sha2::{Digest, Sha256};

use crate::grammar_manifest;

static CACHE_DIR: OnceLock<PathBuf> = OnceLock::new();

pub struct Cache {
    options: arbrisseau::Options,
}

impl Cache {
    pub fn new(root: impl AsRef<Path>) -> Self {
        CACHE_DIR
            .set(root.as_ref().join(grammar_manifest::VERSION))
            .expect("grammar cache already initialized");
        Self {
            options: arbrisseau::Options::new(load_module),
        }
    }
}

fn load_module(language: &str) -> arbrisseau::ModuleFuture<'_> {
    Box::pin(async move {
        let Some((size, hash)) = grammar_manifest::artifact(language) else {
            return Ok(None);
        };
        let path = CACHE_DIR
            .get()
            .expect("grammar cache not initialized")
            .join(format!("{language}.wasm"));
        if let Ok(bytes) = fs::read(&path)
            && verified(&bytes, size, hash)
        {
            return Ok(Some(bytes));
        }
        match download(language, size, hash, &path) {
            Ok(bytes) => Ok(Some(bytes)),
            Err(error) => {
                eprintln!("sorceryd: cannot load {language} grammar: {error:#}");
                Ok(None)
            }
        }
    })
}

fn verified(bytes: &[u8], size: usize, hash: &str) -> bool {
    bytes.len() == size && format!("{:x}", Sha256::digest(bytes)) == hash
}

fn download(language: &str, size: usize, hash: &str, path: &Path) -> Result<Vec<u8>> {
    let url = format!(
        "https://cdn.jsdelivr.net/npm/@arborium/{language}@{}/grammar_bg.wasm",
        grammar_manifest::VERSION,
    );
    let mut response = minreq::get(&url)
        .with_timeout(30)
        .with_max_redirects(3)
        .send_lazy()
        .with_context(|| format!("downloading {url}"))?;
    if response.status_code != 200 {
        anyhow::bail!("downloading {url}: HTTP {}", response.status_code);
    }
    let mut bytes = Vec::with_capacity(size);
    std::io::Read::take(&mut response, size as u64 + 1)
        .read_to_end(&mut bytes)
        .with_context(|| format!("downloading {url}"))?;
    if !verified(&bytes, size, hash) {
        anyhow::bail!("downloading {url}: size or SHA-256 mismatch");
    }
    let dir = path.parent().unwrap();
    fs::create_dir_all(dir).with_context(|| format!("creating grammar cache {}", dir.display()))?;
    // Concurrent sorcery processes may race on the same cache; a per-process
    // temporary and rename keep the final file always complete.
    let temporary = path.with_extension(format!("wasm.tmp-{}", std::process::id()));
    fs::write(&temporary, &bytes)
        .and_then(|()| fs::rename(&temporary, path))
        .with_context(|| format!("installing grammar cache {}", path.display()))?;
    Ok(bytes)
}

pub struct Highlighter {
    cache: Arc<Cache>,
    highlighter: arbrisseau::Highlighter,
}

impl Highlighter {
    pub fn new(cache: Arc<Cache>) -> Self {
        Self {
            cache,
            highlighter: arbrisseau::Highlighter::new(),
        }
    }

    pub fn highlight(&mut self, language: &str, source: &str) -> Result<String> {
        let spans = pollster::block_on(self.highlighter.highlight(
            language,
            source,
            &self.cache.options,
        ))?;
        Ok(spans_to_html(source, spans, &HtmlFormat::default()))
    }

    pub fn highlight_lines(&mut self, language: &str, source: &str) -> Result<Vec<String>> {
        let spans = pollster::block_on(self.highlighter.highlight(
            language,
            source,
            &self.cache.options,
        ))?;
        Ok(render_lines(source, spans))
    }
}

fn render_lines(source: &str, spans: Vec<Span>) -> Vec<String> {
    if source.is_empty() {
        return Vec::new();
    }
    let tokens = spans_to_flat_tokens(source, spans);
    let mut token = 0;
    let mut offset = 0;
    let mut lines = source.split('\n').collect::<Vec<_>>();
    if source.ends_with('\n') {
        lines.pop();
    }

    lines
        .into_iter()
        .map(|raw_line| {
            let line = raw_line.strip_suffix('\r').unwrap_or(raw_line);
            let start = offset;
            let end = start + line.len();
            offset += raw_line.len() + 1;
            while tokens
                .get(token)
                .is_some_and(|span| span.end as usize <= start)
            {
                token += 1;
            }

            let mut html = String::new();
            let mut position = start;
            let mut i = token;
            while let Some(span) = tokens.get(i).filter(|span| (span.start as usize) < end) {
                let span_start = (span.start as usize).max(start);
                let span_end = (span.end as usize).min(end);
                if position < span_start {
                    html.push_str(&html_escape(&source[position..span_start]));
                }
                if span_start < span_end {
                    html.push_str(&format!(
                        "<a-{tag}>{text}</a-{tag}>",
                        tag = span.tag,
                        text = html_escape(&source[span_start..span_end]),
                    ));
                }
                position = span_end;
                i += 1;
            }
            if position < end {
                html.push_str(&html_escape(&source[position..end]));
            }
            html
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use arbrisseau::Span;

    use super::render_lines;

    #[test]
    fn renders_highlighted_lines_independently() {
        let source = "/* <\n * x */\r\nlet x = 1;\n";
        let spans = vec![
            Span {
                start: 0,
                end: 12,
                capture: "comment".into(),
                pattern_index: 0,
            },
            Span {
                start: 14,
                end: 17,
                capture: "keyword".into(),
                pattern_index: 0,
            },
        ];
        assert_eq!(
            render_lines(source, spans),
            [
                "<a-c>/* &lt;</a-c>",
                "<a-c> * x */</a-c>",
                "<a-k>let</a-k> x = 1;"
            ]
        );
    }
}