# patent_1_cryptographic_validation.py
import hashlib
import hmac
import json
from typing import Any, Dict, List, Tuple, Union
import unittest


class AIDecision:
    """
    Represents a single AI decision with its input, output, and HMAC.
    """

    def __init__(self, input_data: Any, output_data: Any, hmac_key: bytes):
        """
        Initializes an AIDecision object.

        Args:
            input_data: The input data used for the decision.
            output_data: The output data generated by the AI.
            hmac_key: The secret key used for HMAC generation.
        """
        self.input_data = input_data
        self.output_data = output_data
        self.hmac_key = hmac_key
        self.hmac = self._generate_hmac()

    def _generate_hmac(self) -> str:
        """
        Generates an HMAC-SHA256 hash for the decision.

        Returns:
            The HMAC-SHA256 hash as a hexadecimal string.
        """
        message = json.dumps({"input": self.input_data, "output": self.output_data}).encode(
            "utf-8"
        )
        hmac_obj = hmac.new(self.hmac_key, message, hashlib.sha256)
        return hmac_obj.hexdigest()

    def verify_hmac(self) -> bool:
        """
        Verifies the HMAC-SHA256 hash of the decision.

        Returns:
            True if the HMAC is valid, False otherwise.
        """
        return self.hmac == self._generate_hmac()

    def to_dict(self) -> Dict[str, Any]:
        """
        Returns a dictionary representation of the decision.

        Returns:
            A dictionary containing the input, output, and HMAC.
        """
        return {
            "input": self.input_data,
            "output": self.output_data,
            "hmac": self.hmac,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any], hmac_key: bytes) -> "AIDecision":
        """
        Creates an AIDecision object from a dictionary.

        Args:
            data: A dictionary containing the input, output, and HMAC.
            hmac_key: The secret key used for HMAC generation.

        Returns:
            An AIDecision object.
        """
        decision = cls(data["input"], data["output"], hmac_key)
        decision.hmac = data["hmac"]
        return decision


class AIDecisionChain:
    """
    Represents a chain of AI decisions with cryptographic integrity.
    """

    def __init__(self, hmac_key: bytes):
        """
        Initializes an AIDecisionChain object.

        Args:
            hmac_key: The secret key used for HMAC generation.
        """
        self.chain: List[AIDecision] = []
        self.hmac_key = hmac_key
        self.previous_hmac: Union[str, None] = None

    def add_decision(self, input_data: Any, output_data: Any) -> None:
        """
        Adds a new decision to the chain.

        Args:
            input_data: The input data used for the decision.
            output_data: The output data generated by the AI.
        """
        decision = AIDecision(input_data, output_data, self.hmac_key)

        # Include the previous HMAC in the current decision's input
        if self.previous_hmac:
            decision.input_data = {"previous_hmac": self.previous_hmac, "input": input_data}

        decision.hmac = decision._generate_hmac()  # Recalculate HMAC after modification

        self.chain.append(decision)
        self.previous_hmac = decision.hmac

    def verify_chain(self) -> bool:
        """
        Verifies the integrity of the entire decision chain.

        Returns:
            True if the chain is valid, False otherwise.
        """
        if not self.chain:
            return True

        previous_hmac = None
        for i, decision in enumerate(self.chain):
            if not decision.verify_hmac():
                print(f"HMAC verification failed for decision at index {i}")
                return False

            if i > 0:
                # Verify that the previous HMAC matches what's stored in the current decision's input
                stored_previous_hmac = self.chain[i].input_data.get("previous_hmac")
                if stored_previous_hmac != previous_hmac:
                    print(f"Previous HMAC mismatch at index {i}")
                    return False

            previous_hmac = decision.hmac

        return True

    def to_list(self) -> List[Dict[str, Any]]:
        """
        Returns a list of dictionaries representing the decision chain.

        Returns:
            A list of dictionaries, each containing the input, output, and HMAC of a decision.
        """
        return [decision.to_dict() for decision in self.chain]

    @classmethod
    def from_list(cls, data: List[Dict[str, Any]], hmac_key: bytes) -> "AIDecisionChain":
        """
        Creates an AIDecisionChain object from a list of dictionaries.

        Args:
            data: A list of dictionaries, each containing the input, output, and HMAC of a decision.
            hmac_key: The secret key used for HMAC generation.

        Returns:
            An AIDecisionChain object.
        """
        chain = cls(hmac_key)
        chain.chain = [AIDecision.from_dict(d, hmac_key) for d in data]

        # Reconstruct the previous_hmac for chain verification
        if chain.chain:
            chain.previous_hmac = chain.chain[-1].hmac

        return chain

    def detect_tampering(self) -> bool:
        """
        Detects if the chain has been tampered with.

        Returns:
            True if tampering is detected, False otherwise.
        """
        return not self.verify_chain()


