use std::collections::HashMap;
use std::fmt::{self, Write};
use std::path::{Path, PathBuf};

use itertools::Itertools;

use crate::codespan;
use crate::diagnostic::{Diagnostic, DiagnosticCode};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("io error: {0}")]
    Io(#[from] std::io::Error),

    #[error("invalid path: {path}")]
    InvalidPath { path: PathBuf },

    #[error("cannot find project root")]
    CannotFindProjectRoot,

    #[error("cannot read source file at {file}:\n  {io}")]
    CannotReadSourceFile {
        file: std::path::PathBuf,
        io: std::io::Error,
    },

    #[error("{}", DisplayMessages(.diagnostics))]
    Compile { diagnostics: Vec<DiagnosticMessage> },
}

impl Error {
    pub(crate) fn from_diagnostics(
        diagnostics: Vec<Diagnostic>,
        sources: &impl crate::project::SourceProvider,
    ) -> Self {
        let diagnostics = compose_diagnostic_messages(diagnostics, sources);
        Error::Compile { diagnostics }
    }
}

#[derive(Debug)]
pub struct DiagnosticMessage {
    diagnostic: Diagnostic,
    display: String,
    range: codespan::Range,
    additional_ranges: Vec<Option<codespan::Range>>,
}

impl DiagnosticMessage {
    pub fn code(&self) -> &'static str {
        self.diagnostic.code.get()
    }

    pub fn message(&self) -> &str {
        &self.diagnostic.message
    }

    pub fn span(&self) -> &Option<crate::Span> {
        &self.diagnostic.span
    }

    pub fn additional(&self) -> Vec<Additional<'_>> {
        self.diagnostic
            .additional
            .iter()
            .enumerate()
            .map(|(i, additional)| Additional {
                additional,
                range: self.additional_ranges.get(i).and_then(|x| x.as_ref()),
            })
            .collect()
    }

    pub fn display(&self) -> &str {
        &self.display
    }

    pub fn range(&self) -> &codespan::Range {
        &self.range
    }
}

#[derive(Debug)]
pub struct Additional<'d> {
    additional: &'d crate::diagnostic::Additional,
    range: Option<&'d codespan::Range>,
}

impl<'d> Additional<'d> {
    pub fn message(&self) -> &'d str {
        &self.additional.message
    }

    pub fn span(&self) -> Option<codespan::Span> {
        self.additional.span
    }

    pub fn range(&self) -> Option<&'d codespan::Range> {
        self.range
    }
}

fn compose_diagnostic_messages(
    diagnostics: Vec<Diagnostic>,
    sources: &impl crate::project::SourceProvider,
) -> Vec<DiagnosticMessage> {
    let mut cache = FileTreeCache::new(sources);

    let mut messages = Vec::with_capacity(diagnostics.len());
    for diagnostic in diagnostics {
        let Some(span) = diagnostic.span else {
            panic!(
                "missing diagnostic span: [{:?}] {}, {:#?}",
                diagnostic.code, diagnostic.message, diagnostic.additional
            );
        };

        let range = compose_range(span, sources, &mut cache);

        let display = compose_display(&diagnostic, sources, &mut cache);

        let mut additional_ranges = Vec::with_capacity(diagnostic.additional.len());
        for a in &diagnostic.additional {
            let range = a.span.map(|s| compose_range(s, sources, &mut cache));
            additional_ranges.push(range);
        }

        messages.push(DiagnosticMessage {
            diagnostic,
            display,
            range,
            additional_ranges,
        });
    }
    messages
}

fn compose_display<S>(
    diagnostic: &Diagnostic,
    sources: &impl crate::project::SourceProvider,
    cache: &mut FileTreeCache<S>,
) -> String
where
    S: crate::project::SourceProvider,
{
    use ariadne::{Config, Label, Report, ReportKind};

    let config = Config::default().with_color(false);

    let span = diagnostic.span.unwrap();
    let (source_path, _) = sources.get_by_id(span.source_id).unwrap();
    let span = std::ops::Range::from(span);

    let kind = match diagnostic.code.get_severity() {
        crate::diagnostic::Severity::Warning => ReportKind::Warning,
        crate::diagnostic::Severity::Error => ReportKind::Error,
    };

    let mut report = Report::build(kind, source_path, span.start)
        .with_config(config)
        .with_label(Label::new((source_path, span)).with_message(&diagnostic.message));

    if diagnostic.code != DiagnosticCode::CUSTOM {
        report = report.with_code(diagnostic.code.get());
    }

    let mut notes = String::new();
    for additional in &diagnostic.additional {
        if let Some(span) = additional.span {
            let span = std::ops::Range::from(span);
            report.add_label(Label::new((source_path, span)).with_message(&diagnostic.message))
        } else {
            notes += "\n   │ ";
            notes += &additional.message;
        }
    }
    if !notes.is_empty() {
        report.set_note(notes);
    }

    let mut out = Vec::new();
    report.finish().write(cache, &mut out).unwrap();
    let out = String::from_utf8(out).unwrap();
    out.lines().map(|l| l.trim_end()).join("\n")
}

fn compose_range<S>(
    span: codespan::Span,
    sources: &impl crate::project::SourceProvider,
    cache: &mut FileTreeCache<S>,
) -> codespan::Range
where
    S: crate::project::SourceProvider,
{
    use ariadne::Cache;

    let (source_path, _content) = sources.get_by_id(span.source_id).unwrap();
    let source = cache.fetch(&source_path).unwrap();
    let source_len = source.len();

    let Some(start) = source.get_byte_line(span.start as usize) else {
        panic!("span {span:?} is out of bounds of the source (len = {source_len})",);
    };
    let start = codespan::LineColumn {
        line: start.1 as u32,
        column: start.2 as u32,
    };

    let Some(end) = source.get_byte_line(span.start as usize + span.len as usize) else {
        panic!("span {span:?} is out of bounds of the source (len = {source_len})",);
    };
    let end = codespan::LineColumn {
        line: end.1 as u32,
        column: end.2 as u32,
    };
    codespan::Range { start, end }
}

struct FileTreeCache<'a, S: crate::project::SourceProvider> {
    provider: &'a S,
    cache: HashMap<PathBuf, ariadne::Source>,
}
impl<'a, S: crate::project::SourceProvider> FileTreeCache<'a, S> {
    fn new(file_tree: &'a S) -> Self {
        FileTreeCache {
            provider: file_tree,
            cache: HashMap::new(),
        }
    }
}

impl<'a, S: crate::project::SourceProvider> ariadne::Cache<&Path> for FileTreeCache<'a, S> {
    type Storage = String;
    fn fetch(&mut self, path: &&Path) -> Result<&ariadne::Source, Box<dyn fmt::Debug + '_>> {
        let (_, content) = match self.provider.get_by_path(path) {
            Some(v) => v,
            None => return Err(Box::new(format!("Unknown file `{path:?}`"))),
        };

        Ok(self
            .cache
            .entry((*path).to_owned())
            .or_insert_with(|| ariadne::Source::from(content.to_string())))
    }

    fn display<'b>(&self, id: &&'b Path) -> Option<Box<dyn fmt::Display + 'b>> {
        if id.as_os_str().is_empty() {
            Some(Box::new(
                self.provider.get_root().file_name()?.display().to_string(),
            ))
        } else {
            Some(Box::new(id.display()))
        }
    }
}

struct DisplayMessages<'a>(&'a Vec<DiagnosticMessage>);

impl<'a> std::fmt::Display for DisplayMessages<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for d in self.0 {
            f.write_str(&d.display)?;
            f.write_char('\n')?;
        }
        Ok(())
    }
}
