use rustc_data_structures::fx::FxIndexMap;
use rustc_hir::intravisit::{self, Visitor};
use rustc_hir::{self as hir, LifetimeSource};
use rustc_session::{declare_lint, declare_lint_pass};
use rustc_span::Span;
use rustc_span::def_id::LocalDefId;
use tracing::instrument;

use crate::{LateContext, LateLintPass, LintContext, lints};

declare_lint! {
    /// The `mismatched_lifetime_syntaxes` lint detects when the same
    /// lifetime is referred to by different syntaxes between function
    /// arguments and return values.
    ///
    /// The three kinds of syntaxes are:
    ///
    /// 1. Named lifetimes. These are references (`&'a str`) or paths
    ///    (`Person<'a>`) that use a lifetime with a name, such as
    ///    `'static` or `'a`.
    ///
    /// 2. Elided lifetimes. These are references with no explicit
    ///    lifetime (`&str`), references using the anonymous lifetime
    ///    (`&'_ str`), and paths using the anonymous lifetime
    ///    (`Person<'_>`).
    ///
    /// 3. Hidden lifetimes. These are paths that do not contain any
    ///    visual indication that it contains a lifetime (`Person`).
    ///
    /// ### Example
    ///
    /// ```rust,compile_fail
    /// #![deny(mismatched_lifetime_syntaxes)]
    ///
    /// pub fn mixing_named_with_elided(v: &'static u8) -> &u8 {
    ///     v
    /// }
    ///
    /// struct Person<'a> {
    ///     name: &'a str,
    /// }
    ///
    /// pub fn mixing_hidden_with_elided(v: Person) -> Person<'_> {
    ///     v
    /// }
    ///
    /// struct Foo;
    ///
    /// impl Foo {
    ///     // Lifetime elision results in the output lifetime becoming
    ///     // `'static`, which is not what was intended.
    ///     pub fn get_mut(&'static self, x: &mut u8) -> &mut u8 {
    ///         unsafe { &mut *(x as *mut _) }
    ///     }
    /// }
    /// ```
    ///
    /// {{produces}}
    ///
    /// ### Explanation
    ///
    /// Lifetime elision is useful because it frees you from having to
    /// give each lifetime its own name and show the relation of input
    /// and output lifetimes for common cases. However, a lifetime
    /// that uses inconsistent syntax between related arguments and
    /// return values is more confusing.
    ///
    /// In certain `unsafe` code, lifetime elision combined with
    /// inconsistent lifetime syntax may result in unsound code.
    pub MISMATCHED_LIFETIME_SYNTAXES,
    Warn,
    "detects when a lifetime uses different syntax between arguments and return values"
}

declare_lint_pass!(LifetimeSyntax => [MISMATCHED_LIFETIME_SYNTAXES]);

impl<'tcx> LateLintPass<'tcx> for LifetimeSyntax {
    #[instrument(skip_all)]
    fn check_fn(
        &mut self,
        cx: &LateContext<'tcx>,
        _: intravisit::FnKind<'tcx>,
        fd: &'tcx hir::FnDecl<'tcx>,
        _: &'tcx hir::Body<'tcx>,
        _: Span,
        _: LocalDefId,
    ) {
        check_fn_like(cx, fd);
    }

    #[instrument(skip_all)]
    fn check_trait_item(&mut self, cx: &LateContext<'tcx>, ti: &'tcx hir::TraitItem<'tcx>) {
        match ti.kind {
            hir::TraitItemKind::Const(..) => {}
            hir::TraitItemKind::Fn(fn_sig, _trait_fn) => check_fn_like(cx, fn_sig.decl),
            hir::TraitItemKind::Type(..) => {}
        }
    }

    #[instrument(skip_all)]
    fn check_foreign_item(&mut self, cx: &LateContext<'tcx>, fi: &'tcx hir::ForeignItem<'tcx>) {
        match fi.kind {
            hir::ForeignItemKind::Fn(fn_sig, _idents, _generics) => check_fn_like(cx, fn_sig.decl),
            hir::ForeignItemKind::Static(..) => {}
            hir::ForeignItemKind::Type => {}
        }
    }
}

