hamelin_lib 0.10.2

use std::error::Error;
use std::fmt::{Debug, Formatter, Write};
use std::io::Cursor;
use std::ops::RangeInclusive;
use std::sync::Arc;

use ariadne::{Color, ColorGenerator, Config, Label, Report, ReportKind, Source};
use derive_more::Display;
use serde::Serialize;
use thiserror::Error;
use tsify::Tsify;

use crate::antlr::fit;
use crate::tree::ast::identifier::Identifier;
use crate::tree::ast::node::Spannable;
use crate::types::Type;

/// Trait for contexts that can accumulate errors in a diagnostic sink
///
/// This trait is implemented by both `ParseContext` (for parse-phase errors)
/// and `StatementTranslationContext` (for type-check-phase errors), allowing
/// the `ErrorBuilder` to work generically with either context.
pub trait DiagnosticSink {
    /// Add an error to the diagnostic sink
    fn add_error(&mut self, error: Arc<TranslationError>);
}

#[derive(Debug, Clone, Default, Display, Serialize, PartialEq, Eq, Hash, Tsify)]
#[tsify(into_wasm_abi)]
pub enum Level {
    #[default]
    Error,
    Warning,
    Info,
}

#[derive(Debug, Default, Display, Clone, Copy, Eq, PartialEq, Hash, Serialize, Tsify)]
#[tsify(into_wasm_abi)]
pub enum LanguageArea {
    #[default]
    #[display("finding matching function definitions")]
    FunctionCall,
    #[display("finding matching operator definitions")]
    Operator,
    #[display("performing field lookup (the dot operator)")]
    FieldLookup,
    #[display("performing index access (the square brackets)")]
    IndexAccess,
    #[display("parsing")]
    Parsing,
}

#[derive(Debug, Default, Display, Clone, Copy, Eq, PartialEq, Hash, Serialize, Tsify)]
#[tsify(into_wasm_abi)]
pub enum Stage {
    #[default]
    #[display("translation")]
    Translation,
    #[display("parsing")]
    Parsing,
    #[display("semantic analysis")]
    SemanticAnalysis,
    #[display("normalization")]
    Normalization,
    #[display("lowering to IR")]
    Lowering,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash, thiserror::Error, Serialize, Tsify)]
#[tsify(into_wasm_abi)]
#[error("{}@{:?}: {}\n{}", level, primary.interval, primary.message, source_desc.clone().unwrap_or_default())]
pub struct TranslationError {
    pub area: Option<LanguageArea>,
    pub stage: Stage,
    pub level: Level,
    pub primary: Context,
    pub supporting: Option<Vec<Context>>,
    pub source_desc: Option<String>,
}

#[derive(Debug, thiserror::Error, PartialEq, Eq, Serialize, Tsify, Clone)]
#[tsify(into_wasm_abi)]
#[error("{pretty}")]
pub struct ContextualTranslationError {
    pub error: TranslationError,
    pub pretty: String,
}

impl ContextualTranslationError {
    pub fn new(error: TranslationError, pretty: String) -> Self {
        Self { error, pretty }
    }
}

#[derive(Debug, Clone, Serialize, PartialEq, Eq, Hash, Tsify)]
#[tsify(into_wasm_abi)]
pub struct Context {
    pub interval: RangeInclusive<usize>,
    pub message: String,
}

#[derive(Debug, Clone, thiserror::Error, Default, Serialize, PartialEq, Eq, Tsify)]
#[tsify(into_wasm_abi)]
pub struct TranslationErrors(pub Vec<TranslationError>);

#[derive(Debug, thiserror::Error, Serialize, PartialEq, Eq, Tsify)]
#[tsify(into_wasm_abi)]
pub struct ContextualTranslationErrors {
    pub hamelin: String,
    pub errors: Vec<ContextualTranslationError>,
}

impl Display for ContextualTranslationErrors {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for e in &self.errors {
            writeln!(f, "{}", e)?;
        }

        Ok(())
    }
}

impl ContextualTranslationErrors {
    pub fn new(hamelin: String, errors: TranslationErrors) -> Self {
        Self::with_options(hamelin, None, true, errors)
    }

    pub fn new_with_source_name(
        hamelin: String,
        source_name: Option<String>,
        errors: TranslationErrors,
    ) -> Self {
        Self::with_options(hamelin, source_name, true, errors)
    }

    pub fn with_options(
        hamelin: String,
        source_name: Option<String>,
        color: bool,
        errors: TranslationErrors,
    ) -> Self {
        let prettified_errors = errors
            .into_iter()
            .map(|mut e| {
                e.primary.interval = fit(&e.primary.interval, &hamelin);
                if let Some(supporting) = &mut e.supporting {
                    for c in supporting.iter_mut() {
                        c.interval = fit(&c.interval, &hamelin);
                    }
                }

                let pretty = e
                    .make_pretty(&hamelin, source_name.as_deref(), color)
                    .unwrap_or_default();
                ContextualTranslationError::new(e, pretty)
            })
            .collect();

        Self {
            hamelin,
            errors: prettified_errors,
        }
    }

