"""Tests for the two-call boundary and prefix-quality policy."""

from __future__ import annotations

import json

import pytest

from swe_data_processing.audit import (
    compute_prefix_quality,
    derive_prefix_safety,
    effective_boundary_policy,
    materialize_prefix,
    validate_boundary,
    validate_prefix_quality,
)
from swe_data_processing.policy import PolicyViolation
from swe_data_processing.workflow import (
    prepare_boundary_payload,
    prepare_prefix_quality_payload,
)


def _record() -> dict:
    return {
        "trajectory_id": "sample-1",
        "resolved": 0,
        "tools": [],
        "trajectory": [
            {"role": "system", "content": "system"},
            {"role": "user", "content": "Fix the bug."},
            {"role": "assistant", "content": "I will inspect the code."},
            {"role": "tool", "content": "relevant.py"},
            {"role": "assistant", "content": "I will apply the harmful patch."},
            {"role": "tool", "content": "tests failed"},
        ],
    }


def _boundary() -> dict:
    return {
        "sample_id": "sample-1",
        "checks": {
            "task_coverage": "COMPLETE",
            "final_patch_scope": "CLEAN",
            "constraints": "RESPECTED",
            "claims_vs_observations": "CONSISTENT",
        },
        "decision": "TRUNCATE",
        "truncate_before_turn": 5,
        "category": "PERSISTENT_WRONG_IMPLEMENTATION",
        "severity": "MAJOR",
        "state_effect": "UNRECOVERED",
        "evidence_turns": [5],
        "reason": "The patch fails and is not repaired.",
    }


def _quality() -> dict:
    return {
        "sample_id": "sample-1",
        "behavior_issues": [],
        "dimensions": {
            "planning": 14,
            "investigation": 15,
            "tool_use_and_observation": 14,
            "progress": 12,
            "clarity_and_efficiency": 10,
        },
        "evidence_turns": [],
        "reason": "Useful investigation.",
    }


def test_valid_boundary_materializes_exact_prefix() -> None:
    validate_boundary(_record(), _boundary())
    prefix = materialize_prefix(_record(), 5)
    assert prefix == _record()["trajectory"][:4]


def test_boundary_must_be_an_assistant_turn() -> None:
    result = _boundary()
    result["truncate_before_turn"] = 4
    with pytest.raises(PolicyViolation, match="assistant turn"):
        validate_boundary(_record(), result)


def test_boundary_evidence_must_include_boundary_turn() -> None:
    result = _boundary()
    result["evidence_turns"] = [6]
    with pytest.raises(PolicyViolation, match="excluded assistant turn"):
        validate_boundary(_record(), result)


def test_keep_full_cannot_contain_a_boundary() -> None:
    result = _boundary()
    result.update(
        decision="KEEP_FULL",
        truncate_before_turn=None,
        category="NONE",
        severity="NONE",
        state_effect="NONE",
        evidence_turns=[],
    )
    validate_boundary(_record(), result)
    result["truncate_before_turn"] = 5
    with pytest.raises(PolicyViolation, match="cannot contain a boundary"):
        validate_boundary(_record(), result)


def test_prefix_quality_cannot_reference_suffix() -> None:
    prefix = materialize_prefix(_record(), 5)
    result = _quality()
    result["evidence_turns"] = [5]
    with pytest.raises(PolicyViolation, match="outside the prefix"):
        validate_prefix_quality("sample-1", prefix, result)


def test_quality_score_and_tier_are_computed_locally() -> None:
    score = compute_prefix_quality(_quality())
    assert score["educational_quality_score"] == 65
    assert score["quality_tier"] == "MEDIUM"


def test_unrecovered_major_issue_is_locally_invalid() -> None:
    result = _quality()
    result["behavior_issues"] = [
        {
            "turn_id": 3,
            "kind": "ERROR",
            "severity": "MAJOR",
            "recovered": False,
            "reason": "The prefix leaves a known broken edit.",
        }
    ]
    result["evidence_turns"] = [3]
    validate_prefix_quality("sample-1", materialize_prefix(_record(), 5), result)
    assert derive_prefix_safety(result)["prefix_valid"] is False
    assert compute_prefix_quality(result)["quality_tier"] == "REJECT"


@pytest.mark.parametrize("resolved", [0, -1])
def test_non_success_outcome_is_capped_before_first_stateful_turn(resolved: int) -> None:
    record = _record()
    record["resolved"] = resolved
    record["tools"] = [
        {"type": "function", "function": {"name": "str_replace_editor"}}
    ]
    record["trajectory"][2]["tool_calls"] = [
        {
            "id": "call-1",
            "type": "function",
            "function": {
                "name": "str_replace_editor",
                "arguments": '{"command":"str_replace","path":"src/a.py"}',
            },
        }
    ]
    result = _boundary()
    result.update(decision="KEEP_FULL", truncate_before_turn=None)
    policy = effective_boundary_policy(record, result)
    assert policy["decision"] == "TRUNCATE"
    assert policy["truncate_before_turn"] == 3
    assert policy["source"] == "NON_SUCCESS_FIRST_STATEFUL_CAP"


def test_boundary_payload_hides_outcome_and_patch_metadata() -> None:
    record = _record()
    record["resolved"] = -1
    record["metadata"] = {
        "model_patch": {"patch": "MODEL_PATCH_SECRET"},
        "reference_patch": {"patch": "REFERENCE_PATCH_SECRET"},
    }
    serialized = json.dumps(prepare_boundary_payload(record), ensure_ascii=False)
    assert "resolved" not in serialized
    assert "MODEL_PATCH_SECRET" in serialized
    assert "REFERENCE_PATCH_SECRET" not in serialized


def test_prefix_payload_is_identical_when_only_suffix_and_labels_change() -> None:
    first = _record()
    second = _record()
    second["resolved"] = -1
    second["trajectory"][4]["content"] = "DIFFERENT SUFFIX"
    second["trajectory"][5]["content"] = "DIFFERENT TOOL RESULT"
    second["metadata"] = {"model_patch": {"patch": "DIFFERENT PATCH"}}
    first_prefix = materialize_prefix(first, 5)
    second_prefix = materialize_prefix(second, 5)
    assert prepare_prefix_quality_payload(
        first, first_prefix
    ) == prepare_prefix_quality_payload(second, second_prefix)