use std::collections::hash_map::Entry;
use std::fmt::Debug;
use std::mem;
use std::ops::{ControlFlow, Deref};

use rustc_data_structures::fx::{FxHashMap, FxHashSet};
use rustc_hir::CRATE_HIR_ID;
use rustc_hir::attrs::lang_items::LangItem;
use rustc_hir::def::Namespace;
use rustc_hir::def_id::{CRATE_DEF_ID, DefId, LOCAL_CRATE};
use rustc_infer::infer::canonical::query_response::make_query_region_constraints;
use rustc_infer::infer::canonical::{
    Canonical, CanonicalExt as _, CanonicalQueryInput, CanonicalVarKind, CanonicalVarValues,
    QueryRegionConstraint,
};
use rustc_infer::infer::{InferCtxt, RegionVariableOrigin, SubregionOrigin, TyCtxtInferExt};
use rustc_infer::traits::solve::{
    ComputeGoalFastPathOutcome, FetchEligibleAssocItemResponse, Goal, SucceededInErased,
};
use rustc_lint_defs::builtin::RECURSION_DEPTH_EXCEEDING_LIMIT;
use rustc_middle::traits::query::NoSolution;
use rustc_middle::traits::solve::{Certainty, MaybeInfo};
use rustc_middle::ty::print::{FmtPrinter, Print};
use rustc_middle::ty::{
    self, CanonicalizerState, MayBeErased, Ty, TyCtxt, TypeFlags, TypeFoldable, TypeSuperVisitable,
    TypeVisitable, TypeVisitableExt, TypeVisitor, TypingMode,
};
use rustc_next_trait_solver::solve::{GoalStalledOn, GoalStalledOnOpaques, TyOrConstInferVar};
use rustc_span::{DUMMY_SP, Span};
use rustc_structures::Limit;
use thin_vec::{ThinVec, thin_vec};

use super::inspect::InferCtxtProofTreeExt;
use crate::solve::inspect::{self, InspectConfig, ProofTreeVisitor};
use crate::traits::{EvaluateConstErr, ObligationCause, sizedness_fast_path, specialization_graph};

#[repr(transparent)]
pub struct SolverDelegate<'tcx>(InferCtxt<'tcx>);

impl<'a, 'tcx> From<&'a InferCtxt<'tcx>> for &'a SolverDelegate<'tcx> {
    fn from(infcx: &'a InferCtxt<'tcx>) -> Self {
        // SAFETY: `repr(transparent)`
        unsafe { std::mem::transmute(infcx) }
    }
}

impl<'tcx> Deref for SolverDelegate<'tcx> {
    type Target = InferCtxt<'tcx>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<'tcx> SolverDelegate<'tcx> {
    fn known_no_opaque_types_in_storage(&self) -> bool {
        self.inner.borrow_mut().opaque_types().is_empty()
            // in erased mode, observing that opaques are empty aren't enough to give a result
            // here, so let's try the slow path instead.
            && !self.typing_mode_raw().is_erased_not_coherence()
    }
}

/// Create a [`ComputeGoalFastPathOutcome`] signalling the goal is stalled
/// on a list of [`ty::GenericArg`]
fn goal_stalled_on_args<'tcx>(
    stalled_vars: ThinVec<TyOrConstInferVar>,
) -> ComputeGoalFastPathOutcome<'tcx> {
    ComputeGoalFastPathOutcome::TriviallyStalled {
        stalled_on: GoalStalledOn {
            stalled_vars,
            sub_roots: ThinVec::new(),
            stalled_maybe_info: MaybeInfo::AMBIGUOUS,
            opaques: GoalStalledOnOpaques::No,
        },
    }
}

