use std::collections::HashMap;

use ariadne::{Cache, Label, Report, ReportKind, Source};
use std::ops::Range;

/// Stable identifier for each diagnostic kind.
///
/// The default severity for each code is defined by [`DiagCode::default_severity`].
/// Users can override any code's severity via [`DiagConfig`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum DiagCode {
    // --- Parse phase ---
    ParseIntImaginary,
    ParseUnknownDType,
    ParseSyntax,
    ParseUndefinedVar,
    ParseReservedIdentifier,

    // --- Lower phase ---
    LowerUndefinedVar,
    LowerTypeMismatch,
    LowerIntOutOfRange,
    LowerUnknownFunction,
    LowerArgCount,
    LowerDivByZero,
    LowerTupleArity,

    // --- Lower phase ---
    /// Implicit type promotion applied to operands (default: Info).
    LowerImplicitPromotion,
    /// Float literal loses precision when stored in a narrower dtype (default: Warning).
    LowerFloatPrecisionLoss,
    /// Mixing u64 with signed integers forces a potentially lossy cast to i64 (default: Warning).
    LowerUnsafeUIntToSInt,
    /// Multiple sequential casts that could likely be simplified (default: Warning).
    LowerRedundantCastChain,

    // --- Opt phase ---
    /// Type promotion applied during expression reconstruction after optimisation (default: Info).
    OptPromotion,
    /// Constant folding encountered a potentially unsafe operation (default: Warning).
    ConstFoldUnsafe,

    // --- Compose phase ---
    /// Symbol name appears in multiple composed expressions (default: Warning).
    ComposeSymbolReuse,
}

impl DiagCode {
    /// Built-in severity used when the user has not overridden this code.
    pub fn default_severity(self) -> Severity {
        use DiagCode::*;
        match self {
            // Parse errors that are always fatal and cannot be suppressed.
            ParseIntImaginary | ParseUnknownDType | ParseSyntax => Severity::Fatal,
            // Parse-time name resolution errors: treat as normal errors so users can configure them.
            ParseUndefinedVar | ParseReservedIdentifier => Severity::Error,
            // Lowering hard errors (may be reclassified by configuration).
            LowerUndefinedVar | LowerTypeMismatch | LowerIntOutOfRange | LowerUnknownFunction
            | LowerArgCount | LowerDivByZero | LowerTupleArity => Severity::Error,
            // Soft diagnostics.
            LowerImplicitPromotion => Severity::Info,
            LowerFloatPrecisionLoss => Severity::Warning,
            LowerUnsafeUIntToSInt => Severity::Warning,
            LowerRedundantCastChain => Severity::Warning,
            // OptPromotion: Allow by default; L301 already reports implicit promotions with precise spans.
            OptPromotion => Severity::Allow,
            ConstFoldUnsafe => Severity::Warning,
            ComposeSymbolReuse => Severity::Warning,
        }
    }

