"""Tests for prooftrail.py: decorator-based tamper-evident audit trails."""

import pytest

from ads_foundation.prooftrail import ProofTrailMonitor, Receipt, Session


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def make_monitor_and_fn():
    """Create a monitor with a simple tracked function."""
    monitor = ProofTrailMonitor()

    @monitor.track_tool_call
    def query_db(customer_id: str) -> dict:
        return {"balance": 15420.50, "id": customer_id}

    return monitor, query_db


# ---------------------------------------------------------------------------
# Decorator basics
# ---------------------------------------------------------------------------

class TestDecoratorBasics:
    def test_returns_same_result(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            result = query_db("cust_1")
        assert result == {"balance": 15420.50, "id": "cust_1"}

    def test_preserves_function_name(self):
        monitor, query_db = make_monitor_and_fn()
        assert query_db.__name__ == "query_db"

    def test_no_tracking_outside_session(self):
        monitor, query_db = make_monitor_and_fn()
        result = query_db("cust_1")
        assert result == {"balance": 15420.50, "id": "cust_1"}
        # No session was active, so no receipts anywhere
        with monitor.session(agent_id="bot") as session:
            assert session.trail() == []


# ---------------------------------------------------------------------------
# Receipt generation
# ---------------------------------------------------------------------------

class TestReceiptGeneration:
    def test_receipt_per_call(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("cust_1")
            query_db("cust_2")
            query_db("cust_3")
        assert len(session.trail()) == 3

    def test_receipt_fields(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="support_bot") as session:
            query_db("cust_987")
            receipt = session.get_latest_receipt()

        assert receipt.index == 0
        assert receipt.agent_id == "support_bot"
        assert receipt.tool_name == "query_db"
        assert isinstance(receipt.timestamp, str)
        assert len(receipt.query_hash) == 32
        assert len(receipt.response_hash) == 32
        assert len(receipt.record_hash) == 32
        assert len(receipt.root_hash) == 32

    def test_get_receipt_by_index(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("a")
            query_db("b")
        r0 = session.get_receipt(0)
        r1 = session.get_receipt(1)
        assert r0.index == 0
        assert r1.index == 1

    def test_get_receipt_missing_index(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("a")
        with pytest.raises(KeyError):
            session.get_receipt(99)

    def test_get_latest_receipt_empty(self):
        monitor = ProofTrailMonitor()
        with monitor.session(agent_id="bot") as session:
            with pytest.raises(ValueError):
                session.get_latest_receipt()

    def test_receipt_is_frozen(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("a")
            receipt = session.get_latest_receipt()
        with pytest.raises(AttributeError):
            receipt.index = 999  # type: ignore[misc]


# ---------------------------------------------------------------------------
# Verification
# ---------------------------------------------------------------------------

class TestVerification:
    def test_valid_receipt_verifies(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("cust_1")
            receipt = session.get_latest_receipt()
            assert session.verify_receipt(receipt)

    def test_monitor_verify_receipt(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("cust_1")
            receipt = session.get_latest_receipt()
        # Verify outside the session via monitor
        assert monitor.verify_receipt(receipt)

    def test_tampered_record_hash_fails(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("cust_1")
            receipt = session.get_latest_receipt()

        # Tamper with the record hash
        tampered = Receipt(
            index=receipt.index,
            timestamp=receipt.timestamp,
            agent_id=receipt.agent_id,
            tool_name=receipt.tool_name,
            query_hash=receipt.query_hash,
            response_hash=receipt.response_hash,
            record_hash=b"\xff" * 32,  # tampered
            root_hash=receipt.root_hash,
            proof=receipt.proof,
        )
        assert not monitor.verify_receipt(tampered)

    def test_tampered_proof_fails(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("a")
            query_db("b")
            query_db("c")
            receipt = session.get_receipt(0)

        # Tamper with a sibling hash in the proof
        from ads_foundation.merkle import MerkleProof

        if receipt.proof.path:
            side, sib = receipt.proof.path[0]
            bad_sib = bytearray(sib)
            bad_sib[0] ^= 0xFF
            tampered_proof = MerkleProof(
                path=[(side, bytes(bad_sib))] + receipt.proof.path[1:],
                root_hash=receipt.proof.root_hash,
            )
            tampered = Receipt(
                index=receipt.index,
                timestamp=receipt.timestamp,
                agent_id=receipt.agent_id,
                tool_name=receipt.tool_name,
                query_hash=receipt.query_hash,
                response_hash=receipt.response_hash,
                record_hash=receipt.record_hash,
                root_hash=receipt.root_hash,
                proof=tampered_proof,
            )
            assert not monitor.verify_receipt(tampered)


# ---------------------------------------------------------------------------
# Chain integrity
# ---------------------------------------------------------------------------

class TestChainIntegrity:
    def test_sequential_indices(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            for i in range(5):
                query_db(f"cust_{i}")
        receipts = session.trail()
        assert [r.index for r in receipts] == [0, 1, 2, 3, 4]

    def test_each_receipt_has_unique_record_hash(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            query_db("a")
            query_db("a")  # same args, different timestamp -> different hash
        r0, r1 = session.trail()
        assert r0.record_hash != r1.record_hash

    def test_all_receipts_verify(self):
        monitor, query_db = make_monitor_and_fn()
        with monitor.session(agent_id="bot") as session:
            for i in range(8):
                query_db(f"cust_{i}")
        for receipt in session.trail():
            assert monitor.verify_receipt(receipt)


# ---------------------------------------------------------------------------
# Multiple sessions
# ---------------------------------------------------------------------------

class TestMultipleSessions:
    def test_different_agent_ids(self):
        monitor, query_db = make_monitor_and_fn()

        with monitor.session(agent_id="agent_a") as s1:
            query_db("x")
        with monitor.session(agent_id="agent_b") as s2:
            query_db("y")

        assert s1.get_latest_receipt().agent_id == "agent_a"
        assert s2.get_latest_receipt().agent_id == "agent_b"

    def test_sessions_share_chain(self):
        monitor, query_db = make_monitor_and_fn()

        with monitor.session(agent_id="a") as s1:
            query_db("x")
        with monitor.session(agent_id="b") as s2:
            query_db("y")

        # Second session's receipt should have index 1 (continues the chain)
        assert s2.get_latest_receipt().index == 1

    def test_session_isolation(self):
        monitor, query_db = make_monitor_and_fn()

        with monitor.session(agent_id="a") as s1:
            query_db("x")

        with monitor.session(agent_id="b") as s2:
            query_db("y")

        assert len(s1.trail()) == 1
        assert len(s2.trail()) == 1


# ---------------------------------------------------------------------------
# Multiple decorated functions
# ---------------------------------------------------------------------------

class TestMultipleFunctions:
    def test_different_tools_same_session(self):
        monitor = ProofTrailMonitor()

        @monitor.track_tool_call
        def search_web(query: str) -> dict:
            return {"results": 42}

        @monitor.track_tool_call
        def run_code(code: str) -> dict:
            return {"output": "hello"}

        with monitor.session(agent_id="multi_bot") as session:
            search_web("test query")
            run_code("print('hello')")

        receipts = session.trail()
        assert len(receipts) == 2
        assert receipts[0].tool_name == "search_web"
        assert receipts[1].tool_name == "run_code"
        assert monitor.verify_receipt(receipts[0])
        assert monitor.verify_receipt(receipts[1])


# ---------------------------------------------------------------------------
# API key
# ---------------------------------------------------------------------------

class TestAPIKey:
    def test_api_key_accepted(self):
        monitor = ProofTrailMonitor(api_key="pt_key_xyz")

        @monitor.track_tool_call
        def fn() -> str:
            return "ok"

        with monitor.session(agent_id="bot") as session:
            result = fn()
        assert result == "ok"
        assert session.verify_receipt(session.get_latest_receipt())

    def test_api_key_none_by_default(self):
        monitor = ProofTrailMonitor()
        assert monitor._api_key is None
