use std::io::{self, Read};
use std::path::PathBuf;
use std::sync::OnceLock;

use chrono::Datelike;
use clap::Parser;
use typst::diag::{FileError, FileResult};
use typst::foundations::{Bytes, Datetime};
use typst::layout::PagedDocument;
use typst::syntax::{FileId, Source, VirtualPath};
use typst::text::{Font, FontBook};
use typst::utils::LazyHash;
use typst::{Library, World};

#[derive(Parser, Debug)]
#[command(name = "mathrender", about = "Render LaTeX or Typst math to SVG")]
struct Args {
    input: Option<String>,

    #[arg(long, short = 't')]
    typst: bool,

    #[arg(long, short = 'i')]
    inline: bool,

    #[arg(long, default_value_t = 16.0)]
    font_size: f64,

    #[arg(long, default_value = "#000000")]
    color: String,
}

#[derive(Debug)]
enum Error {
    Compile(String),
    NoPages,
}

impl std::fmt::Display for Error {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Error::Compile(msg) => write!(f, "Typst compile error(s):\n  {}", msg),
            Error::NoPages => write!(f, "Typst produced no pages"),
        }
    }
}

impl std::error::Error for Error {}

struct GlobalFonts {
    book: LazyHash<FontBook>,
    fonts: Vec<Font>,
}

fn get_global_fonts() -> &'static GlobalFonts {
    static FONTS: OnceLock<GlobalFonts> = OnceLock::new();
    FONTS.get_or_init(|| {
        let mut book = FontBook::new();
        let mut fonts = Vec::new();

        for font_data in typst_assets::fonts() {
            let bytes = Bytes::new(font_data);
            for index in 0u32.. {
                match Font::new(bytes.clone(), index) {
                    Some(font) => {
                        book.push(font.info().clone());
                        fonts.push(font);
                    }
                    None => break,
                }
            }
        }
        GlobalFonts { book: LazyHash::new(book), fonts }
    })
}

fn safe_tex2typst(s: &str, buf: &mut String) {
    let s = s.trim();
    if s.is_empty() {
        return;
    }
    // Assertion enforcing pre-condition bounds
    debug_assert!(s.len() < 10000, "Input to safe_tex2typst is abnormally large");
    
    buf.push_str(&tex2typst_rs::tex2typst(s));
}

fn env_matrix_convert(env: &str, content: &str, out_typst_string: &mut String) {
    let delim = match env {
        "pmatrix" | "pmatrix*"            => "(",
        "bmatrix" | "bmatrix*"            => "[",
        "vmatrix" | "vmatrix*"            => "|",
        "Bmatrix" | "Bmatrix*"            => "{",
        "Vmatrix" | "Vmatrix*"            => "‖",
        "matrix"  | "matrix*"             => "",
        "smallmatrix"                     => "(",
        _                                 => "(",
    };

    match delim {
        ""  => out_typst_string.push_str("mat(delim: #none, "),
        "(" => out_typst_string.push_str("mat("),
        d   => { out_typst_string.push_str("mat(delim: \""); out_typst_string.push_str(d); out_typst_string.push_str("\", "); }
    }

    let mut first_row = true;
    for row in content.split("\\\\").take(256) {
        if row.trim().is_empty() { continue; }
        if !first_row { out_typst_string.push_str("; "); }
        first_row = false;

        let mut first_cell = true;
        for cell in row.split('&').take(256) {
            let cell = cell.trim();
            if !first_cell { out_typst_string.push_str(", "); }
            first_cell = false;
            safe_tex2typst(cell, out_typst_string);
        }
    }
    out_typst_string.push(')');
}

fn env_cases_convert(content: &str, out_typst_string: &mut String) {
    out_typst_string.push_str("cases(");
    let mut first_row = true;
    for row in content.split("\\\\").take(256) {
        let row = row.trim();
        if row.is_empty() { continue; }
        if !first_row { out_typst_string.push_str(", "); }
        first_row = false;

        let mut parts = row.splitn(2, '&');
        let expr = parts.next().expect("Cases row must have content").trim();
        let cond = parts.next().map(|s| s.trim());
        
        safe_tex2typst(expr, out_typst_string);
        if let Some(c) = cond {
            out_typst_string.push_str(", ");
            safe_tex2typst(c, out_typst_string);
        }
    }
    out_typst_string.push(')');
}

fn process_latex_env(env: &str, content: &str, result: &mut String) {
    debug_assert!(!env.is_empty(), "Environment string cannot be empty");
    
    match env {
        "pmatrix"  | "pmatrix*"  |
        "bmatrix"  | "bmatrix*"  |
        "vmatrix"  | "vmatrix*"  |
        "Bmatrix"  | "Bmatrix*"  |
        "Vmatrix"  | "Vmatrix*"  |
        "matrix"   | "matrix*"   |
        "smallmatrix"            => env_matrix_convert(env, content, result),

        "cases" | "dcases"       => env_cases_convert(content, result),

        "align"    | "align*"    |
        "aligned"  | "alignat"   | "alignat*" |
        "gather"   | "gather*"   |
        "multline" | "multline*" |
        "split"    | "equation"  | "equation*" => safe_tex2typst(content, result),

        _ => {
            result.push_str("\\begin{");
            result.push_str(env);
            result.push('}');
            result.push_str(content);
            result.push_str("\\end{");
            result.push_str(env);
            result.push('}');
        }
    }
}