    /// Short stable diagnostic code (e.g. `P001`, `L002`) suitable for humans.
    ///
    /// Prefix encodes the phase, first digit encodes default severity:
    /// - `P` = parse
    /// - `L` = lower/IR
    /// - `O` = optimisation
    /// - `0` = fatal (cannot be suppressed)
    /// - `1` = error
    /// - `2` = warning
    /// - `3` = suggestion / info
    pub fn short_code(self) -> &'static str {
        use DiagCode::*;
        match self {
            // Parse phase (fatal)
            ParseIntImaginary => "P001",
            ParseUnknownDType => "P002",
            ParseSyntax => "P003",
            ParseUndefinedVar => "P101",
            ParseReservedIdentifier => "P102",

            // Lower / IR phase
            LowerUndefinedVar => "L101",
            LowerTypeMismatch => "L102",
            LowerIntOutOfRange => "L103",
            LowerUnknownFunction => "L104",
            LowerArgCount => "L105",
            LowerDivByZero => "L106",
            LowerTupleArity => "L107",
            LowerFloatPrecisionLoss => "L201",
            LowerUnsafeUIntToSInt => "L202",
            LowerRedundantCastChain => "L203",
            LowerImplicitPromotion => "L301",

            // Optimisation phase
            ConstFoldUnsafe => "O201",
            OptPromotion => "O301",

            // Compose phase
            ComposeSymbolReuse => "C201",
        }
    }

    /// Default help message for this diagnostic code, if any.
    ///
    /// Call-sites can override this with a more specific, context-aware hint by
    /// passing an explicit `hint` when constructing [`Diag`]. When `hint` is
    /// `None`, [`Diag::new`] will fall back to this default.
    pub fn default_help(self) -> Option<&'static str> {
        use DiagCode::*;
        match self {
            // Parse phase
            ParseIntImaginary => {
                Some("use a float imaginary literal instead, e.g. `2.0j` instead of `2j`")
            }
            ParseUnknownDType => Some(
                "use a supported dtype name such as `i32`, `i64`, `u8`, `f32`, `f64`, or `c128`",
            ),
            ParseSyntax => {
                Some("check the surrounding tokens and parentheses for typos or missing separators")
            }
            ParseUndefinedVar => {
                Some("declare the variable earlier or fix the name to match an existing binding")
            }
            ParseReservedIdentifier => Some(
                "choose a different identifier; reserved words cannot be used as variable names",
            ),

            // Lower / IR hard errors
            LowerUndefinedVar => {
                Some("ensure the variable is defined as a parameter or let-binding before use")
            }
            LowerTypeMismatch => Some(
                "make operand dtypes compatible, for example by inserting explicit casts using `@dtype`",
            ),
            LowerIntOutOfRange => Some(
                "adjust the literal value or choose a wider integer dtype that can represent it",
            ),
            LowerUnknownFunction => Some(
                "define the function or call a supported builtin such as `exp`, `log`, `sqrt`, `floor`, or `ceil`",
            ),
            LowerArgCount => {
                Some("add or remove arguments so the call matches the function's parameter count")
            }
            LowerDivByZero => Some(
                "ensure the divisor cannot be zero, e.g. by guarding with an `if` or changing the expression",
            ),
            LowerTupleArity => Some(
                "use a scalar expression here, or destructure the tuple before using it in a scalar position",
            ),

            // Lower / IR soft diagnostics
            LowerImplicitPromotion => Some(
                "consider inserting an explicit cast with `@dtype` on one operand to make the promotion clear",
            ),
            LowerFloatPrecisionLoss => Some(
                "use a wider float dtype (e.g. `f32` or `f64`) if you need to preserve full precision",
            ),
            LowerUnsafeUIntToSInt => Some(
                "avoid mixing `u64` with signed ints, or cast to a wider float/complex dtype if large values are expected",
            ),
            LowerRedundantCastChain => Some(
                "simplify the expression to a single cast like `expr@dtype` unless an intermediate cast is required",
            ),

            // Optimisation phase
            OptPromotion => Some(
                "if this promotion is surprising, make the intended dtype explicit in the source using `@dtype` casts",
            ),
            ConstFoldUnsafe => Some(
                "rewrite the constant expression to avoid overflow, or choose a wider integer/float dtype for the result",
            ),

            // Compose phase
            ComposeSymbolReuse => Some(
                "rename one of the symbols if you intended the outer and inner bindings to be distinct",
            ),
        }
    }

    /// Short kebab-case identifier (used in CLI flags like `--allow=lower-type-mismatch`).
    pub fn as_str(self) -> &'static str {
        use DiagCode::*;
        match self {
            ParseIntImaginary => "parse-int-imaginary",
            ParseUnknownDType => "parse-unknown-dtype",
            ParseSyntax => "parse-syntax",
            ParseUndefinedVar => "parse-undefined-var",
            ParseReservedIdentifier => "parse-reserved-identifier",
            LowerUndefinedVar => "lower-undefined-var",
            LowerTypeMismatch => "lower-type-mismatch",
            LowerIntOutOfRange => "lower-int-out-of-range",
            LowerUnknownFunction => "lower-unknown-function",
            LowerArgCount => "lower-arg-count",
            LowerDivByZero => "lower-div-by-zero",
            LowerTupleArity => "lower-tuple-arity",
            LowerImplicitPromotion => "lower-implicit-promotion",
            LowerFloatPrecisionLoss => "lower-float-precision-loss",
            LowerUnsafeUIntToSInt => "lower-unsafe-uint-to-sint",
            LowerRedundantCastChain => "lower-redundant-cast-chain",
            OptPromotion => "opt-promotion",
            ConstFoldUnsafe => "const-fold-unsafe",
            ComposeSymbolReuse => "compose-symbol-reuse",
        }
    }
}