/// Create a [`ComputeGoalFastPathOutcome`] signalling the  goal is stalled
/// on a list of [`ty::GenericArg`] *or* the opaque type storage being nonempty.
///
fn goal_stalled_on_args_or_nonempty_opaques<'tcx>(
    stalled_vars: ThinVec<TyOrConstInferVar>,
) -> ComputeGoalFastPathOutcome<'tcx> {
    ComputeGoalFastPathOutcome::TriviallyStalled {
        stalled_on: GoalStalledOn {
            stalled_vars,
            sub_roots: ThinVec::new(),
            stalled_maybe_info: MaybeInfo::AMBIGUOUS,
            opaques: GoalStalledOnOpaques::Yes {
                num_opaques_in_storage: 0,
                // This function should only be called when not in erased mode,
                // otherwise this is wrong. The `compute_goal_fast_path` does this
                // through `known_no_opaque_types_in_storage`
                previously_succeeded_in_erased: SucceededInErased::No,
            },
        },
    }
}

struct CollectNonRegionInfer<'tcx> {
    infers: ThinVec<ty::GenericArg<'tcx>>,
    visited: FxHashSet<Ty<'tcx>>,
}

impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for CollectNonRegionInfer<'tcx> {
    fn visit_ty(&mut self, ty: Ty<'tcx>) {
        if self.visited.contains(&ty) {
            return;
        }

        match ty.kind() {
            ty::Infer(_) => self.infers.push(ty.into()),
            _ => ty.super_visit_with(self),
        }

        self.visited.insert(ty);
    }

    fn visit_const(&mut self, ct: ty::Const<'tcx>) {
        match ct.kind() {
            ty::ConstKind::Infer(_) => self.infers.push(ct.into()),
            _ => ct.super_visit_with(self),
        }
    }
}

impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<'tcx> {
    type Infcx = InferCtxt<'tcx>;
    type Interner = TyCtxt<'tcx>;

    fn cx(&self) -> TyCtxt<'tcx> {
        self.0.tcx
    }

    fn build_with_canonical<V>(
        interner: TyCtxt<'tcx>,
        canonical: &CanonicalQueryInput<'tcx, V>,
    ) -> (Self, V, CanonicalVarValues<'tcx>)
    where
        V: TypeFoldable<TyCtxt<'tcx>>,
    {
        let (infcx, value, vars) = interner
            .infer_ctxt()
            .with_next_trait_solver(true)
            .build_with_canonical(DUMMY_SP, canonical);
        (SolverDelegate(infcx), value, vars)
    }

    fn compute_goal_fast_path(
        &self,
        goal: Goal<'tcx, ty::Predicate<'tcx>>,
        span: Span,
    ) -> ComputeGoalFastPathOutcome<'tcx> {
        use ComputeGoalFastPathOutcome as Outcome;

        // FIXME(-Zassumptions-on-binders): actually handle fast path
        if self.tcx.assumptions_on_binders() {
            return Outcome::NoFastPath;
        }

        let pred = goal.predicate.kind();
        match pred.skip_binder() {
            ty::PredicateKind::Clause(ty::ClauseKind::Trait(trait_pred)) => {
                let trait_pred = pred.rebind(trait_pred);

                let self_ty = self.shallow_resolve(trait_pred.self_ty().skip_binder());
                if let Some(vid) = self_ty.ty_vid()
                // We don't do this fast path when opaques are defined since we may
                // eventually use opaques to incompletely guide inference via ty var
                // self types.
                // FIXME: Properly consider opaques here.
                && self.known_no_opaque_types_in_storage()
                {
                    goal_stalled_on_args_or_nonempty_opaques(thin_vec![TyOrConstInferVar::Ty(vid)])
                } else if trait_pred.polarity() == ty::ClausePolarity::Positive {
                    match self.0.tcx.as_lang_item(trait_pred.def_id()) {
                        Some(LangItem::Sized) | Some(LangItem::MetaSized) => {
                            let predicate = self.resolve_vars_if_possible(goal.predicate);
                            if sizedness_fast_path(self.tcx, predicate, goal.param_env) {
                                Outcome::TriviallyHolds
                            } else {
                                Outcome::NoFastPath
                            }
                        }
                        Some(LangItem::Copy | LangItem::Clone) => {
                            let self_ty =
                                self.resolve_vars_if_possible(trait_pred.self_ty().skip_binder());
                            // Unlike `Sized` traits, which always prefer the built-in impl,
                            // `Copy`/`Clone` may be shadowed by a param-env candidate which
                            // could force a lifetime error or guide inference. While that's
                            // not generally desirable, it is observable, so for now let's
                            // ignore this fast path for types that have regions or infer.
                            if !self_ty
                                .has_type_flags(TypeFlags::HAS_FREE_REGIONS | TypeFlags::HAS_INFER)
                                && self_ty.is_trivially_pure_clone_copy()
                            {
                                Outcome::TriviallyHolds
                            } else {
                                Outcome::NoFastPath
                            }
                        }
                        _ => Outcome::NoFastPath,
                    }
                } else {
                    Outcome::NoFastPath
                }
            }
            ty::PredicateKind::DynCompatible(def_id) if self.0.tcx.is_dyn_compatible(def_id) => {
                Outcome::TriviallyHolds
            }
            ty::PredicateKind::Clause(ty::ClauseKind::RegionOutlives(outlives)) => {
                if outlives.has_escaping_bound_vars() {
                    return Outcome::NoFastPath;
                }

                self.0.sub_regions(
                    SubregionOrigin::RelateRegionParamBound(span, None),
                    outlives.1,
                    outlives.0,
                    ty::VisibleForLeakCheck::Yes,
                );
                Outcome::TriviallyHolds
            }
            ty::PredicateKind::Clause(ty::ClauseKind::TypeOutlives(outlives)) => {
                if outlives.has_escaping_bound_vars() {
                    return Outcome::NoFastPath;
                }

                let ty = self.resolve_vars_if_possible(outlives.0);
                let mut infer_collector = CollectNonRegionInfer {
                    infers: Default::default(),
                    visited: Default::default(),
                };
                ty.visit_with(&mut infer_collector);
                let infers = infer_collector.infers;
                if !infers.is_empty() {
                    return goal_stalled_on_args(
                        infers
                            .into_iter()
                            .map(|i| {
                                TyOrConstInferVar::maybe_from_generic_arg::<Self::Interner>(i)
                                    .unwrap()
                            })
                            .collect(),
                    );
                }

                if ty.has_non_rigid_aliases() {
                    return Outcome::NoFastPath;
                }

                self.0.register_type_outlives_constraint(
                    outlives.0,
                    outlives.1,
                    &ObligationCause::dummy_with_span(span),
                );

                Outcome::TriviallyHolds
            }
            ty::PredicateKind::Subtype(ty::SubtypePredicate { a, b, .. })
            | ty::PredicateKind::Coerce(ty::CoercePredicate { a, b }) => {
                if a.has_escaping_bound_vars() || b.has_escaping_bound_vars() {
                    return Outcome::NoFastPath;
                }

                match (self.shallow_resolve(a).kind(), self.shallow_resolve(b).kind()) {
                    (&ty::Infer(ty::TyVar(a_vid)), &ty::Infer(ty::TyVar(b_vid))) => {
                        self.sub_unify_ty_vids_raw(a_vid, b_vid);
                        goal_stalled_on_args(thin_vec![
                            TyOrConstInferVar::Ty(a_vid),
                            TyOrConstInferVar::Ty(b_vid),
                        ])
                    }
                    _ => Outcome::NoFastPath,
                }
            }
            ty::PredicateKind::Clause(ty::ClauseKind::ConstArgHasType(ct, _)) => {
                if ct.has_escaping_bound_vars() {
                    return Outcome::NoFastPath;
                }

                let arg = self.shallow_resolve_const(ct);
                if let Some(vid) = arg.ct_vid() {
                    goal_stalled_on_args(thin_vec![TyOrConstInferVar::Const(vid)])
                } else {
                    Outcome::NoFastPath
                }
            }
            ty::PredicateKind::Clause(ty::ClauseKind::WellFormed(arg)) => {
                if arg.has_escaping_bound_vars() {
                    return Outcome::NoFastPath;
                }

                let arg = self.shallow_resolve_term(arg);
                if arg.is_trivially_wf(self.tcx) {
                    Outcome::TriviallyHolds
                } else if arg.is_infer() {
                    goal_stalled_on_args(thin_vec![
                        TyOrConstInferVar::maybe_from_term::<TyCtxt<'tcx>>(arg)
                            .expect("its an infer var"),
                    ])
                } else {
                    Outcome::NoFastPath
                }
            }
            _ => Outcome::NoFastPath,
        }
    }

    fn fresh_var_for_kind(
        &self,
        arg: ty::GenericArg<'tcx>,
        span: Span,
        universe: ty::UniverseIndex,
    ) -> ty::GenericArg<'tcx> {
        match arg.kind() {
            ty::GenericArgKind::Lifetime(_) => {
                self.next_region_var_in_universe(RegionVariableOrigin::Misc(span), universe).into()
            }
            ty::GenericArgKind::Type(_) => self.next_ty_var_in_universe(span, universe).into(),
            ty::GenericArgKind::Const(_) => self.next_const_var_in_universe(span, universe).into(),
        }
    }

    fn leak_check(&self, max_input_universe: ty::UniverseIndex) -> Result<(), NoSolution> {
        self.0.leak_check(max_input_universe, None).map_err(|_| NoSolution)
    }

    fn evaluate_const<E: Debug>(
        &self,
        param_env: ty::ParamEnv<'tcx>,
        alias_const: ty::AliasConst<'tcx>,
        normalize_ty: impl FnOnce(ty::Unnormalized<'tcx, Ty<'tcx>>) -> Result<Ty<'tcx>, E>,
    ) -> Result<Option<ty::Const<'tcx>>, E> {
        let ct = ty::Const::new_alias(self.tcx, ty::IsRigid::No, alias_const);

        match crate::traits::try_evaluate_const(&self.0, ct, param_env, normalize_ty) {
            Ok(ct) => Ok(Some(ct)),
            Err(EvaluateConstErr::EvaluationFailure(e)) => {
                Ok(Some(ty::Const::new_error(self.tcx, e)))
            }
            Err(
                EvaluateConstErr::InvalidConstParamTy(_) | EvaluateConstErr::HasGenericsOrInfers,
            ) => Ok(None),
            Err(EvaluateConstErr::FailedNormalization(e)) => Err(e),
        }
    }

    fn well_formed_goals(
        &self,
        param_env: ty::ParamEnv<'tcx>,
        term: ty::Term<'tcx>,
    ) -> Option<Vec<Goal<'tcx, ty::Predicate<'tcx>>>> {
        crate::traits::wf::unnormalized_obligations(
            &self.0,
            param_env,
            term,
            DUMMY_SP,
            CRATE_DEF_ID,
        )
        .map(|obligations| obligations.into_iter().map(|obligation| obligation.as_goal()).collect())
    }

    fn make_deduplicated_region_constraints(
        &self,
    ) -> Vec<(ty::RegionConstraint<'tcx>, ty::VisibleForLeakCheck)> {
        // Cannot use `take_registered_region_obligations` as we may compute the response
        // inside of a `probe` whenever we have multiple choices inside of the solver.
        let region_obligations = self.0.inner.borrow().region_obligations().to_owned();
        let region_assumptions = self.0.inner.borrow().region_assumptions().to_owned();
        let region_constraints = self.0.with_region_constraints(|region_constraints| {
            make_query_region_constraints(
                region_obligations,
                region_constraints,
                region_assumptions,
            )
        });

        let mut seen = FxHashMap::default();
        let mut constraints = vec![];
        for QueryRegionConstraint { constraint: outlives, visible_for_leak_check: vis, .. } in
            region_constraints.constraints
        {
            match seen.entry(outlives) {
                Entry::Occupied(occupied) => {
                    let idx = occupied.get();
                    let (_, prev_vis): &mut (_, ty::VisibleForLeakCheck) =
                        constraints.get_mut(*idx).unwrap();
                    *prev_vis = (*prev_vis).or(vis);
                }
                Entry::Vacant(vacant) => {
                    vacant.insert(constraints.len());
                    constraints.push((outlives, vis));
                }
            }
        }
        constraints
    }

    fn instantiate_canonical<V>(
        &self,
        canonical: Canonical<'tcx, V>,
        values: CanonicalVarValues<'tcx>,
    ) -> V
    where
        V: TypeFoldable<TyCtxt<'tcx>>,
    {
        canonical.instantiate(self.tcx, &values)
    }

    fn instantiate_canonical_var(
        &self,
        kind: CanonicalVarKind<'tcx>,
        span: Span,
        var_values: &[ty::GenericArg<'tcx>],
        universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
    ) -> ty::GenericArg<'tcx> {
        self.0.instantiate_canonical_var(span, kind, var_values, universe_map)
    }

    fn add_item_bounds_for_hidden_type(
        &self,
        def_id: DefId,
        args: ty::GenericArgsRef<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
        hidden_ty: Ty<'tcx>,
        goals: &mut Vec<Goal<'tcx, ty::Predicate<'tcx>>>,
    ) {
        self.0.add_item_bounds_for_hidden_type(def_id, args, param_env, hidden_ty, goals);
    }

    fn fetch_eligible_assoc_item(
        &self,
        goal_trait_ref: ty::TraitRef<'tcx>,
        trait_assoc_def_id: DefId,
        impl_def_id: DefId,
    ) -> FetchEligibleAssocItemResponse<'tcx> {
        let node_item =
            match specialization_graph::assoc_def(self.tcx, impl_def_id, trait_assoc_def_id) {
                Ok(i) => i,
                Err(guar) => return FetchEligibleAssocItemResponse::Err(guar),
            };

        let typing_mode = self.typing_mode_raw();

        let eligible = if node_item.is_final() {
            // Non-specializable items are always projectable.
            true
        } else {
            // Only reveal a specializable default if we're past type-checking
            // and the obligation is monomorphic, otherwise passes such as
            // transmute checking and polymorphic MIR optimizations could
            // get a result which isn't correct for all monomorphizations.
            match typing_mode {
                TypingMode::Coherence
                | TypingMode::Typeck { .. }
                | TypingMode::PostTypeckUntilBorrowck { .. }
                | TypingMode::Reflection
                | TypingMode::PostBorrowck { .. } => false,
                TypingMode::PostAnalysis | TypingMode::Codegen => {
                    let poly_trait_ref = self.resolve_vars_if_possible(goal_trait_ref);
                    !poly_trait_ref.still_further_specializable()
                }
                TypingMode::ErasedNotCoherence(MayBeErased) => {
                    return FetchEligibleAssocItemResponse::NotFoundBecauseErased;
                }
            }
        };

        // FIXME: Check for defaultness here may cause diagnostics problems.
        if eligible {
            FetchEligibleAssocItemResponse::Found(node_item.item.def_id)
        } else {
            // We know it's not erased since then we'd have returned in the match above,
            // or node_item.final() was true and eligible is always true.
            FetchEligibleAssocItemResponse::NotFound(typing_mode.assert_not_erased())
        }
    }

    // FIXME: This actually should destructure the `Result` we get from transmutability and
    // register candidates. We probably need to register >1 since we may have an OR of ANDs.
    fn is_transmutable(
        &self,
        src: Ty<'tcx>,
        dst: Ty<'tcx>,
        assume: ty::Const<'tcx>,
    ) -> Result<Certainty, NoSolution> {
        // Erase regions because we compute layouts in `rustc_transmute`,
        // which will ICE for region vars.
        let (dst, src) = self.tcx.erase_and_anonymize_regions((dst, src));

        let Some(assume) = rustc_transmute::Assume::from_const(self.tcx, assume) else {
            return Err(NoSolution);
        };

        // FIXME(transmutability): This really should be returning nested goals for `Answer::If*`
        match rustc_transmute::TransmuteTypeEnv::new(self.0.tcx).is_transmutable(src, dst, assume) {
            rustc_transmute::Answer::Yes => Ok(Certainty::Yes),
            rustc_transmute::Answer::No(_) | rustc_transmute::Answer::If(_) => Err(NoSolution),
        }
    }

    fn obtain_canonicalizer_state(&self) -> CanonicalizerState<Self::Interner> {
        // We temporarily take the canonicalizer state.
        mem::take(&mut self.canonicalizer_state.borrow_mut())
    }

    fn release_canonicalizer_state(&self, mut state: CanonicalizerState<Self::Interner>) {
        // Clear (don't deallocate) the state for later reuse.
        state.clear();
        *self.canonicalizer_state.borrow_mut() = state;
    }

    fn emit_next_solver_overflow_fcw(&self, goal: Goal<'tcx, ty::Predicate<'tcx>>, span: Span) {
        let tcx = self.tcx;
        let goal = self.resolve_vars_if_possible(goal);
        let mut visitor = OverflowedGoalChain {
            span,
            predicates: vec![],
            recursion_limit: usize::min(16, tcx.recursion_limit().0),
        };
        let _ = self
            .with_disabled_next_solver_overflow_fcw(|| self.visit_proof_tree(goal, &mut visitor));
        tcx.emit_node_span_lint(
            RECURSION_DEPTH_EXCEEDING_LIMIT,
            CRATE_HIR_ID,
            span,
            rustc_errors::DiagDecorator(|diag| {
                // FIXME: share this with overflow error in fulfillment instead of duplicating.
                let pred_str = |pred: ty::Predicate<'tcx>| {
                    let s = pred.to_string();
                    if s.len() > 80 {
                        let mut p: FmtPrinter<'_, '_> =
                            FmtPrinter::new_with_limit(tcx, Namespace::TypeNS, Limit(10));
                        pred.print(&mut p).unwrap();
                        p.into_buffer()
                    } else {
                        s
                    }
                };
                diag.primary_message(format!(
                    "overflow evaluating the requirement `{}`",
                    pred_str(goal.predicate),
                ));
                for p in visitor.predicates.into_iter().skip(1) {
                    diag.note(format!("which requires `{}`", pred_str(p)));
                }
                diag.note("and so on...");
                diag.help(
                    "consider adding a manual `impl` of auto traits like `Send` for intermediate types, if auto traits are involved",
                );
                diag.help(format!(
                    "or consider increasing the recursion limit by adding a \
                     `#![recursion_limit = \"{}\"]` attribute to your crate (`{}`)",
                    tcx.recursion_limit() * 2,
                    tcx.crate_name(LOCAL_CRATE),
                ));
                diag.note("this lint is attached to the whole crate and can't be disabled on a per-function basis");
            }),
        )
    }
}

struct OverflowedGoalChain<'tcx> {
    span: Span,
    predicates: Vec<ty::Predicate<'tcx>>,
    recursion_limit: usize,
}

impl<'tcx> ProofTreeVisitor<'tcx> for OverflowedGoalChain<'tcx> {
    type Result = ControlFlow<()>;

    fn span(&self) -> Span {
        self.span
    }

    fn config(&self) -> InspectConfig {
        InspectConfig { max_depth: self.recursion_limit }
    }

    fn visit_goal(&mut self, goal: &inspect::InspectGoal<'_, 'tcx>) -> Self::Result {
        self.predicates.push(goal.goal().predicate);
        if let Some(cand) = goal.unique_applicable_candidate() {
            goal.infcx().probe(|_| {
                if let Some(nested_goal_with_largest_required_depth) = cand
                    .instantiate_nested_goals(self.span)
                    .into_iter()
                    .max_by_key(|g| g.required_depth())
                {
                    nested_goal_with_largest_required_depth.visit_with(self)
                } else {
                    ControlFlow::Continue(())
                }
            })?;
        }
        ControlFlow::Continue(())
    }

    fn on_recursion_limit(&mut self) -> Self::Result {
        ControlFlow::Break(())
    }
}