/*
 * Copyright Cedar Contributors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! This module is as straightforward a translation as possible of
//! <https://github.com/cedar-policy/cedar-spec/blob/main/cedar-lean/Cedar/SymCC.lean>.

mod authorizer;
pub mod bitvec;
mod compiler;
mod concretizer;
mod decoder;
mod encoder;
mod enforcer;
mod entity_tag;
mod env;
pub mod ext;
pub mod extension_types;
mod extfun;
mod extractor;
pub mod factory;
mod function;
mod interpretation;
pub mod op;
mod result;
mod smtlib_script;
pub mod solver;
mod symbolizer;
mod tags;
pub mod term;
pub mod term_type;
pub mod type_abbrevs;
pub mod verifier;

use cedar_policy::Schema;
use decoder::{parse_sexpr, IdMaps};
use env::to_validator_request_env;

use cedar_policy_core::ast::{Expr, ExprBuilder, Policy, PolicySet};
use cedar_policy_core::validator::{typecheck::Typechecker, types::RequestEnv, ValidationMode};
use encoder::Encoder;
use solver::{Decision, Solver};
use verifier::Asserts;

use crate::err::{Error, Result};

pub use bitvec::BitVecError;
pub use concretizer::ConcretizeError;
pub use concretizer::Env;
pub use decoder::DecodeError;
pub use encoder::EncodeError;
pub use env::{Environment, SymEnv};
pub use interpretation::Interpretation;
pub use result::CompileError;
pub use smtlib_script::SmtLibScript;
pub use solver::SolverError;
pub use verifier::{
    verify_always_allows, verify_always_denies, verify_disjoint, verify_equivalent, verify_implies,
    verify_never_errors,
};

/// Internal symbolic compiler.
/// Corresponds to the `SolverM` monad in Lean.
#[derive(Clone, Debug)]
pub struct SymCompiler<S> {
    solver: S,
}

impl<S> SymCompiler<S> {
    pub fn new(solver: S) -> Self {
        Self { solver }
    }

    pub fn solver(&self) -> &S {
        &self.solver
    }

    pub fn solver_mut(&mut self) -> &mut S {
        &mut self.solver
    }
}

impl<S: Solver> SymCompiler<S> {
    /// Given a verification condition generator `vc`,
    /// calls the SMT solver on an SMTLib encoding of the conditions generated by `vc`
    /// and returns `true` if the result is unsatisfiable. The function `vc` is
    /// expected to produce a list of terms type `Bool` that are well-formed
    /// with respect to the `symenv`.
    pub async fn check_unsat(
        &mut self,
        vc: impl FnOnce(&SymEnv) -> std::result::Result<Asserts, CompileError>,
        symenv: &SymEnv,
    ) -> Result<bool> {
        let asserts = vc(symenv)?;
        if asserts.iter().any(|assert| *assert == false.into()) {
            // some assert has been compiled to constant-false by symcc.
            // skip encoding and calling the solver.
            Ok(true)
        } else if asserts.iter().all(|assert| *assert == true.into()) {
            // all asserts have been compiled to constant-true by symcc.
            // skip encoding and calling the solver.
            Ok(false)
        } else {
            self.solver
                .smtlib_input()
                .reset()
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            self.solver
                .smtlib_input()
                .set_logic("ALL")
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            let mut encoder = Encoder::new(symenv, self.solver.smtlib_input())?;
            encoder.encode(asserts.iter()).await?;
            match self.solver.check_sat().await? {
                Decision::Unsat => Ok(true),
                Decision::Sat => Ok(false),
                Decision::Unknown => Err(Error::SolverUnknown),
            }
        }
    }

    /// Checks satisfiability, and then returns the model (as a pair of [`Request`] and [`Entities`]);
    /// For soundness, `footprint` must include all expressions used for generating the verification condition.
    pub async fn check_sat(
        &mut self,
        vc: impl FnOnce(&SymEnv) -> std::result::Result<Asserts, CompileError>,
        symenv: &SymEnv,
        footprint: impl Iterator<Item = &Expr>,
    ) -> Result<Option<Env>> {
        let asserts = vc(symenv)?;
        if asserts.iter().any(|assert| *assert == false.into()) {
            // some assert has been reduced to constant-false by the symcc process.
            // skip encoding and calling the solver.
            Ok(None)
        } else if asserts.iter().all(|assert| *assert == true.into()) {
            let interp = Interpretation::default(symenv);
            Ok(Some(symenv.interpret(&interp).concretize(footprint)?))
        } else {
            self.solver
                .smtlib_input()
                .reset()
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            self.solver
                .smtlib_input()
                .set_option("produce-models", "true")
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            self.solver
                .smtlib_input()
                .set_logic("ALL")
                .await
                .map_err(|err| Error::EncodeError(err.into()))?;
            let mut encoder =
                Encoder::new(symenv, self.solver.smtlib_input()).map_err(Error::EncodeError)?;
            encoder
                .encode(asserts.iter())
                .await
                .map_err(Error::EncodeError)?;
            let id_maps = IdMaps::from_encoder(&encoder);
            match self.solver.check_sat().await? {
                Decision::Unsat => Ok(None),
                Decision::Sat => {
                    let Some(model_str) = self.solver.get_model().await? else {
                        return Ok(None);
                    };
                    let model = parse_sexpr(model_str.as_bytes())?;
                    let interp = model.decode_model(symenv, &id_maps)?;
                    Ok(Some(symenv.extract(footprint, &interp)?))
                }
                Decision::Unknown => Err(Error::SolverUnknown),
            }
        }
    }

    /// Returns true iff `policy` does not error on any well-formed input in the `symenv`.
    pub async fn check_never_errors(&mut self, policy: &Policy, symenv: &SymEnv) -> Result<bool> {
        self.check_unsat(|symenv| verify_never_errors(policy, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_never_errors`] is false.
    pub async fn check_never_errors_with_counterexample(
        &mut self,
        policy: &Policy,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        self.check_sat(
            |symenv| verify_never_errors(policy, symenv),
            symenv,
            std::iter::once(&policy.condition()),
        )
        .await
    }

    /// Returns true iff the authorization decision of `policies1` implies that
    /// of `policies2` for every well-formed input in the `symenv`. That is,
    /// every input allowed by `policies1` is allowed by `policies2`;
    /// `policies2` is either more permissive than, or equivalent to, `policies1`.
    pub async fn check_implies(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(
            |symenv| verify_implies(policies1, policies2, symenv),
            symenv,
        )
        .await
    }

    /// Returns some counterexample iff [`Self::check_implies`] is false.
    pub async fn check_implies_with_counterexample(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies1
            .policies()
            .chain(policies2.policies())
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat(
            |symenv| verify_implies(policies1, policies2, symenv),
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff `policies` allows all well-formed inputs in the `symenv`.
    pub async fn check_always_allows(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(|symenv| verify_always_allows(policies, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_always_allows`] is false.
    pub async fn check_always_allows_with_counterexample(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies
            .policies()
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat(
            |symenv| verify_always_allows(policies, symenv),
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff `policies` denies all well-formed inputs in the `symenv`.
    pub async fn check_always_denies(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(|symenv| verify_always_denies(policies, symenv), symenv)
            .await
    }

    /// Returns some counterexample iff [`Self::check_always_denies`] is false.
    pub async fn check_always_denies_with_counterexample(
        &mut self,
        policies: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies
            .policies()
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat(
            |symenv| verify_always_denies(policies, symenv),
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff `policies1` and `policies2` produce the same
    /// authorization decision on all well-formed inputs in the `symenv`.
    pub async fn check_equivalent(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(
            |symenv| verify_equivalent(policies1, policies2, symenv),
            symenv,
        )
        .await
    }

    /// Returns some counterexample iff [`Self::check_equivalent`] is false.
    pub async fn check_equivalent_with_counterexample(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies1
            .policies()
            .chain(policies2.policies())
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat(
            |symenv| verify_equivalent(policies1, policies2, symenv),
            symenv,
            footprint.iter(),
        )
        .await
    }

    /// Returns true iff there is no well-formed input in the `symenv` that is allowed by both
    /// `policies1` and `policies2`. If this returns `false`, then there is at least one well-formed
    /// input that is allowed by both `policies1` and `policies2`.
    pub async fn check_disjoint(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<bool> {
        self.check_unsat(
            |symenv| verify_disjoint(policies1, policies2, symenv),
            symenv,
        )
        .await
    }

    /// Returns some counterexample iff [`Self::check_disjoint`] is false.
    pub async fn check_disjoint_with_counterexample(
        &mut self,
        policies1: &PolicySet,
        policies2: &PolicySet,
        symenv: &SymEnv,
    ) -> Result<Option<Env>> {
        let footprint = policies1
            .policies()
            .chain(policies2.policies())
            .map(|p| p.condition())
            .collect::<Vec<_>>();
        self.check_sat(
            |symenv| verify_disjoint(policies1, policies2, symenv),
            symenv,
            footprint.iter(),
        )
        .await
    }
}

/// The Cedar symbolic compiler assumes that it receives well-typed policies.  This
/// function enforces this requirement by calling Cedar's typechecker. Specifically,
/// given a policy `p` and type environment `Γ`, this function calls the Cedar
/// typechecker to obtain a policy `p'` that is semantically equivalent to `p` and
/// well-typed with respect to `Γ`.
///
/// All `check` functions defined in this file _must_ be called on the _output_ of
/// this function (or `well_typed_policies`) to ensure that symbolic compilation
/// succeeds. Applying the symbolic compiler directly to a policy `p` may result in
/// type errors---that is, the compiler rejecting the policy because it does not
/// satisfy the `WellTyped` constraints that are assumed by the compiler, and
/// enforced by the typechecker through policy transformation.
///
/// Note that unlike the corresponding Lean implementation which returns `Option` instead of `Result`.
/// We return a `Result` here to surface any type checking errors. In theory it should be impossible
/// to get an error since the policies are validated before calling SymCC.
pub fn well_typed_policy(
    policy: &Policy,
    env: &cedar_policy::RequestEnv,
    schema: &Schema,
) -> Result<Policy> {
    let env = to_validator_request_env(env, schema.as_ref())
        .ok_or_else(|| Error::ActionNotInSchema(env.action().to_string()))?;
    well_typed_policy_inner(policy, &env, schema)
}

fn well_typed_policy_inner(
    policy: &Policy,
    env: &RequestEnv<'_>,
    schema: &Schema,
) -> Result<Policy> {
    let validator_schema = schema.as_ref();
    // We need to perform these three checks like what the validator does here: https://github.com/cedar-policy/cedar/blob/82784864c01b5096cb73885dd2df5643074355ed/cedar-policy-core/src/validator.rs#L178-L185
    // We don't need `validate_template_action_application` because existence of `env` already serves as evidence
    let errs: Vec<_> = cedar_policy_core::validator::Validator::validate_entity_types(
        schema.as_ref(),
        policy.template(),
    )
    .chain(
        cedar_policy_core::validator::Validator::validate_enum_entity(
            schema.as_ref(),
            policy.template(),
        ),
    )
    .chain(
        cedar_policy_core::validator::Validator::validate_action_ids(
            schema.as_ref(),
            policy.template(),
        ),
    )
    .collect();
    if !errs.is_empty() {
        return Err(Error::PolicyNotWellTyped { errs });
    }
    let type_checker = Typechecker::new(validator_schema, ValidationMode::Strict);
    let policy_check = type_checker.typecheck_by_single_request_env(policy.template(), env);

    use cedar_policy_core::validator::typecheck::PolicyCheck::*;
    match policy_check {
        Success(expr) => Ok(Policy::from_when_clause(
            policy.effect(),
            expr.into_expr::<ExprBuilder<()>>(),
            policy.id().clone(),
            #[allow(
                clippy::useless_conversion,
                reason = "Additional `into()` to avoid build failure when the feature `cedar-policy-core:raw-parsing` is enabled"
            )]
            policy.loc().cloned().map(|l| l.into()),
        )),
        Irrelevant(errs, expr) =>
        // A policy could be irrelevant just for this environment, so unless there were errors we don't want to fail.
        // Note that if the policy was irrelevant for all environments schema validation would have caught this
        // before SymCC. The Lean implementation needs to be updated to match this behavior.
        {
            if errs.is_empty() {
                Ok(Policy::from_when_clause(
                    policy.effect(),
                    expr.into_expr::<ExprBuilder<()>>(),
                    policy.id().clone(),
                    #[allow(
                        clippy::useless_conversion,
                        reason = "Additional `into()` to avoid build failure when the feature `cedar-policy-core:raw-parsing` is enabled"
                    )]
                    policy.loc().cloned().map(|l| l.into()),
                ))
            } else {
                Err(Error::PolicyNotWellTyped { errs })
            }
        }
        Fail(errs) => Err(Error::PolicyNotWellTyped { errs }),
    }
}

pub fn well_typed_policies(
    policies: &PolicySet,
    env: &cedar_policy::RequestEnv,
    schema: &Schema,
) -> Result<PolicySet> {
    if policies.policies().any(|p| !p.is_static()) {
        return Err(CompileError::UnsupportedFeature(
            "template-linked policies are not supported".to_string(),
        )
        .into());
    }
    let env = to_validator_request_env(env, schema.as_ref())
        .ok_or_else(|| Error::ActionNotInSchema(env.action().to_string()))?;
    let typed_policies: Result<Vec<Policy>> = policies
        .static_policies()
        .map(|p| well_typed_policy_inner(p, &env, schema))
        .collect();
    match typed_policies {
        Ok(ps) => {
            let mut res = PolicySet::new();
            // PANIC SAFETY
            #[allow(
                clippy::unwrap_used,
                reason = "adding well-typed policy should not error"
            )]
            ps.into_iter().for_each(|p| res.add(p).unwrap());
            Ok(res)
        }
        Err(err) => Err(err),
    }
}