/// Diagnostic severity level, ordered from least to most severe.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Severity {
    /// Suppress entirely – the diagnostic is not emitted.
    Allow,
    /// Informational hint; compilation always continues.
    Info,
    /// Warning; compilation continues unless `warnings_as_errors` is set.
    Warning,
    /// Error; compilation stops (but may be demoted/promoted by configuration).
    Error,
    /// Fatal error; compilation stops and this diagnostic cannot be suppressed.
    Fatal,
}

impl From<Severity> for ReportKind<'static> {
    fn from(sev: Severity) -> Self {
        match sev {
            Severity::Allow => {
                panic!(
                    "Severity::Allow should never be converted to ReportKind; filter it earlier"
                );
            }
            Severity::Info => ReportKind::Advice,
            Severity::Warning => ReportKind::Warning,
            Severity::Error | Severity::Fatal => ReportKind::Error,
        }
    }
}

/// Per-session diagnostic configuration.
///
/// Start from [`DiagConfig::default()`] (all codes use their built-in defaults)
/// and customise using the builder methods.
///
/// ```
/// use gonidium::diag::{DiagCode, DiagConfig, Severity};
///
/// let cfg = DiagConfig::default()
///     .set(DiagCode::LowerImplicitPromotion, Severity::Allow)  // silence promotions
///     .set(DiagCode::LowerFloatPrecisionLoss, Severity::Error) // make precision loss fatal
///     .with_warnings_as_errors();                              // treat remaining warnings as errors
/// ```
#[derive(Debug, Clone, Default)]
pub struct DiagConfig {
    overrides: HashMap<DiagCode, Severity>,
    warnings_as_errors: bool,
}

impl DiagConfig {
    pub fn new() -> Self {
        Self::default()
    }

    /// Override the severity of `code`.
    pub fn set(mut self, code: DiagCode, severity: Severity) -> Self {
        self.overrides.insert(code, severity);
        self
    }

    /// Treat all `Warning`-severity diagnostics as `Error` (like `-Werror`).
    ///
    /// Explicit per-code overrides take precedence over this flag.
    pub fn with_warnings_as_errors(mut self) -> Self {
        self.warnings_as_errors = true;
        self
    }

    /// Query the effective severity for `code`, respecting overrides.
    pub fn severity(&self, code: DiagCode) -> Severity {
        let default = code.default_severity();
        // Fatal diagnostics are never configurable: they always remain fatal.
        if default == Severity::Fatal {
            return Severity::Fatal;
        }

        let base = self.overrides.get(&code).copied().unwrap_or(default);
        if self.warnings_as_errors && base == Severity::Warning {
            Severity::Error
        } else {
            base
        }
    }
}

/// A single compiler diagnostic (info, warning, or error).
///
/// Construct via the `Diag::new` constructor or the phase-specific helpers in
/// `ir::error` and `parse::error`.  Use [`Diag::report_string`] to render
/// an ariadne-formatted report, passing the effective [`Severity`] from a
/// [`DiagConfig`] (or call [`Diag::report_string_default`] to use the code's
/// built-in default severity).
/// Optional secondary labels for multi-span diagnostics: `(start, end, message)`.
pub type DiagLabel = (usize, usize, String);

#[derive(Debug, Clone)]
pub struct Diag {
    pub code: DiagCode,
    pub message: String,
    /// Source byte-offset span `(start, end)`, if available.
    pub span: Option<(usize, usize)>,
    /// Secondary labels (e.g. "left operand", "op", "right operand") for precise highlighting.
    pub labels: Option<Vec<DiagLabel>>,
    pub hint: Option<String>,
}

