"""Comprehensive tests for RLang execution proof bundle generation."""

import json

import pytest

from rlang.bor import PipelineProofBundle, StepExecutionRecord, run_program_with_proof


def test_single_step_proof():
    """Test proof generation for single-step pipeline."""
    src = """
fn inc(x: Int) -> Int;

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

    assert isinstance(bundle, PipelineProofBundle)
    assert bundle.entry_pipeline == "main"
    assert bundle.input_value == 10
    assert bundle.output_value == 11
    assert len(bundle.steps) == 1

    step = bundle.steps[0]
    assert step.index == 0
    assert step.step_name == "inc"
    assert step.input_snapshot == 10
    assert step.output_snapshot == 11


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

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

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

    assert bundle.output_value == 12  # (5 + 1) * 2
    assert len(bundle.steps) == 2
    assert [s.step_name for s in bundle.steps] == ["inc", "double"]
    assert bundle.steps[0].output_snapshot == 6
    assert bundle.steps[1].input_snapshot == 6


def test_mock_function_fallback():
    """Test that mock functions are used when function not in registry."""
    src = """
fn step1(x: Int) -> Int;

pipeline main(Int) -> Int { step1 }
"""
    bundle = run_program_with_proof(
        source=src,
        input_value=42,
        fn_registry=None,
    )

    assert len(bundle.steps) == 1
    step = bundle.steps[0]
    assert step.step_name == "step1"
    # Default mock returns {"mock": "step1", "input": 42}
    assert isinstance(step.output_snapshot, dict)
    assert step.output_snapshot["mock"] == "step1"
    assert step.output_snapshot["input"] == 42
    assert bundle.output_value == step.output_snapshot


def test_proof_json_determinism():
    """Test that proof JSON is deterministic."""
    src = """
fn inc(x: Int) -> Int;

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

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

    json1 = bundle1.to_json()
    json2 = bundle2.to_json()

    assert json1 == json2

    data = json.loads(json1)
    assert data["input"] == 7
    assert data["output"] == 8
    assert len(data["steps"]) == 1


def test_no_entry_pipeline_error():
    """Test that missing entry pipeline raises ValueError."""
    src = "fn f(x: Int) -> Int;"

    with pytest.raises(ValueError) as exc_info:
        run_program_with_proof(source=src, input_value=1)

    assert "No entry pipeline" in str(exc_info.value)


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

pipeline main(Int) -> Int { add }
"""
    bundle = run_program_with_proof(
        source=src,
        input_value=5,
        fn_registry={"add": lambda x: x + 3},
    )

    assert bundle.version == "v0"
    assert bundle.language == "rlang"
    assert bundle.entry_pipeline == "main"
    assert bundle.input_value == 5
    assert bundle.output_value == 8
    assert len(bundle.steps) == 1

    # Verify program IR is included
    assert bundle.program_ir is not None
    assert bundle.program_ir.entry_pipeline == "main"


def test_step_execution_record_structure():
    """Test that step execution records have correct structure."""
    src = """
fn multiply(x: Int) -> Int;

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

    step = bundle.steps[0]
    assert isinstance(step, StepExecutionRecord)
    assert step.index == 0
    assert step.step_name == "multiply"
    assert step.template_id.startswith("fn:")
    assert step.input_snapshot == 6
    assert step.output_snapshot == 42

    # Verify to_dict works
    step_dict = step.to_dict()
    assert step_dict["index"] == 0
    assert step_dict["step_name"] == "multiply"
    assert step_dict["input"] == 6
    assert step_dict["output"] == 42


def test_proof_bundle_to_dict():
    """Test that proof bundle to_dict() works correctly."""
    src = """
fn square(x: Int) -> Int;

pipeline main(Int) -> Int { square }
"""
    bundle = run_program_with_proof(
        source=src,
        input_value=9,
        fn_registry={"square": lambda x: x * x},
    )

    bundle_dict = bundle.to_dict()

    assert bundle_dict["version"] == "v0"
    assert bundle_dict["language"] == "rlang"
    assert bundle_dict["entry_pipeline"] == "main"
    assert bundle_dict["input"] == 9
    assert bundle_dict["output"] == 81
    assert len(bundle_dict["steps"]) == 1
    assert "program" in bundle_dict


def test_proof_bundle_to_json():
    """Test that proof bundle to_json() produces valid JSON."""
    src = """
fn identity(x: Int) -> Int;

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

    json_str = bundle.to_json()
    assert isinstance(json_str, str)

    # Should be valid JSON
    data = json.loads(json_str)
    assert data["input"] == 100
    assert data["output"] == 100


def test_mixed_mock_and_real_functions():
    """Test proof generation with mix of mock and real functions."""
    src = """
fn real_fn(x: Int) -> Int;
fn mock_fn(x: Int) -> Int;

pipeline main(Int) -> Int { real_fn -> mock_fn }
"""
    bundle = run_program_with_proof(
        source=src,
        input_value=10,
        fn_registry={"real_fn": lambda x: x * 2},
    )

    assert len(bundle.steps) == 2

    # First step uses real function
    assert bundle.steps[0].step_name == "real_fn"
    assert bundle.steps[0].output_snapshot == 20

    # Second step uses mock function
    assert bundle.steps[1].step_name == "mock_fn"
    assert bundle.steps[1].input_snapshot == 20
    assert isinstance(bundle.steps[1].output_snapshot, dict)
    assert bundle.steps[1].output_snapshot["mock"] == "mock_fn"
    assert bundle.steps[1].output_snapshot["input"] == 20


def test_complex_pipeline_proof():
    """Test proof generation 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,
    }

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

    assert len(bundle.steps) == 3
    assert bundle.input_value == 5
    assert bundle.output_value == 16  # (5 + 1) * 3 - 2 = 16

    # Verify step-by-step execution
    assert bundle.steps[0].input_snapshot == 5
    assert bundle.steps[0].output_snapshot == 6
    assert bundle.steps[1].input_snapshot == 6
    assert bundle.steps[1].output_snapshot == 18
    assert bundle.steps[2].input_snapshot == 18
    assert bundle.steps[2].output_snapshot == 16


