#!/usr/bin/env python3
"""
test_inference_receipt.py - Unit tests for inference receipt schema

Run:
    python -m pytest tests/test_inference_receipt.py -v
    python tests/test_inference_receipt.py  # Standalone
"""

import sys
import json
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))

from helix_cdc.receipts.inference_receipt_v1 import (
    InferenceReceipt,
    Trail,
    ISATrace,
    ISAOp,
    Routing,
    SeInput,
    DecisionEvidence,
    TopKEntry,
    DecisionMargins,
    Guardrails,
    GateCheck,
    Seal,
    CDRAnchor,
    MerkleProof,
    Provenance,
    canonicalize,
    sha256_canonical,
    SCHEMA_VERSION,
)


def test_canonicalize_basic():
    """Canonicalization should produce deterministic output."""
    obj = {"b": 2, "a": 1, "c": 3}
    result = canonicalize(obj)
    assert result == '{"a":1,"b":2,"c":3}'


def test_canonicalize_nested():
    """Nested objects should be canonicalized."""
    obj = {"z": {"y": 2, "x": 1}, "a": [3, 1, 2]}
    result = canonicalize(obj)
    # Keys sorted, arrays preserved
    assert '"a":[3,1,2]' in result
    assert '"x":1' in result


def test_canonicalize_floats():
    """Floats should be consistently formatted."""
    obj = {"a": 1.0, "b": 1.5, "c": 1.123456789012345}
    result = canonicalize(obj)
    # 1.0 should become 1 (int form)
    assert '"a":1' in result
    # Precision should be limited
    parsed = json.loads(result)
    assert abs(parsed["c"] - 1.123456789) < 1e-9


def test_isa_op_creation():
    """ISAOp should be created correctly."""
    op = ISAOp(
        op="ADD",
        op_hex="0x03",
        args=[0, 1, 2],
        cost_cycles=2,
        cost_flops=1,
        cost_bytes=24,
    )
    assert op.op == "ADD"
    assert op.cost_cycles == 2


def test_isa_trace_auto_summary():
    """ISATrace should auto-compute summary."""
    ops = [
        ISAOp("LOAD", "0x01", [0, 100], cost_cycles=3, cost_bytes=8),
        ISAOp("ADD", "0x03", [0, 1, 2], cost_cycles=2, cost_flops=1, cost_bytes=24),
    ]
    trace = ISATrace(ops_executed=ops)

    assert trace.ops_summary["total_ops"] == 2
    assert trace.ops_summary["total_cycles"] == 5
    assert trace.ops_summary["total_flops"] == 1
    assert trace.ops_summary["total_bytes"] == 32


def test_routing_auto_compute_se():
    """Routing should auto-compute Se if not provided."""
    routing = Routing(
        se_input=SeInput(H_entropy=5.0, C_coherence=0.8, D_depth=10),
        se_computed=0,  # Will be computed
        route_decision="core_0",
        route_reason="test",
    )
    assert routing.se_computed == 40.0  # 5 * 0.8 * 10


def test_decision_evidence_classify():
    """Decision confidence classification."""
    assert DecisionEvidence.classify_confidence(0.3, 2.5) == "near_tie"
    assert DecisionEvidence.classify_confidence(4.0, 0.5) == "clear_winner"
    assert DecisionEvidence.classify_confidence(2.0, 1.5) == "moderate"


def test_guardrails_verdict():
    """Guardrails verdict computation."""
    # All pass
    gr = Guardrails(
        gates_checked=[
            GateCheck("g1", True, 1, 1, ""),
            GateCheck("g2", True, 2, 2, ""),
        ],
        verdict=""
    )
    assert gr.compute_verdict() == "PASS"

    # Hard failure
    gr = Guardrails(
        gates_checked=[
            GateCheck("g1", True, 1, 1, ""),
            GateCheck("g2", False, 2, 3, ""),
        ],
        verdict=""
    )
    assert gr.compute_verdict() == "FAIL"

    # Soft failure only
    gr = Guardrails(
        gates_checked=[
            GateCheck("g1", True, 1, 1, ""),
            GateCheck("soft_g2", False, 2, 3, ""),
        ],
        verdict=""
    )
    assert gr.compute_verdict() == "ACCEPTABLE_WITH_TAIL_RISK"


def test_seal_chain_hash():
    """Seal chain hash computation."""
    seal = Seal(
        trail_hash="a" * 64,
        state_hash="b" * 64,
        inputs_hash="c" * 64,
        outputs_hash="d" * 64,
    )

    chain_hash = seal.compute_chain_hash("parent123")
    assert len(chain_hash) == 64
    assert seal.chain_hash == chain_hash


def test_receipt_creation():
    """Full receipt creation."""
    trail = Trail(
        isa_trace=ISATrace(ops_executed=[
            ISAOp("ADD", "0x03", [0, 1, 2], 2, 1, 24),
        ]),
    )

    receipt = InferenceReceipt.create(
        step_id="step_1",
        trail=trail,
        inputs={"prompt": "test"},
        outputs={"result": 42},
        cdr_state_hash="abc123",
    )

    assert receipt.step_id == "step_1"
    assert receipt.schema_version == SCHEMA_VERSION
    assert len(receipt.seal.trail_hash) == 64
    assert len(receipt.seal.chain_hash) == 64