    pub fn with_source_name(self, source_name: Option<String>, color: bool) -> Self {
        let Self { hamelin, errors } = self;
        let errors = errors
            .into_iter()
            .map(|ce| {
                let pretty = ce
                    .error
                    .make_pretty(&hamelin, source_name.as_deref(), color)
                    .unwrap_or_default();
                ContextualTranslationError::new(ce.error, pretty)
            })
            .collect();
        Self { hamelin, errors }
    }

    pub fn take(mut self, n: usize) -> Self {
        self.errors.truncate(n);
        self
    }
}

impl Context {
    pub fn new(interval: RangeInclusive<usize>, message: &str) -> Self {
        Self {
            interval,
            message: message.to_string(),
        }
    }

    pub fn from(spannable: &impl Spannable, message: &str) -> Self {
        Self::new(spannable.span().unwrap_or(0..=0), message)
    }
}

impl TranslationError {
    pub fn new(ctx: Context) -> Self {
        Self {
            primary: ctx,
            supporting: None,
            source_desc: None,
            area: Default::default(),
            stage: Default::default(),
            level: Default::default(),
        }
    }

    pub fn msg(tree: &impl Spannable, message: &str) -> Self {
        Self {
            primary: Context::from(tree, message),
            area: Default::default(),
            level: Default::default(),
            stage: Default::default(),
            supporting: None,
            source_desc: None,
        }
    }

    pub fn wrap<E>(tree: &impl Spannable, source: E) -> Self
    where
        E: Error + Send + Sync + 'static,
    {
        Self {
            primary: Context::from(tree, "error while translating"),
            source_desc: Some(source.to_string()),
            area: Default::default(),
            stage: Default::default(),
            level: Default::default(),
            supporting: None,
        }
    }

    pub fn wrap_box(tree: &impl Spannable, source: Box<dyn Error + Send + Sync + 'static>) -> Self {
        Self {
            primary: Context::from(tree, "error while translating"),
            source_desc: Some(source.to_string()),
            area: Default::default(),
            stage: Default::default(),
            level: Default::default(),
            supporting: None,
        }
    }

    pub fn fatal(hamelin: &str, e: Box<dyn Error + Send + Sync + 'static>) -> Self {
        Self::new(Context::new(0..=hamelin.len() - 1, "fatal error")).with_source_boxed(e)
    }

    pub fn with_level(mut self, level: Level) -> Self {
        self.level = level;
        self
    }

    pub fn with_area(mut self, area: LanguageArea) -> Self {
        self.area = Some(area);
        self
    }

    pub fn with_stage(mut self, stage: Stage) -> Self {
        self.stage = stage;
        self
    }

    pub fn add_context(&mut self, span: RangeInclusive<usize>, message: &str) {
        match self.supporting {
            Some(ref mut ctx) => ctx.push(Context::new(span, message)),
            None => self.supporting = Some(vec![Context::new(span, message)]),
        }
    }

    pub fn with_context(mut self, tree: &impl Spannable, message: &str) -> Self {
        self.add_context(tree.span().unwrap_or(0..=0), message);
        self
    }

    pub fn with_context_vec(mut self, context: Vec<Context>) -> Self {
        self.supporting = Some(context);
        self
    }

    pub fn with_source<E>(mut self, source: E) -> Self
    where
        E: Error + Send + Sync + 'static,
    {
        self.source_desc = Some(source.to_string());
        self
    }

    pub fn with_source_boxed(mut self, source: Box<dyn Error + Send + Sync + 'static>) -> Self {
        self.source_desc = Some(source.to_string());
        self
    }

    /// A more explicit Into<TranslationErrors> for when inference can't make that work.
    pub fn single(self) -> TranslationErrors {
        TranslationErrors(vec![self])
    }

    pub fn single_result<R>(self) -> Result<R, TranslationErrors> {
        Err(self.single())
    }