fn find_env_end<'a>(rest: &'a str, env: &str) -> Option<(&'a str, &'a str)> {
    let end_search = "\\end{";
    let mut search_rest = rest;
    let mut current_offset = 0;
    
    while let Some(end_begin) = search_rest.find(end_search) {
        let tag_content_start = end_begin + end_search.len();
        if search_rest[tag_content_start..].starts_with(env) && 
           search_rest[tag_content_start + env.len()..].starts_with('}') {
            let content = &rest[..current_offset + end_begin];
            let advance = current_offset + end_begin + end_search.len() + env.len() + 1;
            return Some((content, &rest[advance..]));
        }
        let advance = end_begin + end_search.len();
        search_rest = &search_rest[advance..];
        current_offset += advance;
    }
    None
}

fn latex_to_typst(input: &str) -> String {
    let mut result = String::with_capacity(input.len() * 2);
    let mut rest   = input;

    while let Some(begin_pos) = rest.find("\\begin{") {
        safe_tex2typst(&rest[..begin_pos], &mut result);
        rest = &rest[begin_pos + 7..];

        let Some(brace_end) = rest.find('}') else {
            result.push_str("\\begin{");
            continue;
        };
        let env  = &rest[..brace_end];
        rest     = &rest[brace_end + 1..];

        if let Some((content, new_rest)) = find_env_end(rest, env) {
            process_latex_env(env, content, &mut result);
            rest = new_rest;
        } else {
            result.push_str("\\begin{");
            result.push_str(env);
            result.push('}');
        }
    }

    safe_tex2typst(rest, &mut result);
    result
}

fn get_global_library() -> &'static LazyHash<Library> {
    static LIBRARY: OnceLock<LazyHash<Library>> = OnceLock::new();
    LIBRARY.get_or_init(|| {
        LazyHash::new(typst_library::LibraryBuilder::default().build())
    })
}

struct MathWorld {
    source: Source,
}

impl MathWorld {
    fn new(source_text: String) -> Self {
        let main_id = FileId::new(None, VirtualPath::new("/main.typ"));
        let source  = Source::new(main_id, source_text);

        Self {
            source,
        }
    }
}

impl World for MathWorld {
    fn library(&self) -> &LazyHash<Library> { get_global_library() }
    fn book(&self)    -> &LazyHash<FontBook> { &get_global_fonts().book }
    fn main(&self)    -> FileId              { self.source.id() }

    fn source(&self, id: FileId) -> FileResult<Source> {
        if id == self.source.id() {
            Ok(self.source.clone())
        } else {
            Err(FileError::NotFound(PathBuf::from("<virtual>")))
        }
    }

    fn file(&self, _id: FileId) -> FileResult<Bytes> {
        Err(FileError::NotFound(PathBuf::from("<virtual>")))
    }

    fn font(&self, index: usize) -> Option<Font> {
        get_global_fonts().fonts.get(index).cloned()
    }

    fn today(&self, offset: Option<i64>) -> Option<Datetime> {
        let now = chrono::Local::now();
        let _ = offset;
        Datetime::from_ymd(now.year(), now.month() as u8, now.day() as u8)
    }
}

fn build_typst_source(math_expr: &str, display: bool, font_size: f64, color: &str) -> String {
    let wrapped = if display {
        format!("$ {} $", math_expr)
    } else {
        format!("${}$", math_expr)
    };

    // 0.5em margin prevents clipping of subscripts and matrices.
    format!(
        "#set page(width: auto, height: auto, margin: 0.5em, fill: none)\n\
         #set text(size: {font_size}pt, fill: rgb(\"{color}\"))\n\
         {wrapped}"
    )
}

fn render_to_svg(
    input: &str,
    is_typst: bool,
    display: bool,
    font_size: f64,
    color: &str,
) -> Result<String, Error> {

    let typst_math = if is_typst {
        input.to_owned()
    } else {
        latex_to_typst(input)
    };

    let source   = build_typst_source(&typst_math, display, font_size, color);
    let world    = MathWorld::new(source);
    let compiled = typst::compile::<PagedDocument>(&world);

    let document: PagedDocument = compiled.output.map_err(|errors| {
        let msgs: Vec<String> = errors.iter().map(|e| e.message.to_string()).collect();
        Error::Compile(msgs.join("\n  "))
    })?;

    let page = document.pages.first().ok_or(Error::NoPages)?;
    Ok(typst_svg::svg(page))
}

fn main() {
    let args = Args::parse();

    let input = match args.input {
        Some(s) => s,
        None => {
            let mut buf = String::new();
            io::stdin().read_to_string(&mut buf).expect("Failed to read stdin");
            buf.trim().to_owned()
        }
    };

    if input.is_empty() {
        eprintln!("Error: no input provided.");
        std::process::exit(1);
    }

    match render_to_svg(&input, args.typst, !args.inline, args.font_size, &args.color) {
        Ok(svg) => print!("{}", svg),
        Err(e)  => {
            eprintln!("Error: {e}");
            std::process::exit(1);
        }
    }
}

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

    #[test]
    fn test_safe_tex2typst() {
        let mut buf = String::new();
        safe_tex2typst("x = 1", &mut buf);
        assert!(buf.contains("x = 1"));
    }

    #[test]
    fn test_convert_matrix() {
        let mut buf = String::new();
        env_matrix_convert("pmatrix", "1 & 2 \\\\ 3 & 4", &mut buf);
        assert_eq!(buf, "mat(1, 2; 3, 4)");
    }

    #[test]
    fn test_convert_cases() {
        let mut buf = String::new();
        env_cases_convert("x & y \\\\ a & b", &mut buf);
        assert!(buf.contains("cases("));
        assert!(buf.contains("x"));
        assert!(buf.contains("y"));
    }

    #[test]
    fn test_latex_to_typst() {
        let input = "x + \\begin{pmatrix} 1 & 2 \\\\ 3 & 4 \\end{pmatrix} = y";
        let result = latex_to_typst(input);
        assert!(result.contains("mat(1, 2; 3, 4)"));
    }
}