/// Simple cache that exposes a single in-memory source to ariadne.
struct SingleSourceCache<'src> {
    id: &'src str,
    source: Source<&'src str>,
}

impl<'src> Cache<&'src str> for SingleSourceCache<'src> {
    type Storage = &'src str;

    fn fetch(&mut self, id: &&'src str) -> Result<&Source<Self::Storage>, impl std::fmt::Debug> {
        if *id == self.id {
            Ok(&self.source)
        } else {
            Err(format!("unknown source id: {id}"))
        }
    }

    fn display<'a>(&self, id: &'a &'src str) -> Option<impl std::fmt::Display + 'a> {
        Some(id)
    }
}

impl Diag {
    pub fn new(
        code: DiagCode,
        message: impl Into<String>,
        span: Option<(usize, usize)>,
        hint: Option<String>,
    ) -> Self {
        let hint = match hint {
            Some(h) => Some(h),
            None => code.default_help().map(|s| s.to_string()),
        };
        Self {
            code,
            message: message.into(),
            span,
            labels: None,
            hint,
        }
    }

    /// Render this diagnostic using the given `severity`.
    /// Returns an empty string when `severity == Severity::Allow`.
    pub fn report_string(&self, source_id: &str, src: &str, severity: Severity) -> String {
        if severity == Severity::Allow {
            return String::new();
        }
        let range: Range<usize> = match self.span {
            Some((start, end)) => start..end,
            None => 0..src.len().min(1),
        };
        let kind: ReportKind<'static> = severity.into();
        let mut out = Vec::new();
        let mut builder = Report::build(kind, (source_id, range.clone()))
            .with_code(self.code.short_code())
            .with_message(&self.message);
        if let Some(labels) = &self.labels {
            for (order, (start, end, msg)) in labels.iter().enumerate() {
                let r = *start..*end;
                builder = builder.with_label(
                    Label::new((source_id, r))
                        .with_message(msg.as_str())
                        .with_order(order as i32),
                );
            }
        } else {
            builder = builder.with_label(Label::new((source_id, range)));
        }
        if let Some(hint) = &self.hint {
            builder = builder.with_help(hint);
        }
        let cache = SingleSourceCache {
            id: source_id,
            source: Source::from(src),
        };
        let _ = builder.finish().write(cache, &mut out);
        String::from_utf8_lossy(&out).to_string()
    }

    /// Render using the code's built-in default severity.
    pub fn report_string_default(&self, source_id: &str, src: &str) -> String {
        self.report_string(source_id, src, self.code.default_severity())
    }
}

impl std::fmt::Display for Diag {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "[{}] {}", self.code.short_code(), self.message)
    }
}

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

/// Accumulates diagnostics during a compilation phase according to a [`DiagConfig`].
///
/// - `Severity::Allow` → diagnostic is silently dropped.
/// - `Severity::Info` / `Severity::Warning` → pushed to [`Self::diags`].
/// - `Severity::Error` / `Severity::Fatal` → `emit` returns `Err(diag)` so the caller can propagate it.
pub struct DiagCollector {
    pub diags: Vec<Diag>,
    config: DiagConfig,
}

impl DiagCollector {
    pub fn new(config: DiagConfig) -> Self {
        Self {
            diags: Vec::new(),
            config,
        }
    }

    /// Emit a diagnostic.  Returns `Err(diag)` if the effective severity is `Error` or `Fatal`.
    pub fn emit(&mut self, diag: Diag) -> Result<(), Diag> {
        match self.config.severity(diag.code) {
            Severity::Allow => Ok(()),
            Severity::Info | Severity::Warning => {
                self.diags.push(diag);
                Ok(())
            }
            Severity::Error | Severity::Fatal => Err(diag),
        }
    }

    pub fn config(&self) -> &DiagConfig {
        &self.config
    }

    pub fn severity(&self, code: DiagCode) -> Severity {
        self.config.severity(code)
    }

    pub fn into_diags(self) -> Vec<Diag> {
        self.diags
    }
}

impl Default for DiagCollector {
    fn default() -> Self {
        Self::new(DiagConfig::default())
    }
}
