"""Comprehensive tests for RLang → BoR cryptographic integration."""

import json

import pytest

from rlang.bor import RLangBoRCrypto, run_program_with_proof
from rlang.bor.crypto import compute_subproofs, hash_master, hash_step
from rlang.utils.canonical_json import canonical_dumps

# Try to import RichProofBundle, fall back to mock if not available
try:
    from bor.bundle import RichProofBundle
except ImportError:
    from rlang.bor.crypto import RichProofBundle


def test_step_hashing_correctness():
    """Test that step hashing produces correct hashes."""
    src = """
fn inc(x: Int) -> Int;

pipeline main(Int) -> Int { inc }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry={"inc": lambda x: x + 1},
    )

    crypto = RLangBoRCrypto(proof)
    step_records = crypto.compute_step_hashes()

    assert len(step_records) == 1

    # Verify step hash matches direct hash computation
    step = proof.steps[0]
    step_dict = step.to_dict()
    expected_hash = hash_step(step_dict)

    assert step_records[0].step_hash == expected_hash
    assert step_records[0].index == 0
    assert step_records[0].template_id == step.template_id


def test_step_hashing_multi_step():
    """Test step hashing for multi-step pipeline."""
    src = """
fn inc(x: Int) -> Int;
fn double(x: Int) -> Int;

pipeline main(Int) -> Int { inc -> double }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=5,
        fn_registry={"inc": lambda x: x + 1, "double": lambda x: x * 2},
    )

    crypto = RLangBoRCrypto(proof)
    step_records = crypto.compute_step_hashes()

    assert len(step_records) == 2

    # Verify each step hash matches direct computation
    for i, step in enumerate(proof.steps):
        step_dict = step.to_dict()
        expected_hash = hash_step(step_dict)
        assert step_records[i].step_hash == expected_hash
        assert step_records[i].index == step.index


def test_HMASTER_correctness():
    """Test that HMASTER computation is correct."""
    src = """
fn step1(x: Int) -> Int;
fn step2(x: Int) -> Int;

pipeline main(Int) -> Int { step1 -> step2 }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=1,
        fn_registry={"step1": lambda x: x + 1, "step2": lambda x: x * 2},
    )

    crypto = RLangBoRCrypto(proof)
    step_records = crypto.compute_step_hashes()

    # Compute HMASTER
    HMASTER = crypto.compute_HMASTER(step_records)

    # Verify HMASTER matches direct computation
    step_hashes = [r.step_hash for r in step_records]
    expected_master = hash_master(step_hashes)

    assert HMASTER == expected_master
    assert isinstance(HMASTER, str)
    assert len(HMASTER) == 64  # SHA-256 hex digest length


def test_HRICH_correctness():
    """Test that HRICH computation is correct."""
    src = """
fn test(x: Int) -> Int;

pipeline main(Int) -> Int { test }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=42,
        fn_registry={"test": lambda x: x},
    )

    crypto = RLangBoRCrypto(proof)
    step_records = crypto.compute_step_hashes()
    HMASTER = crypto.compute_HMASTER(step_records)

    # Compute HRICH
    HRICH = crypto.compute_HRICH(HMASTER)

    # Verify HRICH matches direct computation via subproofs
    from rlang.bor.crypto import compute_subproof_hashes, compute_HRICH_from_subproof_hashes
    subproofs = compute_subproofs(primary_hash=HMASTER)
    subproof_hashes = compute_subproof_hashes(subproofs)
    expected_hrich = compute_HRICH_from_subproof_hashes(subproof_hashes)

    assert HRICH == expected_hrich
    assert isinstance(HRICH, str)
    assert len(HRICH) == 64  # SHA-256 hex digest length


