//! Pretty-printing compilation errors.
use core::fmt::{self, Display, Formatter};
use jaq_core::{compile, load};

/// File and corresponding error reports.
pub type FileReports<P = ()> = (load::File<String, P>, Vec<Report>);

/// Report errors that may occur when loading a module.
pub fn load_errors<P>(errs: load::Errors<&str, P>) -> Vec<FileReports<P>> {
    use load::Error;

    let errs = errs.into_iter().map(|(file, err)| {
        let code = file.code;
        let err = match err {
            Error::Io(errs) => errs.into_iter().map(|e| report_io(code, e)).collect(),
            Error::Lex(errs) => errs.into_iter().map(|e| report_lex(code, e)).collect(),
            Error::Parse(errs) => errs.into_iter().map(|e| report_parse(code, e)).collect(),
        };
        (file.map_code(|s| s.into()), err)
    });
    errs.collect()
}

/// Report errors that may occur when compiling a module.
pub fn compile_errors<P>(errs: compile::Errors<&str, P>) -> Vec<FileReports<P>> {
    let errs = errs.into_iter().map(|(file, errs)| {
        let code = file.code;
        let errs = errs.into_iter().map(|e| report_compile(code, e)).collect();
        (file.map_code(|s| s.into()), errs)
    });
    errs.collect()
}

type StringColors = Vec<(String, Option<Color>)>;

/// Error report.
#[derive(Debug)]
pub struct Report {
    /// error summary
    message: String,
    labels: Vec<(core::ops::Range<usize>, StringColors, Color)>,
}

/// Error color.
#[derive(Copy, Clone, Debug)]
pub enum Color {
    /// used for most errors
    Red = 31,
    /// used for unclosed delimiters
    Yellow = 33,
}

impl Color {
    /// Format a string with ANSI colors.
    pub fn ansi(self, f: &mut Formatter, text: &dyn Display) -> fmt::Result {
        write!(f, "\x1b[{}m{}", self as usize, text)?;
        write!(f, "\x1b[{}m", 0)
    }
}

fn report_io(code: &str, (path, error): (&str, String)) -> Report {
    let path_range = load::span(code, path);
    Report {
        message: format!("could not load file {path}: {error}"),
        labels: [(path_range, [(error, None)].into(), Color::Red)].into(),
    }
}

fn report_lex(code: &str, (expected, found): load::lex::Error<&str>) -> Report {
    // truncate found string to its first character
    let found = &found[..found.char_indices().nth(1).map_or(found.len(), |(i, _)| i)];

    let found_range = load::span(code, found);
    let found = match found {
        "" => [("unexpected end of input".to_string(), None)].into(),
        c => [("unexpected character ", None), (c, Some(Color::Red))]
            .map(|(s, c)| (s.into(), c))
            .into(),
    };
    let label = (found_range, found, Color::Red);

    let labels = match expected {
        load::lex::Expect::Delim(open) => {
            let text = [("unclosed delimiter ", None), (open, Some(Color::Yellow))]
                .map(|(s, c)| (s.into(), c));
            Vec::from([(load::span(code, open), text.into(), Color::Yellow), label])
        }
        _ => Vec::from([label]),
    };

    Report {
        message: format!("expected {}", expected.as_str()),
        labels,
    }
}

fn report_parse(code: &str, (expected, found): load::parse::Error<&str>) -> Report {
    let found_range = load::span(code, found);

    let found = if found.is_empty() {
        "unexpected end of input"
    } else {
        "unexpected token"
    };
    let found = [(found.to_string(), None)].into();

    Report {
        message: format!("expected {}", expected.as_str()),
        labels: Vec::from([(found_range, found, Color::Red)]),
    }
}

fn report_compile(code: &str, (found, undefined): compile::Error<&str>) -> Report {
    use compile::Undefined::Filter;
    let found_range = load::span(code, found);
    let wnoa = |exp, got| format!("wrong number of arguments (expected {exp}, found {got})");
    let message = match (found, undefined) {
        ("reduce", Filter(arity)) => wnoa("2", arity),
        ("foreach", Filter(arity)) => wnoa("2 or 3", arity),
        (_, undefined) => format!("undefined {}", undefined.as_str()),
    };
    let found = [(message.clone(), None)].into();

    Report {
        message,
        labels: Vec::from([(found_range, found, Color::Red)]),
    }
}

type CodeBlock = codesnake::Block<codesnake::CodeWidth<String>, Box<dyn Display>, Option<Color>>;

/// Function to apply color to snakes/text.
pub type Paint = fn(&mut Formatter, &Option<Color>, &dyn Display) -> fmt::Result;

struct FromFn<F>(F);

fn from_fn<F: Fn(&mut Formatter) -> fmt::Result>(f: F) -> FromFn<F> {
    FromFn(f)
}

impl<F: Fn(&mut Formatter) -> fmt::Result> Display for FromFn<F> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        (self.0)(f)
    }
}

impl Report {
    /// Convert report to a code block.
    pub fn to_block(&self, idx: &codesnake::LineIndex, paint: Paint) -> CodeBlock {
        use codesnake::{Block, CodeWidth, Label};
        let labels = self.labels.iter().cloned().map(|(range, text, color)| {
            Label::new(range)
                .with_style(Some(color))
                .with_text(Box::new(from_fn(move |f| {
                    text.iter().try_for_each(|(text, col)| paint(f, col, text))
                })) as Box<dyn Display>)
        });
        let block = Block::new(idx, labels).unwrap();
        block.with_paint(paint).map_code(|c| {
            let c = c.replace('\t', "    ");
            let w = unicode_width::UnicodeWidthStr::width(&*c);
            CodeWidth::new(c, core::cmp::max(w, 1))
        })
    }
}

/// Pretty-printer for file reports.
pub struct FileReportsDisp<'a, P> {
    file_reports: &'a FileReports<P>,
    paint: Paint,
    path: fn(&P) -> String,
}

impl<'a, P> FileReportsDisp<'a, P> {
    /// Construct a new pretty-printer for file reports.
    ///
    /// By default, this does not apply any colors and does not print file paths.
    pub fn new(file_reports: &'a FileReports<P>) -> Self {
        Self {
            file_reports,
            paint: |f, _style, disp| disp.fmt(f),
            path: |_| "".into(),
        }
    }

    /// Set a function that determines how colors should be applied to text.
    pub fn with_paint(mut self, paint: Paint) -> Self {
        self.paint = paint;
        self
    }

    /// Set a function that determines how the file path should be printed.
    pub fn with_path(mut self, path: fn(&P) -> String) -> Self {
        self.path = path;
        self
    }
}

impl<'a, P> Display for FileReportsDisp<'a, P> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let (file, reports) = &self.file_reports;
        let path = (self.path)(&file.path);
        let idx = codesnake::LineIndex::new(&file.code);
        reports.iter().try_for_each(|e| {
            writeln!(f, "Error: {}", e.message)?;
            let block = e.to_block(&idx, self.paint);
            writeln!(f, "{}{}", block.prologue(), path)?;
            writeln!(f, "{}", block.space_vert())?;
            writeln!(f, "{}{}", block, block.epilogue())
        })
    }
}