fn check_fn_like<'tcx>(cx: &LateContext<'tcx>, fd: &'tcx hir::FnDecl<'tcx>) {
    if fd.inputs.is_empty() {
        return;
    }
    let hir::FnRetTy::Return(output) = fd.output else {
        return;
    };

    let mut map: FxIndexMap<hir::LifetimeKind, LifetimeGroup<'_>> = FxIndexMap::default();

    LifetimeInfoCollector::collect(output, |info| {
        let group = map.entry(info.lifetime.kind).or_default();
        group.outputs.push(info);
    });
    if map.is_empty() {
        return;
    }

    for input in fd.inputs {
        LifetimeInfoCollector::collect(input, |info| {
            if let Some(group) = map.get_mut(&info.lifetime.kind) {
                group.inputs.push(info);
            }
        });
    }

    for LifetimeGroup { ref inputs, ref outputs } in map.into_values() {
        if inputs.is_empty() {
            continue;
        }
        if !lifetimes_use_matched_syntax(inputs, outputs) {
            emit_mismatch_diagnostic(cx, inputs, outputs);
        }
    }
}

#[derive(Default)]
struct LifetimeGroup<'tcx> {
    inputs: Vec<Info<'tcx>>,
    outputs: Vec<Info<'tcx>>,
}

#[derive(Debug, Copy, Clone, PartialEq)]
enum LifetimeSyntaxCategory {
    Hidden,
    Elided,
    Named,
}

impl LifetimeSyntaxCategory {
    fn new(lifetime: &hir::Lifetime) -> Option<Self> {
        use LifetimeSource::*;
        use hir::LifetimeSyntax::*;

        match (lifetime.syntax, lifetime.source) {
            // E.g. `&T`.
            (Implicit, Reference) |
            // E.g. `&'_ T`.
            (ExplicitAnonymous, Reference) |
            // E.g. `ContainsLifetime<'_>`.
            (ExplicitAnonymous, Path { .. }) |
            // E.g. `+ '_`, `+ use<'_>`.
            (ExplicitAnonymous, OutlivesBound | PreciseCapturing) => Some(Self::Elided),

            // E.g. `ContainsLifetime`.
            (Implicit, Path { .. }) => Some(Self::Hidden),

            // E.g. `&'a T`.
            (ExplicitBound, Reference) |
            // E.g. `ContainsLifetime<'a>`.
            (ExplicitBound, Path { .. }) |
            // E.g. `+ 'a`, `+ use<'a>`.
            (ExplicitBound, OutlivesBound | PreciseCapturing) => Some(Self::Named),

            (Implicit, OutlivesBound | PreciseCapturing) | (_, Other) => None,
        }
    }
}

#[derive(Debug, Default)]
pub struct LifetimeSyntaxCategories<T> {
    pub hidden: T,
    pub elided: T,
    pub named: T,
}

impl<T> LifetimeSyntaxCategories<T> {
    fn select(&mut self, category: LifetimeSyntaxCategory) -> &mut T {
        use LifetimeSyntaxCategory::*;

        match category {
            Elided => &mut self.elided,
            Hidden => &mut self.hidden,
            Named => &mut self.named,
        }
    }
}

impl<T> LifetimeSyntaxCategories<Vec<T>> {
    pub fn len(&self) -> LifetimeSyntaxCategories<usize> {
        LifetimeSyntaxCategories {
            hidden: self.hidden.len(),
            elided: self.elided.len(),
            named: self.named.len(),
        }
    }

    pub fn iter_unnamed(&self) -> impl Iterator<Item = &T> {
        let Self { hidden, elided, named: _ } = self;
        std::iter::chain(hidden, elided)
    }
}

impl std::ops::Add for LifetimeSyntaxCategories<usize> {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self {
            hidden: self.hidden + rhs.hidden,
            elided: self.elided + rhs.elided,
            named: self.named + rhs.named,
        }
    }
}

fn lifetimes_use_matched_syntax(input_info: &[Info<'_>], output_info: &[Info<'_>]) -> bool {
    let (first, inputs) = input_info.split_first().unwrap();
    std::iter::chain(inputs, output_info).all(|info| info.syntax_category == first.syntax_category)
}