def test_receipt_json_roundtrip():
    """Receipt should survive JSON roundtrip."""
    trail = Trail(
        isa_trace=ISATrace(ops_executed=[
            ISAOp("LOAD", "0x01", [0], 3, 0, 8),
            ISAOp("ADD", "0x03", [0, 1, 2], 2, 1, 24),
        ]),
        routing=Routing(
            se_input=SeInput(5.0, 0.8, 10),
            se_computed=40.0,
            route_decision="core_0",
            route_reason="test",
        ),
        decision_evidence=DecisionEvidence(
            task_type="token_selection",
            decision_made=100,
            top_k=[TopKEntry(100, 5.0, 0.8, 1, "the")],
            margins=DecisionMargins(2.0, 0.0, 0.5, 0.01),
            confidence_classification="clear_winner",
        ),
        guardrails=Guardrails(
            gates_checked=[GateCheck("g1", True, 1, 1, "ok")],
            verdict="PASS",
        ),
    )

    original = InferenceReceipt.create(
        step_id="step_1",
        trail=trail,
        inputs={"a": 1},
        outputs={"b": 2},
        cdr_state_hash="xyz",
        cdr_anchor=CDRAnchor("cp_1", 5, "hash", True, 3),
    )

    json_str = original.to_json()
    restored = InferenceReceipt.from_json(json_str)

    assert restored.step_id == original.step_id
    assert restored.seal.trail_hash == original.seal.trail_hash
    assert restored.seal.chain_hash == original.seal.chain_hash
    assert restored.receipt_hash() == original.receipt_hash()


def test_receipt_verify_trail_hash():
    """Trail hash verification."""
    trail = Trail(
        isa_trace=ISATrace(ops_executed=[
            ISAOp("ADD", "0x03", [0, 1, 2], 2, 1, 24),
        ]),
    )

    receipt = InferenceReceipt.create(
        step_id="step_1",
        trail=trail,
        inputs={},
        outputs={},
        cdr_state_hash="abc",
    )

    assert receipt.verify_trail_hash()


def test_receipt_chain_verification():
    """Chain verification between receipts."""
    # Genesis receipt
    receipt1 = InferenceReceipt.create(
        step_id="step_0",
        trail=Trail(),
        inputs={},
        outputs={},
        cdr_state_hash="state0",
    )

    # Chained receipt
    receipt2 = InferenceReceipt.create(
        step_id="step_1",
        trail=Trail(),
        inputs={},
        outputs={},
        cdr_state_hash="state1",
        parent_step_id="step_0",
        parent_chain_hash=receipt1.seal.chain_hash,
    )

    assert receipt1.verify_chain(None)  # Genesis
    assert receipt2.verify_chain(receipt1)  # Chained


def test_merkle_proof_in_receipt():
    """Merkle proof should be stored in receipt."""
    receipt = InferenceReceipt.create(
        step_id="step_1",
        trail=Trail(),
        inputs={},
        outputs={},
        cdr_state_hash="abc",
        merkle_proof=MerkleProof(
            leaf_index=5,
            leaf_hash="a" * 64,
            siblings=[{"hash": "b" * 64, "position": "right"}],
            root_hash="c" * 64,
        ),
    )

    assert receipt.seal.merkle_proof is not None
    assert receipt.seal.merkle_proof.leaf_index == 5


def test_provenance_from_environment():
    """Provenance should capture environment."""
    prov = Provenance.from_environment()

    assert prov.python_version  # Should be set
    assert prov.platform  # Should be set
    assert prov.machine  # Should be set


def test_canonical_json():
    """Canonical JSON should be deterministic."""
    receipt = InferenceReceipt.create(
        step_id="test",
        trail=Trail(),
        inputs={"z": 1, "a": 2},
        outputs={"b": 3},
        cdr_state_hash="xyz",
    )

    canonical1 = receipt.to_canonical_json()
    canonical2 = receipt.to_canonical_json()

    assert canonical1 == canonical2


def test_receipt_hash_determinism():
    """Receipt hash should be deterministic."""
    for _ in range(10):
        receipt = InferenceReceipt.create(
            step_id="test",
            trail=Trail(isa_trace=ISATrace(ops_executed=[])),
            inputs={"a": 1},
            outputs={"b": 2},
            cdr_state_hash="xyz",
        )

        hash1 = receipt.receipt_hash()
        hash2 = receipt.receipt_hash()
        assert hash1 == hash2


# =============================================================================
# STANDALONE TEST RUNNER
# =============================================================================

if __name__ == "__main__":
    import traceback

    tests = [
        test_canonicalize_basic,
        test_canonicalize_nested,
        test_canonicalize_floats,
        test_isa_op_creation,
        test_isa_trace_auto_summary,
        test_routing_auto_compute_se,
        test_decision_evidence_classify,
        test_guardrails_verdict,
        test_seal_chain_hash,
        test_receipt_creation,
        test_receipt_json_roundtrip,
        test_receipt_verify_trail_hash,
        test_receipt_chain_verification,
        test_merkle_proof_in_receipt,
        test_provenance_from_environment,
        test_canonical_json,
        test_receipt_hash_determinism,
    ]

    passed = 0
    failed = 0

    print("Inference Receipt Tests")
    print("=" * 60)

    for test in tests:
        try:
            test()
            print(f"  ✓ {test.__name__}")
            passed += 1
        except AssertionError as e:
            print(f"  ✗ {test.__name__}: {e}")
            failed += 1
        except Exception as e:
            print(f"  ✗ {test.__name__}: {type(e).__name__}: {e}")
            traceback.print_exc()
            failed += 1

    print()
    print(f"Results: {passed} passed, {failed} failed")
    sys.exit(0 if failed == 0 else 1)