# Unit Tests
class TestAIDecisionChain(unittest.TestCase):
    def setUp(self):
        self.hmac_key = b"secret_key"
        self.chain = AIDecisionChain(self.hmac_key)

    def test_add_and_verify_decision(self):
        self.chain.add_decision("input1", "output1")
        self.assertTrue(self.chain.verify_chain())

    def test_add_multiple_decisions(self):
        self.chain.add_decision("input1", "output1")
        self.chain.add_decision("input2", "output2")
        self.chain.add_decision("input3", "output3")
        self.assertTrue(self.chain.verify_chain())

    def test_tampering_detection(self):
        self.chain.add_decision("input1", "output1")
        self.chain.add_decision("input2", "output2")
        chain_data = self.chain.to_list()
        # Tamper with the second decision
        chain_data[1]["output"] = "tampered_output"
        tampered_chain = AIDecisionChain.from_list(chain_data, self.hmac_key)
        self.assertTrue(tampered_chain.detect_tampering())

    def test_serialization_and_deserialization(self):
        self.chain.add_decision("input1", "output1")
        self.chain.add_decision("input2", "output2")
        chain_data = self.chain.to_list()
        new_chain = AIDecisionChain.from_list(chain_data, self.hmac_key)
        self.assertTrue(new_chain.verify_chain())
        self.assertEqual(len(self.chain.chain), len(new_chain.chain))
        self.assertEqual(self.chain.chain[0].input_data, new_chain.chain[0].input_data)
        self.assertEqual(self.chain.chain[1].output_data, new_chain.chain[1].output_data)

    def test_empty_chain(self):
        self.assertTrue(self.chain.verify_chain())  # An empty chain should be valid
        self.assertFalse(self.chain.detect_tampering())  # An empty chain cannot be tampered with

    def test_complex_input_output(self):
        input_data = {"feature1": 1.0, "feature2": "value"}
        output_data = [1, 2, 3]
        self.chain.add_decision(input_data, output_data)
        self.assertTrue(self.chain.verify_chain())


# Usage Example
if __name__ == "__main__":
    # Generate a secure random key for HMAC (important for security)
    import os
    hmac_key = os.urandom(32)  # 32 bytes for SHA256

    # Create an AI decision chain
    decision_chain = AIDecisionChain(hmac_key)

    # Add some decisions to the chain
    decision_chain.add_decision("What is 2 + 2?", "2 + 2 = 4")
    decision_chain.add_decision("Translate 'hello' to French", "Bonjour")
    decision_chain.add_decision({"user_id": 123, "query": "Recommend a movie"}, "Recommend movie 'Inception'")

    # Verify the integrity of the chain
    if decision_chain.verify_chain():
        print("AI decision chain is valid.")
    else:
        print("AI decision chain has been tampered with!")

    # Serialize the chain to a list of dictionaries
    chain_data = decision_chain.to_list()
    print("Serialized chain data:", chain_data)

    # Deserialize the chain from the list of dictionaries
    loaded_chain = AIDecisionChain.from_list(chain_data, hmac_key)

    # Verify the integrity of the loaded chain
    if loaded_chain.verify_chain():
        print("Loaded AI decision chain is valid.")
    else:
        print("Loaded AI decision chain has been tampered with!")

    # Tamper with the chain data (example)
    chain_data[1]["output"] = "Tampered Output"
    tampered_chain = AIDecisionChain.from_list(chain_data, hmac_key)

    # Detect tampering
    if tampered_chain.detect_tampering():
        print("Tampering detected in the AI decision chain!")
    else:
        print("No tampering detected.")

    # Run the unit tests
    unittest.main(argv=['first-arg-is-ignored'], exit=False)  # Prevent SystemExit