fn emit_mismatch_diagnostic<'tcx>(
    cx: &LateContext<'tcx>,
    input_info: &[Info<'_>],
    output_info: &[Info<'_>],
) {
    // There can only ever be zero or one bound lifetime
    // for a given lifetime resolution.
    let mut bound_lifetime = None;

    // We offer the following kinds of suggestions (when appropriate
    // such that the suggestion wouldn't violate the lint):
    //
    // 1. Every lifetime becomes named, when there is already a
    //    user-provided name.
    //
    // 2. A "mixed" signature, where references become implicit
    //    and paths become explicitly anonymous.
    //
    // 3. Every lifetime becomes implicit.
    //
    // 4. Every lifetime becomes explicitly anonymous.
    //
    // Number 2 is arguably the most common pattern and the one we
    // should push strongest. Number 3 is likely the next most common,
    // followed by number 1. Coming in at a distant last would be
    // number 4.
    //
    // Beyond these, there are variants of acceptable signatures that
    // we won't suggest because they are very low-value. For example,
    // we will never suggest `fn(&T1, &'_ T2) -> &T3` even though that
    // would pass the lint.
    //
    // The following collections are the lifetime instances that we
    // suggest changing to a given alternate style.

    // 1. Convert all to named.
    let mut suggest_change_to_explicit_bound = Vec::new();

    // 2. Convert to mixed. We track each kind of change separately.
    let mut suggest_change_to_mixed_implicit = Vec::new();
    let mut suggest_change_to_mixed_explicit_anonymous = Vec::new();

    // 3. Convert all to implicit.
    let mut suggest_change_to_implicit = Vec::new();

    // 4. Convert all to explicit anonymous.
    let mut suggest_change_to_explicit_anonymous = Vec::new();

    // Some styles prevent using implicit syntax at all.
    let mut allow_suggesting_implicit = true;

    // It only makes sense to suggest mixed if we have both sources.
    let mut saw_a_reference = false;
    let mut saw_a_path = false;

    for info in input_info.iter().chain(output_info) {
        use LifetimeSource::*;
        use hir::LifetimeSyntax::*;

        let lifetime = info.lifetime;

        if lifetime.syntax == ExplicitBound {
            bound_lifetime = Some(info);
        }

        match (lifetime.syntax, lifetime.source) {
            // E.g. `&T`.
            (Implicit, Reference) => {
                suggest_change_to_explicit_anonymous.push(info);
                suggest_change_to_explicit_bound.push(info);
            }

            // E.g. `&'_ T`.
            (ExplicitAnonymous, Reference) => {
                suggest_change_to_implicit.push(info);
                suggest_change_to_explicit_bound.push(info);
            }

            // E.g. `ContainsLifetime`.
            (Implicit, Path { .. }) => {
                suggest_change_to_mixed_explicit_anonymous.push(info);
                suggest_change_to_explicit_anonymous.push(info);
                suggest_change_to_explicit_bound.push(info);
            }

            // E.g. `ContainsLifetime<'_>`, `+ '_`, `+ use<'_>`.
            (ExplicitAnonymous, Path { .. } | OutlivesBound | PreciseCapturing) => {
                suggest_change_to_explicit_bound.push(info);
            }

            // E.g. `&'a T`.
            (ExplicitBound, Reference) => {
                suggest_change_to_implicit.push(info);
                suggest_change_to_mixed_implicit.push(info);
                suggest_change_to_explicit_anonymous.push(info);
            }

            // E.g. `ContainsLifetime<'a>`, `+ 'a`, `+ use<'a>`.
            (ExplicitBound, Path { .. } | OutlivesBound | PreciseCapturing) => {
                suggest_change_to_mixed_explicit_anonymous.push(info);
                suggest_change_to_explicit_anonymous.push(info);
            }

            (Implicit, OutlivesBound | PreciseCapturing) => {
                panic!("This syntax / source combination is not possible");
            }

            (_, Other) => {
                panic!("This syntax / source combination has already been skipped");
            }
        }

        if matches!(lifetime.source, Path { .. } | OutlivesBound | PreciseCapturing) {
            allow_suggesting_implicit = false;
        }

        match lifetime.source {
            Reference => saw_a_reference = true,
            Path { .. } => saw_a_path = true,
            _ => {}
        }
    }

    let categorize = |infos: &[Info<'_>]| {
        let mut categories = LifetimeSyntaxCategories::<Vec<_>>::default();
        for info in infos {
            categories.select(info.syntax_category).push(info.reporting_span());
        }
        categories
    };

    let inputs = categorize(input_info);
    let outputs = categorize(output_info);

    let make_implicit_suggestions =
        |infos: &[&Info<'_>]| infos.iter().map(|i| i.removing_span()).collect::<Vec<_>>();

    let explicit_bound_suggestion = bound_lifetime.map(|info| {
        build_mismatch_suggestion(info.lifetime.ident.as_str(), &suggest_change_to_explicit_bound)
    });

    tracing::debug!(?bound_lifetime, ?explicit_bound_suggestion);

    let should_suggest_mixed =
        // Do we have a mixed case?
        (saw_a_reference && saw_a_path) &&
        // Is there anything to change?
        (!suggest_change_to_mixed_implicit.is_empty() ||
         !suggest_change_to_mixed_explicit_anonymous.is_empty()) &&
        // If we have a named lifetime, prefer consistent naming.
        bound_lifetime.is_none();

    let mixed_suggestion = should_suggest_mixed.then(|| {
        let implicit_suggestions = make_implicit_suggestions(&suggest_change_to_mixed_implicit);

        let explicit_anonymous_suggestions = build_mismatch_suggestions_for_lifetime(
            "'_",
            &suggest_change_to_mixed_explicit_anonymous,
        );

        lints::MismatchedLifetimeSyntaxesSuggestion::Mixed {
            implicit_suggestions,
            explicit_anonymous_suggestions,
            optional_alternative: false,
        }
    });

    tracing::debug!(
        ?suggest_change_to_mixed_implicit,
        ?suggest_change_to_mixed_explicit_anonymous,
        ?mixed_suggestion,
    );

    let should_suggest_implicit =
        // Is there anything to change?
        !suggest_change_to_implicit.is_empty() &&
        // We never want to hide the lifetime in a path (or similar).
        allow_suggesting_implicit &&
        // If we have a named lifetime, prefer consistent naming.
        bound_lifetime.is_none();

    let implicit_suggestion = should_suggest_implicit.then(|| {
        let suggestions = make_implicit_suggestions(&suggest_change_to_implicit);

        lints::MismatchedLifetimeSyntaxesSuggestion::Implicit {
            suggestions,
            optional_alternative: false,
        }
    });

    tracing::debug!(
        ?should_suggest_implicit,
        ?suggest_change_to_implicit,
        allow_suggesting_implicit,
        ?implicit_suggestion,
    );

    let should_suggest_explicit_anonymous =
        // Is there anything to change?
        !suggest_change_to_explicit_anonymous.is_empty() &&
        // If we already have a mixed suggestion, avoid overlapping alternatives.
        mixed_suggestion.is_none() &&
        // If we have a named lifetime, prefer consistent naming.
        bound_lifetime.is_none();

    let explicit_anonymous_suggestion = should_suggest_explicit_anonymous
        .then(|| build_mismatch_suggestion("'_", &suggest_change_to_explicit_anonymous));

    tracing::debug!(
        ?should_suggest_explicit_anonymous,
        ?suggest_change_to_explicit_anonymous,
        ?explicit_anonymous_suggestion,
    );

    // We can produce a number of suggestions which may overwhelm
    // the user. Instead, we order the suggestions based on Rust
    // idioms. The "best" choice is shown to the user and the
    // remaining choices are shown to tools only.
    let mut suggestions = Vec::new();
    suggestions.extend(explicit_bound_suggestion);
    suggestions.extend(mixed_suggestion);
    suggestions.extend(implicit_suggestion);
    suggestions.extend(explicit_anonymous_suggestion);

    cx.emit_span_lint(
        MISMATCHED_LIFETIME_SYNTAXES,
        inputs.iter_unnamed().chain(outputs.iter_unnamed()).copied().collect::<Vec<_>>(),
        lints::MismatchedLifetimeSyntaxes { inputs, outputs, suggestions },
    );
}

fn build_mismatch_suggestion(
    lifetime_name: &str,
    infos: &[&Info<'_>],
) -> lints::MismatchedLifetimeSyntaxesSuggestion {
    let lifetime_name = lifetime_name.to_owned();

    let suggestions = build_mismatch_suggestions_for_lifetime(&lifetime_name, infos);

    lints::MismatchedLifetimeSyntaxesSuggestion::Explicit {
        lifetime_name,
        suggestions,
        optional_alternative: false,
    }
}

fn build_mismatch_suggestions_for_lifetime(
    lifetime_name: &str,
    infos: &[&Info<'_>],
) -> Vec<(Span, String)> {
    use hir::{AngleBrackets, LifetimeSource, LifetimeSyntax};

    #[derive(Clone, Copy, PartialEq, Eq, Hash)]
    enum PathSuggestionKind {
        Missing,
        Empty,
        Full,
    }

    let mut suggestions = Vec::new();
    let mut path_counts: FxIndexMap<(hir::HirId, PathSuggestionKind), (Span, usize)> =
        FxIndexMap::default();

    for info in infos {
        let lifetime = info.lifetime;
        if matches!(lifetime.syntax, LifetimeSyntax::Implicit) {
            if let LifetimeSource::Path { angle_brackets } = lifetime.source {
                let (span, kind) = match angle_brackets {
                    AngleBrackets::Missing => {
                        (lifetime.ident.span.shrink_to_hi(), PathSuggestionKind::Missing)
                    }
                    AngleBrackets::Empty => (lifetime.ident.span, PathSuggestionKind::Empty),
                    AngleBrackets::Full => (lifetime.ident.span, PathSuggestionKind::Full),
                };
                let entry = path_counts.entry((info.ty.hir_id, kind)).or_insert((span, 0));
                entry.1 += 1;
                continue;
            }
        }
        suggestions.push(info.suggestion(lifetime_name));
    }

    for ((_ty_hir_id, kind), (span, count)) in path_counts {
        let repeated = std::iter::repeat(lifetime_name).take(count).collect::<Vec<_>>().join(", ");

        let suggestion = match kind {
            PathSuggestionKind::Missing => format!("<{repeated}>"),
            PathSuggestionKind::Empty => repeated,
            PathSuggestionKind::Full => format!("{repeated}, "),
        };

        suggestions.push((span, suggestion));
    }

    suggestions
}

#[derive(Debug)]
struct Info<'tcx> {
    lifetime: &'tcx hir::Lifetime,
    syntax_category: LifetimeSyntaxCategory,
    ty: &'tcx hir::Ty<'tcx>,
}

impl<'tcx> Info<'tcx> {
    /// When reporting a lifetime that is implicit, we expand the span
    /// to include the type. Otherwise we end up pointing at nothing,
    /// which is a bit confusing.
    fn reporting_span(&self) -> Span {
        if self.lifetime.is_implicit() { self.ty.span } else { self.lifetime.ident.span }
    }

    /// When removing an explicit lifetime from a reference,
    /// we want to remove the whitespace after the lifetime.
    ///
    /// ```rust
    /// fn x(a: &'_ u8) {}
    /// ```
    ///
    /// Should become:
    ///
    /// ```rust
    /// fn x(a: &u8) {}
    /// ```
    // FIXME: Ideally, we'd also remove the lifetime declaration.
    fn removing_span(&self) -> Span {
        let mut span = self.lifetime.ident.span;
        if let hir::TyKind::Ref(_, mut_ty) = self.ty.kind {
            span = span.until(mut_ty.ty.span);
        }
        span
    }

    fn suggestion(&self, lifetime_name: &str) -> (Span, String) {
        self.lifetime.suggestion(lifetime_name)
    }
}

struct LifetimeInfoCollector<'tcx, F> {
    info_func: F,
    ty: &'tcx hir::Ty<'tcx>,
}

impl<'tcx, F> LifetimeInfoCollector<'tcx, F>
where
    F: FnMut(Info<'tcx>),
{
    fn collect(ty: &'tcx hir::Ty<'tcx>, info_func: F) {
        let mut this = Self { info_func, ty };

        intravisit::walk_unambig_ty(&mut this, ty);
    }
}

impl<'tcx, F> Visitor<'tcx> for LifetimeInfoCollector<'tcx, F>
where
    F: FnMut(Info<'tcx>),
{
    #[instrument(skip(self))]
    fn visit_lifetime(&mut self, lifetime: &'tcx hir::Lifetime) {
        if let Some(syntax_category) = LifetimeSyntaxCategory::new(lifetime) {
            let info = Info { lifetime, syntax_category, ty: self.ty };
            (self.info_func)(info);
        }
    }

    #[instrument(skip(self))]
    fn visit_ty(&mut self, ty: &'tcx hir::Ty<'tcx, hir::AmbigArg>) -> Self::Result {
        let old_ty = std::mem::replace(&mut self.ty, ty.as_unambig_ty());
        intravisit::walk_ty(self, ty);
        self.ty = old_ty;
    }
}