"""
Extended proofs are proofs that deal with internal precommitments.

This is a convenient building block that is useful for defining complex primitives.

Precommitments are parameters that are not known at the proof instantiation time
and are computed by the proving side. Therefore, they have to be sent explicitly
in a preliminary round for interactive protocols, and as an additional attribute
in non-interactive or simulated transcripts.

The protocol of an extended proof is as follows:

1. A proof statement (:py:class:`ExtendedProofStmt`) is instantiated and its
   attributes are set, but cannot run most methods.
2. Upon processing of the precommitment
   (:py:meth:`ExtendedProofProver.process_precommitment`), a separate and complete
   regular proof statement (:py:class:`zksk.base.ComposableProofStmt`) is constructed inside the
   extended proof (:py:attr`ExtendedProofStmt.constructed_stmt`)
3. :py:meth:`ExtendedProofStmt.get_bases`,
   :py:meth:`ExtendedProofStmt.get_secret_vars` delegate to the internal
   ``constructed_stmt``. The :py:class:`ExtendedProver` and
   :py:class:`ExtendedVerifier` work in the same way, embedding a
   ``constructed_prover`` (resp., ``constructed_verifier``).
"""

import abc

from zksk.base import Prover, Verifier
from zksk.composition import ComposableProofStmt
from zksk.exceptions import StatementSpecError
from zksk.utils.misc import get_default_attr


class ExtendedProofStmt(ComposableProofStmt, metaclass=abc.ABCMeta):
    """
    Proof that deals with precommitments.
    """

    @abc.abstractmethod
    def construct_stmt(self):
        """
        Build internal proof statement for this class

        This function must be overridden. The function should return a
        constructed proof statement. It can use the values that were computed
        by internal_precommit to do so.
        """
        pass

    def precommit(self):
        """
        Compute precommitments. Override if needed.

        Override this function to compute precommitments and set corresponding
        secrets that must be computed before the ZK proof itself can be
        constructed and proven.

        Returns:
            Precommitment
        """
        return []

    def simulate_precommit(self):
        """
        Simulate a precommitment. Override if needed.

        Override this method to enable using this proof in or-proofs.

        It should compute the same output as generated by precommit, but without relying on any
        secrets.
        """
        raise StatementSpecError(
            "Override simulate_precommit in order to " "use or-proofs and simulations"
        )

    def validate(self, precommitment, *args, **kwargs):
        """
        Validate proof's construction. Override if needed.
        """
        pass

    def full_validate(self, *args, **kwargs):
        self.validate(self.precommitment, *args, **kwargs)

    @property
    def constructed_stmt(self):
        return get_default_attr(self, "_constructed_stmt")

    @property
    def precommitment(self):
        return get_default_attr(self, "_precommitment")

    def get_secret_vars(self):
        return self.constructed_stmt.get_secret_vars()

    def get_bases(self):
        return self.constructed_stmt.get_bases()

    def get_prover(self, secrets_dict=None):
        """
        Get a prover object.

        Returns:
            Prover object if all secret values are known, None otherwise.
        """
        if secrets_dict is None:
            secrets_dict = {}

        for k, v in secrets_dict.items():
            k.value = v

        self.secret_values = {}
        self.secret_values.update(secrets_dict)

        return ExtendedProver(self, self.secret_values)

    def get_verifier_cls(self):
        return ExtendedVerifier

    def recompute_commitment(self, challenge, responses):
        return self.constructed_stmt.recompute_commitment(challenge, responses)

    def get_proof_id(self, secret_id_map=None):
        """
        Identifier for the proof statement.
        """
        if self.constructed_stmt is not None:
            proof_id = [
                self.__class__.__name__,
                self.precommitment,
                self.constructed_stmt.get_proof_id(secret_id_map),
            ]
        else:
            raise ValueError("Proof ID unknown before the proof is constructed.")
        return proof_id

    def full_construct_stmt(self, precommitment):
        self._precommitment = precommitment
        self._constructed_stmt = self.construct_stmt(precommitment)
        return self.constructed_stmt

    def prepare_simulate_proof(self):
        self._precommitment = self.simulate_precommit()
        self.full_construct_stmt(self.precommitment)

    def simulate_proof(self, responses_dict=None, challenge=None):
        """
        Simulate the proof.

        Args:
            responses_dict: Mapping from secrets to responses
            challenge: Challenge
        """
        tr = self._constructed_stmt.simulate_proof(
            challenge=challenge, responses_dict=responses_dict
        )
        tr.precommitment = self._precommitment
        return tr

    def _precommit(self):
        self._precommitment = self.precommit()
        return self._precommitment


class ExtendedProver(Prover):
    """
    Prover dealing with precommitments.

    This prover will create a constructed Prover object and delegate to its methods.
    """

    def internal_commit(self, randomizers_dict=None):
        """
        Trigger the internal prover commit.

        Transfers the randomizer_dict if passed. It might be used if the binding
        of the proof is set True.
        """
        if self.stmt.constructed_stmt is None:
            raise StatementSpecError(
                "You need to pre-commit before commiting. The proofs lack parameters otherwise."
            )
        return self.constructed_prover.internal_commit(randomizers_dict)

    def compute_response(self, challenge):
        """
        Wrap the response computation for the inner proof.
        """
        self.challenge = challenge
        self.constructed_prover.challenge = challenge
        self.response = self.constructed_prover.compute_response(challenge)
        return self.response

    def precommit(self):
        self.precommitment = self.stmt._precommit()
        self.process_precommitment()
        try:
            for p in self.constructed_prover.subs:
                p.precommit()
        except AttributeError:
            pass
        return self.precommitment

    def process_precommitment(self):
        """
        Trigger the inner-proof construction and extract a prover given the secrets.
        """
        self.stmt.full_construct_stmt(self.precommitment)
        self.constructed_prover = self.stmt._constructed_stmt.get_prover(
            self.secret_values
        )


class ExtendedVerifier(Verifier):
    """
    Verifier that deals with precommitments.
    """

    def process_precommitment(self, precommitment):
        """
        Receive the precommitment and trigger the inner-verifier construction.
        """
        self.precommitment = precommitment
        self.stmt.full_construct_stmt(precommitment)
        self.constructed_verifier = self.stmt.constructed_stmt.get_verifier()

    def send_challenge(self, com):
        """
        Transfer the commitment to the inner proof, and compute the challenge.
        """

        statement, self.commitment = com
        self.stmt.check_statement(statement)
        self.challenge = self.constructed_verifier.send_challenge(
            self.commitment, ignore_statement_hash_checks=True
        )
        return self.challenge

    def check_responses_consistency(self, responses, responses_dict):
        """
        Wrap the consistency check of the innter proof.
        """
        return self.constructed_verifier.check_responses_consistency(
            responses, responses_dict
        )