def test_full_rich_bundle_generation():
    """Test complete rich bundle generation."""
    src = """
fn multiply(x: Int) -> Int;

pipeline main(Int) -> Int { multiply }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=7,
        fn_registry={"multiply": lambda x: x * 3},
    )

    crypto = RLangBoRCrypto(proof)
    bundle = crypto.to_rich_bundle()

    assert isinstance(bundle, RichProofBundle)
    assert "H_RICH" in bundle.rich
    assert bundle.rich["H_RICH"] == bundle.H_RICH

    # Verify primary data structure
    assert "primary" in bundle.rich
    primary = bundle.rich["primary"]
    assert "master" in primary
    assert "steps" in primary
    assert primary["version"] == "v0"
    assert primary["language"] == "rlang"
    assert primary["entry_pipeline"] == "main"

    # Verify steps in primary
    assert len(primary["steps"]) == 1
    step_data = primary["steps"][0]
    assert "index" in step_data
    assert "template_id" in step_data
    assert "hash" in step_data


def test_rich_bundle_structure():
    """Test that rich bundle has correct structure."""
    src = """
fn add(x: Int) -> Int;
fn square(x: Int) -> Int;

pipeline main(Int) -> Int { add -> square }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=3,
        fn_registry={"add": lambda x: x + 2, "square": lambda x: x * x},
    )

    crypto = RLangBoRCrypto(proof)
    bundle = crypto.to_rich_bundle()

    # Verify bundle structure
    assert hasattr(bundle, "primary")
    assert hasattr(bundle, "H_RICH")
    assert hasattr(bundle, "rich")

    # Verify rich data contains H_RICH and primary
    assert bundle.rich["H_RICH"] == bundle.H_RICH
    assert bundle.rich["primary"] == bundle.primary

    # Verify primary contains master and steps
    assert "master" in bundle.primary
    assert "steps" in bundle.primary
    assert len(bundle.primary["steps"]) == 2


def test_determinism():
    """Test that same input produces identical rich bundle JSON."""
    src = """
fn inc(x: Int) -> Int;

pipeline main(Int) -> Int { inc }
"""
    registry = {"inc": lambda x: x + 1}

    # Generate proof twice
    proof1 = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry=registry,
    )

    proof2 = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry=registry,
    )

    # Generate rich bundles
    crypto1 = RLangBoRCrypto(proof1)
    crypto2 = RLangBoRCrypto(proof2)

    bundle1 = crypto1.to_rich_bundle()
    bundle2 = crypto2.to_rich_bundle()

    # Verify deterministic JSON
    json1 = bundle1.to_json()
    json2 = bundle2.to_json()

    assert json1 == json2

    # Verify data structures are identical
    data1 = json.loads(json1)
    data2 = json.loads(json2)

    assert data1 == data2
    assert data1["H_RICH"] == data2["H_RICH"]
    assert data1["primary"]["master"] == data2["primary"]["master"]


def test_hashed_program_structure():
    """Test HashedProgram structure."""
    src = """
fn test(x: Int) -> Int;

pipeline main(Int) -> Int { test }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=5,
        fn_registry={"test": lambda x: x * 2},
    )

    crypto = RLangBoRCrypto(proof)
    hashed = crypto.to_hashed_program()

    assert hashed.HMASTER is not None
    assert hashed.HRICH is not None
    assert len(hashed.step_hashes) == 1
    assert hashed.primary_data["master"] == hashed.HMASTER
    assert hashed.rich_data["H_RICH"] == hashed.HRICH


def test_empty_steps_error():
    """Test that empty proof bundle raises error."""
    from rlang.bor.proofs import PipelineProofBundle
    from rlang.canonical import PrimaryProgramIR

    # Create proof bundle with no steps
    ir = PrimaryProgramIR(
        version="v0",
        language="rlang",
        entry_pipeline="main",
        step_templates=[],
        pipelines=[],
    )

    proof = PipelineProofBundle(
        version="v0",
        language="rlang",
        entry_pipeline="main",
        program_ir=ir,
        input_value=1,
        output_value=1,
        steps=[],
    )

    crypto = RLangBoRCrypto(proof)

    with pytest.raises(ValueError) as exc_info:
        crypto.to_rich_bundle()

    assert "no steps" in str(exc_info.value).lower()


def test_step_hash_record_structure():
    """Test StepHashRecord structure."""
    src = """
fn double(x: Int) -> Int;

pipeline main(Int) -> Int { double }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=21,
        fn_registry={"double": lambda x: x * 2},
    )

    crypto = RLangBoRCrypto(proof)
    step_records = crypto.compute_step_hashes()

    assert len(step_records) == 1
    record = step_records[0]

    assert record.index == 0
    assert record.template_id.startswith("fn:")
    assert isinstance(record.step_hash, str)
    assert len(record.step_hash) == 64  # SHA-256 hex