    pub fn make_pretty(
        &self,
        hamelin: &String,
        source_name: Option<&str>,
        color: bool,
    ) -> Result<String, std::fmt::Error> {
        let input = hamelin.as_str();
        let source = Source::from(input);
        let mut string_buf = String::new();

        let mut colors = ColorGenerator::default();
        let kind = match self.level {
            Level::Error => ReportKind::Error,
            Level::Warning => ReportKind::Warning,
            Level::Info => ReportKind::Advice,
        };

        let buf = Vec::new();
        let mut cursor = Cursor::new(buf);

        let name = source_name.unwrap_or("<query>").to_string();

        let label_color = |c: &mut ColorGenerator| -> Option<Color> {
            if color {
                Some(c.next())
            } else {
                None
            }
        };

        let mut report = Report::build(kind, (name.clone(), self.primary.interval.clone()))
            .with_config(Config::default().with_color(color))
            .with_message(format!(
                "problem doing {} {}",
                self.stage,
                self.area.map(|a| format!("in {}", a)).unwrap_or_default()
            ))
            .with_label({
                let label = Label::new((name.clone(), self.primary.interval.clone()))
                    .with_message(&self.primary.message);
                match label_color(&mut colors) {
                    Some(c) => label.with_color(c),
                    None => label,
                }
            });

        if let Some(contexts) = &self.supporting {
            for c in contexts {
                let label = Label::new((name.clone(), c.interval.clone()))
                    .with_message(c.message.clone())
                    .with_order(1);
                report = report.with_label(match label_color(&mut colors) {
                    Some(col) => label.with_color(col),
                    None => label,
                });
            }
        }

        if let Some(source_desc) = &self.source_desc {
            report = report.with_note(format!("{}", source_desc));
        }

        report
            .finish()
            .write((name, source), &mut cursor)
            .map_err(|_| std::fmt::Error)?;

        write!(string_buf, "{}", String::from_utf8_lossy(cursor.get_ref()))?;

        if matches!(self.stage, Stage::Normalization | Stage::Lowering) {
            writeln!(string_buf)?;
            writeln!(string_buf, "─── normalized query ───")?;
            writeln!(string_buf, "{hamelin}")?;
        }

        Ok(string_buf)
    }
}

impl TranslationErrors {
    pub fn add(&mut self, error: TranslationError) {
        self.0.push(error);
    }

    pub fn contextualize(self, hamelin: impl Into<String>) -> ContextualTranslationErrors {
        ContextualTranslationErrors::new(hamelin.into(), self)
    }

    pub fn contextualize_plain(self, hamelin: impl Into<String>) -> ContextualTranslationErrors {
        ContextualTranslationErrors::with_options(hamelin.into(), None, false, self)
    }

    /// Alias for add_all to match standard library naming
    pub fn extend(&mut self, mut errors: TranslationErrors) {
        self.0.append(&mut errors.0);
    }

    pub fn with_all(mut self, other: Self) -> Self {
        self.extend(other);
        self
    }

    pub fn or_ok<R>(self, result: R) -> Result<R, TranslationErrors> {
        if self.is_empty() {
            Ok(result)
        } else {
            Err(self)
        }
    }

    pub fn expect<O>(tree: &impl Spannable, opt: Option<O>) -> Result<O, Self> {
        opt.ok_or_else(|| {
            TranslationError::msg(tree, "expected value not present during parse of clause")
                .single()
        })
    }

    pub fn consume_errors<R>(&mut self, result: Result<R, TranslationErrors>) -> Option<R> {
        match result {
            Ok(r) => Some(r),
            Err(e) => {
                self.extend(e);
                None
            }
        }
    }

    pub fn from_vec<R>(
        results: Vec<Result<R, TranslationErrors>>,
    ) -> Result<Vec<R>, TranslationErrors> {
        let mut errors = vec![];
        let mut successes = vec![];
        for r in results {
            match r {
                Ok(r) => successes.push(r),
                Err(e) => errors.push(e),
            }
        }

        if !errors.is_empty() {
            Err(errors.into())
        } else {
            Ok(successes)
        }
    }

    pub fn from_2<R1, R2>(
        result1: Result<R1, TranslationErrors>,
        result2: Result<R2, TranslationErrors>,
    ) -> Result<(R1, R2), TranslationErrors> {
        match (result1, result2) {
            (Ok(r1), Ok(r2)) => Ok((r1, r2)),
            (r1, r2) => {
                let mut errors = Self::default();
                if let Err(e) = r1 {
                    errors.extend(e);
                }
                if let Err(e) = r2 {
                    errors.extend(e);
                }
                Err(errors)
            }
        }
    }

    pub fn from_2_option<R1, R2>(
        err: TranslationError,
        result1: Option<R1>,
        result2: Option<R2>,
    ) -> Result<(R1, R2), TranslationErrors> {
        match (result1, result2) {
            (Some(r1), Some(r2)) => Ok((r1, r2)),
            _ => Err(err.single()),
        }
    }

