use rustc_hash::{FxHashMap as HashMap, FxHashSet as HashSet};
use std::sync::Arc;
use std::sync::atomic::{AtomicU32, Ordering};

use diagnostics::{PatternIssue, UnusedExpressionKind};
use syntax::ast::{BindingId, BindingKind, DeadCodeCause, Span};
use syntax::program::TestFunction;
use syntax::types::Type;

#[derive(Debug, Default)]
pub struct BindingIdAllocator {
    next: AtomicU32,
}

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

    pub fn reserve(&self) -> BindingId {
        self.next.fetch_add(1, Ordering::Relaxed)
    }

    pub fn snapshot(&self) -> BindingId {
        self.next.load(Ordering::Relaxed)
    }
}

#[derive(Debug)]
pub struct Facts {
    allocator: Arc<BindingIdAllocator>,

    // LSP-consumed; reshaping these affects crates/lsp/.
    pub bindings: HashMap<BindingId, BindingFact>,
    pub usages: Vec<Usage>,
    usage_set: HashSet<(Span, Span)>,

    // Lint-support facts: read by reference by passes::lints (mostly
    // from_facts; interface_satisfied_methods by ref_graph).
    pub dead_code: Vec<DeadCodeFact>,
    pub pattern_issues: Vec<PatternIssue>,
    pub unused_expressions: Vec<UnusedExpressionFact>,
    pub discarded_tail_expressions: Vec<DiscardedTailFact>,
    pub overused_references: Vec<OverusedReferenceFact>,
    pub unused_type_params: Vec<UnusedTypeParamFact>,
    pub type_params_only_in_bound: Vec<TypeParamOnlyInBoundFact>,
    pub always_failing_try_blocks: Vec<Span>,
    pub expression_only_fstrings: Vec<Span>,
    pub interface_satisfied_methods: HashMap<(String, String), Vec<String>>,
    pub equality_derivations: Vec<String>,
    pub test_functions: Vec<TestFunction>,

    // Drained by passes::deferred via mem::take.
    pub generic_call_checks: Vec<GenericCallCheck>,
    pub empty_collection_checks: Vec<EmptyCollectionCheck>,
    pub statement_tail_checks: Vec<StatementTailCheck>,

    /// Value-position `match`/`select` arms that did not reconcile, drained and
    /// checked against the use-site result type at the end of `infer_file`.
    pub branch_subsumptions: Vec<BranchSubsumption>,

    /// Value-position selects with one shorthand receive and no default,
    /// checked for exhaustiveness once the result type is pinned.
    pub select_exhaustiveness_checks: Vec<SelectExhaustivenessCheck>,

    /// Suppresses contradictory lints from or-patterns whose binding sets disagree.
    pub or_pattern_error_spans: HashSet<Span>,

    /// Spans of binary expressions the checker rejected, so lints can skip them.
    pub type_error_spans: HashSet<Span>,

    /// Resolved type for each generic-bound annotation, keyed by the
    /// annotation's span. Lets emit render bounds from the resolved type
    /// instead of re-resolving the annotation.
    pub bound_types: HashMap<Span, Type>,
}

