// https://crates.io/crates/salsa
// hello world https://github.com/salsa-rs/salsa/blob/master/examples/hello_world/main.rs

use crate::error::ChalkError;
use crate::interner::ChalkIr;
use crate::lowering::LowerProgram;
use crate::program::Program;
use crate::program_environment::ProgramEnvironment;
use crate::tls;
use chalk_ir::TraitId;
use chalk_solve::clauses::builder::ClauseBuilder;
use chalk_solve::clauses::program_clauses::ToProgramClauses;
use chalk_solve::coherence::orphan;
use chalk_solve::coherence::{CoherenceSolver, SpecializationPriorities};
use chalk_solve::wf;
use chalk_solve::RustIrDatabase;
use chalk_solve::Solver;
use chalk_solve::SolverChoice;
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::Mutex;

#[salsa::query_group(Lowering)]
pub trait LoweringDatabase: RustIrDatabase<ChalkIr> {
    #[salsa::input]
    fn program_text(&self) -> Arc<String>;

    #[salsa::input]
    fn solver_choice(&self) -> SolverChoice;

    fn program_ir(&self) -> Result<Arc<Program>, ChalkError>;

    /// Performs coherence check and computes which impls specialize
    /// one another (the "specialization priorities").
    fn coherence(
        &self,
    ) -> Result<BTreeMap<TraitId<ChalkIr>, Arc<SpecializationPriorities<ChalkIr>>>, ChalkError>;

    fn orphan_check(&self) -> Result<(), ChalkError>;

    /// The lowered IR, with coherence, orphan, and WF checks performed.
    fn checked_program(&self) -> Result<Arc<Program>, ChalkError>;

    /// The program as logic.
    fn environment(&self) -> Result<Arc<ProgramEnvironment>, ChalkError>;

    /// Creates the solver we can use to solve goals. This solver
    /// stores intermediate, cached state, which is why it is behind a
    /// mutex. Moreover, if the set of program clauses change, that
    /// cached state becomes invalid, so the query is marked as
    /// volatile, thus ensuring that the solver is recreated in every
    /// revision (i.e., each time source program changes).
    #[salsa::volatile]
    fn solver(&self) -> Arc<Mutex<Solver<ChalkIr>>>;
}

fn program_ir(db: &impl LoweringDatabase) -> Result<Arc<Program>, ChalkError> {
    let text = db.program_text();
    Ok(Arc::new(chalk_parse::parse_program(&text)?.lower()?))
}

fn orphan_check(db: &impl LoweringDatabase) -> Result<(), ChalkError> {
    let program = db.program_ir()?;

    tls::set_current_program(&program, || -> Result<(), ChalkError> {
        let local_impls = program.local_impl_ids();
        for impl_id in local_impls {
            orphan::perform_orphan_check(db, db.solver_choice(), impl_id)?;
        }
        Ok(())
    })
}

fn coherence(
    db: &impl LoweringDatabase,
) -> Result<BTreeMap<TraitId<ChalkIr>, Arc<SpecializationPriorities<ChalkIr>>>, ChalkError> {
    let program = db.program_ir()?;

    let priorities_map = tls::set_current_program(&program, || -> Result<_, ChalkError> {
        let priorities_map: Result<BTreeMap<_, _>, ChalkError> = program
            .trait_data
            .keys()
            .map(|&trait_id| {
                let solver = CoherenceSolver::new(db, db.solver_choice(), trait_id);
                let priorities = solver.specialization_priorities()?;
                Ok((trait_id, priorities))
            })
            .collect();
        let priorities_map = priorities_map?;
        Ok(priorities_map)
    });
    let () = db.orphan_check()?;
    priorities_map
}

fn checked_program(db: &impl LoweringDatabase) -> Result<Arc<Program>, ChalkError> {
    let program = db.program_ir()?;

    db.coherence()?;

    let () = tls::set_current_program(&program, || -> Result<(), ChalkError> {
        let solver = wf::WfSolver::new(db, db.solver_choice());
        for &id in program.adt_data.keys() {
            solver.verify_adt_decl(id)?;
        }

        for &impl_id in program.impl_data.keys() {
            solver.verify_trait_impl(impl_id)?;
        }

        Ok(())
    })?;

    Ok(program)
}

fn environment(db: &impl LoweringDatabase) -> Result<Arc<ProgramEnvironment>, ChalkError> {
    let program = db.program_ir()?;

    // Construct the set of *clauses*; these are sort of a compiled form
    // of the data above that always has the form:
    //
    //       forall P0...Pn. Something :- Conditions
    let mut program_clauses = program.custom_clauses.clone();

    let builder = &mut ClauseBuilder::new(db, &mut program_clauses);

    program
        .associated_ty_data
        .values()
        .for_each(|d| d.to_program_clauses(builder));

    program
        .trait_data
        .values()
        .for_each(|d| d.to_program_clauses(builder));

    program
        .adt_data
        .values()
        .for_each(|d| d.to_program_clauses(builder));

    for (&auto_trait_id, _) in program
        .trait_data
        .iter()
        .filter(|(_, auto_trait)| auto_trait.is_auto_trait())
    {
        for &adt_id in program.adt_data.keys() {
            chalk_solve::clauses::push_auto_trait_impls(builder, auto_trait_id, adt_id);
        }
    }

    for datum in program.impl_data.values() {
        // If we encounter a negative impl, do not generate any rule. Negative impls
        // are currently just there to deactivate default impls for auto traits.
        if datum.is_positive() {
            datum.to_program_clauses(builder);
            datum
                .associated_ty_value_ids
                .iter()
                .map(|&atv_id| db.associated_ty_value(atv_id))
                .for_each(|atv| atv.to_program_clauses(builder));
        }
    }

    Ok(Arc::new(ProgramEnvironment::new(program_clauses)))
}

fn solver(db: &impl LoweringDatabase) -> Arc<Mutex<Solver<ChalkIr>>> {
    let choice = db.solver_choice();
    Arc::new(Mutex::new(choice.into_solver()))
}