    pub fn from_3<R1, R2, R3>(
        result1: Result<R1, TranslationErrors>,
        result2: Result<R2, TranslationErrors>,
        result3: Result<R3, TranslationErrors>,
    ) -> Result<(R1, R2, R3), TranslationErrors> {
        match (result1, result2, result3) {
            (Ok(r1), Ok(r2), Ok(r3)) => Ok((r1, r2, r3)),
            (r1, r2, r3) => {
                let mut errors = Self::default();
                if let Err(e) = r1 {
                    errors.extend(e);
                }
                if let Err(e) = r2 {
                    errors.extend(e);
                }
                if let Err(e) = r3 {
                    errors.extend(e);
                }
                Err(errors)
            }
        }
    }

    pub fn maybe_with_area(mut self, area: LanguageArea) -> Self {
        for error in self.0.iter_mut() {
            if error.area.is_none() {
                error.area = Some(area.clone());
            }
        }
        self
    }

    pub fn maybe_with_stage(mut self, stage: Stage) -> Self {
        for error in self.0.iter_mut() {
            if error.stage == Stage::Translation {
                error.stage = stage.clone();
            }
        }
        self
    }

    pub fn iter(&self) -> std::slice::Iter<'_, TranslationError> {
        self.0.iter()
    }

    pub fn into_iter(self) -> std::vec::IntoIter<TranslationError> {
        self.0.into_iter()
    }

    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub fn len(&self) -> usize {
        self.0.len()
    }
}

impl From<TranslationError> for TranslationErrors {
    fn from(error: TranslationError) -> Self {
        Self(vec![error])
    }
}

impl From<ContextualTranslationErrors> for TranslationErrors {
    fn from(c: ContextualTranslationErrors) -> Self {
        Self(c.errors.into_iter().map(|ce| ce.error).collect())
    }
}

impl Display for TranslationErrors {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        for error in &self.0 {
            write!(f, "{}\n\n", error)?;
        }
        Ok(())
    }
}

impl From<Vec<TranslationErrors>> for TranslationErrors {
    fn from(errors: Vec<TranslationErrors>) -> Self {
        let mut ret = Self::default();
        for mut error in errors.into_iter() {
            ret.0.append(&mut error.0);
        }
        ret
    }
}

#[derive(thiserror::Error, Debug)]
pub struct UnexpectedType {
    pub expr_type: Type,
    pub expected: Vec<Type>,
}

impl UnexpectedType {
    pub fn new(typ: Type, expected: Vec<Type>) -> Self {
        Self {
            expr_type: typ,
            expected,
        }
    }
}

impl Display for UnexpectedType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        writeln!(f, "unexpected type: {}\n", self.expr_type)?;

        if self.expected.len() == 1 {
            writeln!(f, "expected: {}\n", self.expected[0])?;
        } else {
            writeln!(f, "expected one of:\n")?;
            for exp in &self.expected {
                write!(f, "  {}\n", exp)?;
            }
        }
        Ok(())
    }
}

#[derive(Error, Debug)]
#[error("Invalid cast attempt from {from} to {to}")]
pub struct InvalidCast {
    from: Type,
    to: Type,
}

impl InvalidCast {
    pub fn new(from: Type, to: Type) -> Self {
        Self { from, to }
    }
}

#[derive(thiserror::Error, Debug, Eq, PartialEq)]
pub struct NonMergeableTypes {
    pub base: Type,
    pub other: Type,
    pub at: Option<Identifier>,
}

impl NonMergeableTypes {
    pub fn new(base: Type, other: Type) -> Self {
        Self {
            base,
            other,
            at: None,
        }
    }

    pub fn prepend_at(mut self, at: Identifier) -> Self {
        self.at = match self.at {
            None => Some(at),
            Some(i) => Some(i.prepend(at).into()),
        };
        self
    }
}

impl Display for NonMergeableTypes {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match &self.at {
            Some(at) => write!(
                f,
                "The types found at {} are not mergeable: {}, {}",
                at, self.base, self.other
            ),
            None => write!(
                f,
                "The types are not mergeable: {}, {}",
                self.base, self.other
            ),
        }
    }
}

/// Wrapper holding a compilation output alongside any errors that occurred.
///
/// All parse/type-check functions return this type. Callers choose their
/// error-handling style:
///
/// - `.result()` — strict: `Err` if any errors, `Ok` otherwise
/// - `.output` / `.errors` — access fields directly
/// - Struct destructuring — `let WithTranslationErrors { output, errors } = …;`
pub struct WithTranslationErrors<T> {
    pub output: T,
    pub errors: TranslationErrors,
}

impl<T> WithTranslationErrors<T> {
    pub fn new(output: T, errors: TranslationErrors) -> Self {
        Self { output, errors }
    }

    /// Convert to `Result` — `Err` if any errors, `Ok(output)` otherwise.
    pub fn into_result(self) -> Result<T, TranslationErrors> {
        if self.errors.is_empty() {
            Ok(self.output)
        } else {
            Err(self.errors)
        }
    }
}