#[derive(Debug, Clone)]
pub struct GenericCallCheck {
    /// A type that must be fully resolved once inference finishes; if it still has
    /// unbound type variables, the call's generic parameter couldn't be inferred.
    pub ty: Type,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct EmptyCollectionCheck {
    pub name: String,
    pub ty: Type,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct StatementTailCheck {
    pub expected_ty: Type,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct BranchSubsumption {
    pub result_ty: Type,
    pub arms: Vec<(Type, Span)>,
}

#[derive(Debug, Clone)]
pub struct SelectExhaustivenessCheck {
    pub result_ty: Type,
    pub span: Span,
}

impl Facts {
    pub fn new(allocator: Arc<BindingIdAllocator>) -> Self {
        Self {
            allocator,
            bindings: HashMap::default(),
            dead_code: Vec::new(),
            pattern_issues: Vec::new(),
            unused_expressions: Vec::new(),
            discarded_tail_expressions: Vec::new(),
            overused_references: Vec::new(),
            unused_type_params: Vec::new(),
            type_params_only_in_bound: Vec::new(),
            always_failing_try_blocks: Vec::new(),
            expression_only_fstrings: Vec::new(),
            generic_call_checks: Vec::new(),
            empty_collection_checks: Vec::new(),
            statement_tail_checks: Vec::new(),
            branch_subsumptions: Vec::new(),
            select_exhaustiveness_checks: Vec::new(),
            or_pattern_error_spans: HashSet::default(),
            type_error_spans: HashSet::default(),
            usages: Vec::new(),
            usage_set: HashSet::default(),
            interface_satisfied_methods: HashMap::default(),
            equality_derivations: Vec::new(),
            test_functions: Vec::new(),
            bound_types: HashMap::default(),
        }
    }

    pub fn add_binding(
        &mut self,
        name: String,
        span: Span,
        kind: BindingKind,
        is_typedef: bool,
        is_struct_field: bool,
        is_as_alias: bool,
    ) -> BindingId {
        let id = self.allocator.reserve();
        self.bindings.insert(
            id,
            BindingFact {
                name,
                span,
                kind,
                used: false,
                mutated: false,
                is_typedef,
                is_struct_field,
                is_as_alias,
            },
        );
        id
    }

    pub fn mark_used(&mut self, id: BindingId) {
        if let Some(fact) = self.bindings.get_mut(&id) {
            fact.used = true;
        }
    }

    pub fn mark_mutated(&mut self, id: BindingId) {
        if let Some(fact) = self.bindings.get_mut(&id) {
            fact.mutated = true;
        }
    }

    pub fn binding_checkpoint(&self) -> BindingId {
        self.allocator.snapshot()
    }

    pub fn remove_bindings_from(&mut self, checkpoint: BindingId) {
        self.bindings.retain(|id, _| *id < checkpoint);
    }

    pub fn add_dead_code(&mut self, span: Span, cause: DeadCodeCause) {
        self.dead_code.push(DeadCodeFact { span, cause });
    }

    pub fn add_overused_reference(&mut self, span: Span, name: Option<String>) {
        self.overused_references
            .push(OverusedReferenceFact { span, name });
    }

    pub fn add_always_failing_try_block(&mut self, span: Span) {
        self.always_failing_try_blocks.push(span);
    }

    pub fn add_expression_only_fstring(&mut self, span: Span) {
        self.expression_only_fstrings.push(span);
    }

    pub fn add_usage(&mut self, usage_span: Span, definition_span: Span) {
        if self.usage_set.insert((usage_span, definition_span)) {
            self.usages.push(Usage {
                usage_span,
                definition_span,
            });
        }
    }

    pub fn mark_method_used_for_interface(
        &mut self,
        module_id: String,
        method_name: String,
        impl_type_name: String,
    ) {
        self.interface_satisfied_methods
            .entry((module_id, method_name))
            .or_default()
            .push(impl_type_name);
    }

    /// Whether `type_name`'s `method_name` was found to satisfy an interface method,
    /// so the naming lint keeps its conformance-required spelling.
    pub fn method_satisfies_interface(
        &self,
        module_id: &str,
        method_name: &str,
        type_name: &str,
    ) -> bool {
        self.interface_satisfied_methods
            .get(&(module_id.to_string(), method_name.to_string()))
            .is_some_and(|types| types.iter().any(|t| t == type_name))
    }

    pub fn absorb_local_facts(&mut self, local: LocalFacts) {
        let LocalFacts {
            unused_expressions,
            discarded_tail_expressions,
            unused_type_params,
            type_params_only_in_bound,
        } = local;
        self.unused_expressions.extend(unused_expressions);
        self.discarded_tail_expressions
            .extend(discarded_tail_expressions);
        self.unused_type_params.extend(unused_type_params);
        self.type_params_only_in_bound
            .extend(type_params_only_in_bound);
    }

    pub fn merge(&mut self, other: Facts) {
        debug_assert!(
            Arc::ptr_eq(&self.allocator, &other.allocator),
            "Facts::merge requires a shared BindingIdAllocator",
        );

        let Facts {
            allocator: _,
            bindings,
            dead_code,
            pattern_issues,
            unused_expressions,
            discarded_tail_expressions,
            overused_references,
            unused_type_params,
            type_params_only_in_bound,
            always_failing_try_blocks,
            expression_only_fstrings,
            generic_call_checks,
            empty_collection_checks,
            statement_tail_checks,
            branch_subsumptions,
            select_exhaustiveness_checks,
            or_pattern_error_spans,
            type_error_spans,
            usages,
            usage_set: _,
            interface_satisfied_methods,
            equality_derivations,
            test_functions,
            bound_types,
        } = other;
        self.equality_derivations.extend(equality_derivations);
        self.test_functions.extend(test_functions);
        self.bound_types.extend(bound_types);

        self.bindings.extend(bindings);
        self.dead_code.extend(dead_code);
        self.pattern_issues.extend(pattern_issues);
        self.unused_expressions.extend(unused_expressions);
        self.discarded_tail_expressions
            .extend(discarded_tail_expressions);
        self.overused_references.extend(overused_references);
        self.unused_type_params.extend(unused_type_params);
        self.type_params_only_in_bound
            .extend(type_params_only_in_bound);
        self.always_failing_try_blocks
            .extend(always_failing_try_blocks);
        self.expression_only_fstrings
            .extend(expression_only_fstrings);
        self.generic_call_checks.extend(generic_call_checks);
        self.empty_collection_checks.extend(empty_collection_checks);
        self.statement_tail_checks.extend(statement_tail_checks);
        self.branch_subsumptions.extend(branch_subsumptions);
        self.select_exhaustiveness_checks
            .extend(select_exhaustiveness_checks);
        self.or_pattern_error_spans.extend(or_pattern_error_spans);
        self.type_error_spans.extend(type_error_spans);

        self.usages.reserve(usages.len());
        self.usage_set.reserve(usages.len());
        for Usage {
            usage_span,
            definition_span,
        } in usages
        {
            self.add_usage(usage_span, definition_span);
        }

        for (key, impl_types) in interface_satisfied_methods {
            self.interface_satisfied_methods
                .entry(key)
                .or_default()
                .extend(impl_types);
        }
    }
}

#[derive(Debug, Default)]
pub struct LocalFacts {
    pub unused_expressions: Vec<UnusedExpressionFact>,
    pub discarded_tail_expressions: Vec<DiscardedTailFact>,
    pub unused_type_params: Vec<UnusedTypeParamFact>,
    pub type_params_only_in_bound: Vec<TypeParamOnlyInBoundFact>,
}

impl LocalFacts {
    pub fn merge(&mut self, other: LocalFacts) {
        let LocalFacts {
            unused_expressions,
            discarded_tail_expressions,
            unused_type_params,
            type_params_only_in_bound,
        } = other;
        self.unused_expressions.extend(unused_expressions);
        self.discarded_tail_expressions
            .extend(discarded_tail_expressions);
        self.unused_type_params.extend(unused_type_params);
        self.type_params_only_in_bound
            .extend(type_params_only_in_bound);
    }

    pub fn add_unused_expression(&mut self, span: Span, kind: UnusedExpressionKind) {
        self.unused_expressions
            .push(UnusedExpressionFact { span, kind });
    }

    pub fn add_discarded_tail(
        &mut self,
        span: Span,
        return_type: String,
        expected_span: Span,
        expected_type: String,
    ) {
        self.discarded_tail_expressions.push(DiscardedTailFact {
            span,
            return_type,
            expected_span,
            expected_type,
        });
    }

    pub fn add_unused_type_param(&mut self, name: String, span: Span) {
        self.unused_type_params
            .push(UnusedTypeParamFact { name, span });
    }

    pub fn add_type_param_only_in_bound(&mut self, name: String, span: Span) {
        self.type_params_only_in_bound
            .push(TypeParamOnlyInBoundFact { name, span });
    }
}

#[derive(Debug, Clone)]
pub struct BindingFact {
    pub name: String,
    pub span: Span,
    pub kind: BindingKind,
    pub used: bool,
    pub mutated: bool,
    pub is_typedef: bool,
    /// If true, this binding is a shorthand in a struct pattern (e.g., `Point { x }`)
    pub is_struct_field: bool,
    /// If true, this binding was introduced by an `as` alias (e.g., `Point { .. } as p`)
    pub is_as_alias: bool,
}

#[derive(Debug, Clone)]
pub struct DeadCodeFact {
    pub span: Span,
    pub cause: DeadCodeCause,
}

#[derive(Debug, Clone)]
pub struct UnusedExpressionFact {
    pub span: Span,
    pub kind: UnusedExpressionKind,
}

#[derive(Debug, Clone)]
pub struct DiscardedTailFact {
    pub span: Span,
    pub return_type: String,
    pub expected_span: Span,
    pub expected_type: String,
}

#[derive(Debug, Clone)]
pub struct OverusedReferenceFact {
    pub span: Span,
    pub name: Option<String>,
}

#[derive(Debug, Clone)]
pub struct UnusedTypeParamFact {
    pub name: String,
    pub span: Span,
}

#[derive(Debug, Clone)]
pub struct TypeParamOnlyInBoundFact {
    pub name: String,
    pub span: Span,
}

/// Records a usage of a symbol, linking the usage location to its definition.
/// Used by LSP for find-references.
#[derive(Debug, Clone)]
pub struct Usage {
    pub usage_span: Span,
    pub definition_span: Span,
}

#[cfg(test)]
mod tests {
    use super::*;
    use syntax::ast::BindingKind;

    fn span(offset: u32) -> Span {
        Span::new(0, offset, 1)
    }

    #[test]
    fn merge_preserves_unique_binding_ids_across_tasks() {
        let allocator = Arc::new(BindingIdAllocator::new());
        let mut a = Facts::new(allocator.clone());
        let mut b = Facts::new(allocator.clone());

        let a_id = a.add_binding(
            "a".into(),
            span(0),
            BindingKind::Let { mutable: false },
            false,
            false,
            false,
        );
        let b_id = b.add_binding(
            "b".into(),
            span(1),
            BindingKind::Let { mutable: false },
            false,
            false,
            false,
        );
        assert_ne!(a_id, b_id);

        a.merge(b);
        assert_eq!(a.bindings.len(), 2);
        assert!(a.bindings.contains_key(&a_id));
        assert!(a.bindings.contains_key(&b_id));
    }

    #[test]
    fn merge_extends_vec_facts() {
        let allocator = Arc::new(BindingIdAllocator::new());
        let mut a = Facts::new(allocator.clone());
        let mut b = Facts::new(allocator);

        a.add_always_failing_try_block(span(0));
        b.add_always_failing_try_block(span(1));
        b.add_always_failing_try_block(span(2));

        a.merge(b);
        assert_eq!(a.always_failing_try_blocks.len(), 3);
    }

    #[test]
    fn merge_deduplicates_usages() {
        let allocator = Arc::new(BindingIdAllocator::new());
        let mut a = Facts::new(allocator.clone());
        let mut b = Facts::new(allocator);

        a.add_usage(span(10), span(0));
        b.add_usage(span(10), span(0));
        b.add_usage(span(20), span(0));

        a.merge(b);
        assert_eq!(a.usages.len(), 2);
    }

    #[test]
    fn merge_deduplicates_or_pattern_error_spans() {
        let allocator = Arc::new(BindingIdAllocator::new());
        let mut a = Facts::new(allocator.clone());
        let mut b = Facts::new(allocator);

        a.or_pattern_error_spans.insert(span(0));
        b.or_pattern_error_spans.insert(span(0));
        b.or_pattern_error_spans.insert(span(1));

        a.merge(b);
        assert_eq!(a.or_pattern_error_spans.len(), 2);
    }

    #[test]
    fn absorb_local_facts_extends_all_four_streams() {
        let allocator = Arc::new(BindingIdAllocator::new());
        let mut facts = Facts::new(allocator);

        let mut local = LocalFacts::default();
        local.add_unused_expression(span(0), UnusedExpressionKind::Value);
        local.add_discarded_tail(span(1), "Int".into(), span(2), "Unit".into());
        local.add_unused_type_param("T".into(), span(3));
        local.add_type_param_only_in_bound("U".into(), span(4));

        facts.absorb_local_facts(local);

        assert_eq!(facts.unused_expressions.len(), 1);
        assert_eq!(facts.discarded_tail_expressions.len(), 1);
        assert_eq!(facts.unused_type_params.len(), 1);
        assert_eq!(facts.type_params_only_in_bound.len(), 1);
    }

    #[test]
    fn merge_concatenates_interface_method_impl_types() {
        let allocator = Arc::new(BindingIdAllocator::new());
        let mut a = Facts::new(allocator.clone());
        let mut b = Facts::new(allocator);

        a.mark_method_used_for_interface("m".into(), "f".into(), "A".into());
        b.mark_method_used_for_interface("m".into(), "f".into(), "B".into());
        b.mark_method_used_for_interface("m".into(), "g".into(), "C".into());

        a.merge(b);
        assert_eq!(a.interface_satisfied_methods.len(), 2);
        assert_eq!(
            a.interface_satisfied_methods[&("m".into(), "f".into())].len(),
            2
        );
        assert_eq!(
            a.interface_satisfied_methods[&("m".into(), "g".into())].len(),
            1
        );
    }
}