def test_primary_data_structure():
    """Test that primary data has correct structure."""
    src = """
fn add(x: Int) -> Int;

pipeline main(Int) -> Int { add }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry={"add": lambda x: x + 5},
        version="v1.0",
        language="rlang-v2",
    )

    crypto = RLangBoRCrypto(proof)
    step_records = crypto.compute_step_hashes()
    HMASTER = crypto.compute_HMASTER(step_records)
    primary_data = crypto.build_primary_data(HMASTER, step_records)

    assert primary_data["master"] == HMASTER
    assert primary_data["version"] == "v1.0"
    assert primary_data["language"] == "rlang-v2"
    assert primary_data["entry_pipeline"] == "main"
    assert len(primary_data["steps"]) == 1

    step_data = primary_data["steps"][0]
    assert step_data["index"] == 0
    assert "template_id" in step_data
    assert "hash" in step_data


def test_complex_pipeline_hashing():
    """Test hashing for complex multi-step pipeline."""
    src = """
fn add_one(x: Int) -> Int;
fn multiply(x: Int) -> Int;
fn subtract(x: Int) -> Int;

pipeline main(Int) -> Int { add_one -> multiply -> subtract }
"""
    registry = {
        "add_one": lambda x: x + 1,
        "multiply": lambda x: x * 3,
        "subtract": lambda x: x - 2,
    }

    proof = run_program_with_proof(
        source=src,
        input_value=5,
        fn_registry=registry,
    )

    crypto = RLangBoRCrypto(proof)
    bundle = crypto.to_rich_bundle()

    # Verify all steps are hashed
    assert len(bundle.primary["steps"]) == 3

    # Verify HMASTER and HRICH are computed
    assert bundle.primary["master"] is not None
    assert bundle.H_RICH is not None

    # Verify step hashes are unique
    step_hashes = [s["hash"] for s in bundle.primary["steps"]]
    assert len(set(step_hashes)) == 3  # All unique


def test_rich_bundle_json_serialization():
    """Test that rich bundle can be serialized to JSON."""
    src = """
fn identity(x: Int) -> Int;

pipeline main(Int) -> Int { identity }
"""
    proof = run_program_with_proof(
        source=src,
        input_value=100,
        fn_registry={"identity": lambda x: x},
    )

    crypto = RLangBoRCrypto(proof)
    bundle = crypto.to_rich_bundle()

    # Serialize to JSON
    json_str = bundle.to_json()
    assert isinstance(json_str, str)

    # Parse back
    data = json.loads(json_str)
    assert "H_RICH" in data
    assert "primary" in data
    assert data["primary"]["master"] is not None


def test_hrich_differs_for_different_branches():
    """Test that HRICH differs when different branches are taken."""
    # Program 1: condition true, then-branch
    src_then = """
fn inc(x: Int) -> Int;
fn dec(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (1 == 1) {
    inc
  } else {
    dec
  }
}
"""

    # Program 2: logically inverted condition (force else)
    src_else = """
fn inc(x: Int) -> Int;
fn dec(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (1 == 2) {
    inc
  } else {
    dec
  }
}
"""

    fn_registry = {
        "inc": lambda x: x + 1,
        "dec": lambda x: x - 1,
    }

    # Run both programs with same input
    proof_then = run_program_with_proof(
        source=src_then,
        input_value=10,
        fn_registry=fn_registry,
    )

    proof_else = run_program_with_proof(
        source=src_else,
        input_value=10,
        fn_registry=fn_registry,
    )

    # Build rich bundles
    crypto_then = RLangBoRCrypto(proof_then)
    rich_then = crypto_then.to_rich_bundle()

    crypto_else = RLangBoRCrypto(proof_else)
    rich_else = crypto_else.to_rich_bundle()

    # Extract H_RICH
    H_RICH_then = rich_then.rich["H_RICH"]
    H_RICH_else = rich_else.rich["H_RICH"]

    # Verify branch paths
    assert rich_then.rich["primary"]["branches"][0]["path"] == "then"
    assert rich_else.rich["primary"]["branches"][0]["path"] == "else"

    # Verify HRICH differs
    assert H_RICH_then != H_RICH_else


def test_trp_subproof_contains_branches():
    """Test that TRP subproof contains branch information."""
    src = """