def test_custom_version_and_language():
    """Test that custom version and language are preserved in proof bundle."""
    src = """
fn test(x: Int) -> Int;

pipeline main(Int) -> Int { test }
"""
    bundle = run_program_with_proof(
        source=src,
        input_value=1,
        fn_registry={"test": lambda x: x},
        version="v1.0",
        language="rlang-v2",
    )

    assert bundle.version == "v1.0"
    assert bundle.language == "rlang-v2"

    bundle_dict = bundle.to_dict()
    assert bundle_dict["version"] == "v1.0"
    assert bundle_dict["language"] == "rlang-v2"


def test_proof_bundle_immutability():
    """Test that proof bundle and records are immutable (frozen dataclasses)."""
    src = """
fn test(x: Int) -> Int;

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

    # Should not be able to modify frozen dataclass
    with pytest.raises(Exception):  # dataclass.FrozenInstanceError
        bundle.input_value = 999

    with pytest.raises(Exception):  # dataclass.FrozenInstanceError
        bundle.steps[0].input_snapshot = 999


def test_template_id_in_records():
    """Test that step execution records include correct template_id."""
    src = """
fn step1(x: Int) -> Int;
fn step2(x: Int) -> Int;

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

    assert len(bundle.steps) == 2
    assert bundle.steps[0].template_id.startswith("fn:")
    assert bundle.steps[1].template_id.startswith("fn:")
    assert bundle.steps[0].template_id != bundle.steps[1].template_id

    # Verify template_id is in dict representation
    step_dict = bundle.steps[0].to_dict()
    assert "template_id" in step_dict
    assert step_dict["template_id"].startswith("fn:")


def test_run_program_with_if_then_branch_only():
    """Test runtime execution with if expression taking then branch."""
    source = """
fn double(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (1 == 1) {
    double
  }
}
"""
    fn_registry = {"double": lambda x: x * 2}
    bundle = run_program_with_proof(
        source=source,
        input_value=10,
        fn_registry=fn_registry,
    )

    assert bundle.output_value == 20
    assert len(bundle.steps) == 1
    assert bundle.steps[0].step_name == "double"
    assert bundle.steps[0].input_snapshot == 10
    assert bundle.steps[0].output_snapshot == 20

    # Check branch record
    assert len(bundle.branches) == 1
    assert bundle.branches[0].path == "then"
    assert bundle.branches[0].condition_value is True


def test_run_program_with_if_else_false_branch():
    """Test runtime execution with if expression taking else branch."""
    source = """
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,
    }
    bundle = run_program_with_proof(
        source=source,
        input_value=10,
        fn_registry=fn_registry,
    )

    assert bundle.output_value == 9  # false condition → else branch (dec)
    assert len(bundle.steps) == 1
    assert bundle.steps[0].step_name == "dec"
    assert bundle.steps[0].input_snapshot == 10
    assert bundle.steps[0].output_snapshot == 9

    # Check branch record
    assert len(bundle.branches) == 1
    assert bundle.branches[0].path == "else"
    assert bundle.branches[0].condition_value is False


def test_run_program_with_if_both_branches_same_type():
    """Test runtime execution with if expression having both branches."""
    source = """
fn double(x: Int) -> Int;
fn half(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (5 > 3) {
    double
  } else {
    half
  }
}
"""
    fn_registry = {
        "double": lambda x: x * 2,
        "half": lambda x: x // 2,
    }
    bundle = run_program_with_proof(
        source=source,
        input_value=10,
        fn_registry=fn_registry,
    )

    # Condition is True, so then branch executes
    assert bundle.output_value == 20
    assert len(bundle.steps) == 1
    assert bundle.steps[0].step_name == "double"
    assert len(bundle.branches) == 1
    assert bundle.branches[0].path == "then"
    assert bundle.branches[0].condition_value is True


def test_if_expr_branch_trace_in_json():
    """Test that branch trace appears in proof bundle JSON."""
    source = """
fn double(x: Int) -> Int;

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

    bundle_dict = bundle.to_dict()
    assert "branches" in bundle_dict
    assert len(bundle_dict["branches"]) == 1
    branch = bundle_dict["branches"][0]
    assert branch["path"] == "then"
    assert branch["condition_value"] is True
    assert "index" in branch

def test_if_expr_with_value_identifier_runtime():
    """Test runtime execution with __value identifier in condition."""
    source = """
fn double(x: Int) -> Int;
fn half(x: Int) -> Int;

pipeline main(Int) -> Int {
  if (__value > 10) {
    double
  } else {
    half
  }
}
"""
    fn_registry = {
        "double": lambda x: x * 2,
        "half": lambda x: x // 2,
    }

    # Test with value > 10 (takes then branch)
    bundle1 = run_program_with_proof(
        source=source,
        input_value=20,
        fn_registry=fn_registry,
    )
    assert bundle1.output_value == 40  # 20 * 2
    assert bundle1.branches[0].path == "then"
    assert bundle1.branches[0].condition_value is True

    # Test with value <= 10 (takes else branch)
    bundle2 = run_program_with_proof(
        source=source,
        input_value=5,
        fn_registry=fn_registry,
    )
    assert bundle2.output_value == 2  # 5 // 2
    assert bundle2.branches[0].path == "else"
    assert bundle2.branches[0].condition_value is False