fn inc(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (1 == 1) {
    inc
  }
}
"""

    proof = run_program_with_proof(
        source=src,
        input_value=5,
        fn_registry={"inc": lambda x: x + 1},
    )

    crypto = RLangBoRCrypto(proof)
    rich = crypto.to_rich_bundle()

    # Verify TRP exists and contains branches
    assert "TRP" in rich.rich["subproofs"]
    assert "branches" in rich.rich["subproofs"]["TRP"]
    assert len(rich.rich["subproofs"]["TRP"]["branches"]) == 1
    assert rich.rich["subproofs"]["TRP"]["branches"][0]["path"] == "then"


def test_no_if_program_has_empty_branches_and_same_core_structure():
    """Test that programs without if have empty branches and same core structure."""
    src = """
fn inc(x: Int) -> Int;

pipeline main(Int) -> Int {
  inc
}
"""

    proof = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry={"inc": lambda x: x + 1},
    )

    crypto = RLangBoRCrypto(proof)
    rich = crypto.to_rich_bundle()

    # Verify branches field exists and is empty
    assert "branches" in rich.rich["primary"]
    assert rich.rich["primary"]["branches"] == []

    # Verify TRP has branches field and is empty
    assert "branches" in rich.rich["subproofs"]["TRP"]
    assert rich.rich["subproofs"]["TRP"]["branches"] == []

    # Verify all previously expected keys are still present
    assert "master" in rich.rich["primary"]
    assert "steps" in rich.rich["primary"]
    assert "version" in rich.rich["primary"]
    assert "language" in rich.rich["primary"]
    assert "entry_pipeline" in rich.rich["primary"]

    # Verify TRP structure
    assert "hash" in rich.rich["subproofs"]["TRP"]
    assert "verified" in rich.rich["subproofs"]["TRP"]
    assert "steps" in rich.rich["subproofs"]["TRP"]


def test_branch_tamper_detection():
    """Test that tampering with branch data is detected via HRICH mismatch."""
    src = """
fn inc(x: Int) -> Int;
fn dec(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (1 == 1) {
    inc
  } else {
    dec
  }
}
"""

    proof = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry={"inc": lambda x: x + 1, "dec": lambda x: x - 1},
    )

    crypto = RLangBoRCrypto(proof)
    rich_original = crypto.to_rich_bundle()
    H_RICH_original = rich_original.rich["H_RICH"]

    # Tamper with branch path
    rich_tampered = rich_original.rich.copy()
    rich_tampered["primary"] = rich_original.rich["primary"].copy()
    rich_tampered["primary"]["branches"] = rich_original.rich["primary"]["branches"].copy()
    rich_tampered["primary"]["branches"][0] = rich_original.rich["primary"]["branches"][0].copy()
    rich_tampered["primary"]["branches"][0]["path"] = "else"  # Changed from "then"

    # Recompute HRICH from tampered data
    # We need to recompute TRP hash and then HRICH
    rich_tampered["subproofs"] = rich_original.rich["subproofs"].copy()
    rich_tampered["subproofs"]["TRP"] = rich_original.rich["subproofs"]["TRP"].copy()
    rich_tampered["subproofs"]["TRP"]["branches"] = rich_tampered["primary"]["branches"].copy()

    # Recompute TRP hash
    from rlang.bor.crypto import compute_subproof_hashes, compute_HRICH_from_subproof_hashes
    trp_hash_tampered = compute_subproof_hashes({"TRP": rich_tampered["subproofs"]["TRP"]})["TRP"]

    # Recompute HRICH
    rich_tampered["subproof_hashes"] = rich_original.rich["subproof_hashes"].copy()
    rich_tampered["subproof_hashes"]["TRP"] = trp_hash_tampered
    H_RICH_tampered = compute_HRICH_from_subproof_hashes(rich_tampered["subproof_hashes"])

    # Verify HRICH differs
    assert H_RICH_original != H_RICH_tampered
