import importlib.util
import json
from pathlib import Path
import sys

import pytest

pytest.importorskip("openenv")
from fastapi.testclient import TestClient
from pydantic import ValidationError

from hessian_recipes.agent_rl_tasks import GermanAgentEnv, generate_tasks, score_episode

ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location("hessian_openenv_under_test", ROOT / "examples/agent_gyms/openenv_server.py")
native = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = native
spec.loader.exec_module(native)


@pytest.fixture
def env():
    instance = native.GermanOpenEnv()
    try:
        instance.reset(task=generate_tasks("train", 1, seed=20261020)[0]["task"])
        yield instance
    finally:
        instance.close()


def call(name, arguments):
    return {"role": "assistant", "content": "", "tool_calls": [
        {"type": "function", "function": {"name": name, "arguments": arguments}}]}


def solve(env, row):
    arguments = {"record_id": row["task"]["record_id"]}
    looked_up = env.step(native.GermanAction(operation="tool", tool_name="lookup", arguments=arguments))
    facts = json.loads(looked_up.observation)
    expression = f'{facts["quantity"]}*{facts["unit_price"]}'
    calculated = env.step(native.GermanAction(operation="tool", tool_name="calculator", arguments={"expression": expression}))
    result = json.loads(calculated.observation)["result"]
    transcript = [call("lookup", arguments), {"role": "tool", "content": looked_up.observation},
                  call("calculator", {"expression": expression}), {"role": "tool", "content": calculated.observation},
                  {"role": "assistant", "content": f"<answer>{result}</answer>"}]
    scored = env.step(native.GermanAction(operation="verify", completion=transcript))
    return transcript, scored


def test_native_identity_binds_actual_sdk_and_locked_requirements():
    identity = native.identity()
    assert identity["backend"] == "openenv" and identity["upstream_revision"] == native.UPSTREAM_REVISION
    assert identity["runtime"]["version"] == "0.9.0"
    assert len(identity["runtime"]["python_source_sha256"]) >= 100
    assert identity["runtime"]["python_source_tree_sha256"]


def test_reset_and_state_reveal_no_private_facts_and_state_is_copy(env):
    row = generate_tasks("train", 1, seed=20261020)[0]
    initial = env.reset(task=row["task"])
    state = env.state
    assert initial.observation == "" and initial.score is None
    assert state.episode_id == row["id"]
    assert set(state.model_dump()) == {"episode_id", "step_count", "instance_id", "done"}
    assert not any(key in json.dumps(initial.model_dump() | state.model_dump())
                   for key in ["private_facts", "quantity", "unit_price", "expected_answer"])
    state.done = True
    assert not env.state.done


def test_native_wrapper_uses_exact_strict_reward_and_trace(env):
    row = generate_tasks("train", 1, seed=20261020)[0]
    transcript, final = solve(env, row)
    reference = GermanAgentEnv()
    reference.reset(task=row["task"])
    for message in transcript:
        for action in message.get("tool_calls", []):
            function = action["function"]
            getattr(reference, function["name"])(**function["arguments"])
    assert final.reward == 1.0 and final.done and final.score == score_episode(reference, transcript)
    assert env.state.done and env.state.step_count == 3
    assert env.step(native.GermanAction(operation="verify", completion=transcript)).score == final.score
    with pytest.raises(ValueError, match="replaced"):
        env.step(native.GermanAction(operation="verify", completion="<answer>0</answer>"))
    with pytest.raises(ValueError, match="mutate"):
        env.step(native.GermanAction(operation="tool", tool_name="lookup", arguments={"record_id": row["task"]["record_id"]}))


@pytest.mark.parametrize("name,arguments", [
    ("unknown", {}), ("lookup", {"record_id": 12}), ("lookup", {"record_id": []}),
    ("lookup", {"record_id": "wrong", "extra": 1}),
    ("lookup", {"tool": "calculator", "expression": "2*3"}),
    ("calculator", {"expression": "2*3"}),
])
def test_invalid_model_actions_are_scored_errors_not_schema_infrastructure(env, name, arguments):
    response = env.step(native.GermanAction(operation="tool", tool_name=name, arguments=arguments))
    assert "error" in json.loads(response.observation)
    final = env.step(native.GermanAction(operation="verify", completion=[call(name, arguments),
                              {"role": "assistant", "content": "<answer>6</answer>"}]))
    assert final.done and final.reward == 0.0 and not final.score["success"]


@pytest.mark.parametrize("expression", [123, "2*" + "3" * 130, "__import__('os')", "1/0"])
def test_bad_calculator_values_after_lookup_use_unchanged_environment_errors(env, expression):
    row = generate_tasks("train", 1, seed=20261020)[0]
    env.step(native.GermanAction(operation="tool", tool_name="lookup", arguments={"record_id": row["task"]["record_id"]}))
    result = env.step(native.GermanAction(operation="tool", tool_name="calculator", arguments={"expression": expression}))
    assert json.loads(result.observation)["error"] == "invalid_expression"


def test_failed_reset_discards_previous_private_task_and_later_reset_recovers(env):
    row = generate_tasks("train", 1, seed=20261020)[0]
    with pytest.raises(ValueError):
        env.reset(task={"schema_version": 1})
    with pytest.raises(RuntimeError, match="Reset"):
        env.step(native.GermanAction(operation="tool", tool_name="lookup", arguments={"record_id": row["task"]["record_id"]}))
    assert env.state.episode_id is None
    env.reset(task=row["task"])
    assert solve(env, row)[1].reward == 1.0


def test_explicit_sdk_episode_identity_changes_per_reset_without_changing_task(env):
    row = generate_tasks("train", 1, seed=20261020)[0]
    instance = env.state.instance_id
    for episode in ["reset-one", "reset-two"]:
        env.reset(task=row["task"], episode_id=episode)
        assert env.state.episode_id == episode and env.state.instance_id == instance
        assert solve(env, row)[1].reward == 1.0


def test_wrong_final_with_correct_tools_and_literal_gold_shortcut_fail(env):
    row = generate_tasks("train", 1, seed=20261020)[0]
    env.step(native.GermanAction(operation="tool", tool_name="lookup", arguments={"record_id": row["task"]["record_id"]}))
    facts = row["task"]["private_facts"]
    shortcut = env.step(native.GermanAction(operation="tool", tool_name="calculator", arguments={"expression": str(facts["quantity"] * facts["unit_price"])}))
    assert json.loads(shortcut.observation)["error"] == "must_multiply_retrieved_values"
    assert env.step(native.GermanAction(operation="verify", completion="<answer>6</answer>")).reward == 0.0


def test_envelope_validation_is_strict_but_inner_argument_values_are_not_coerced():
    assert native.GermanAction(operation="tool", tool_name="lookup", arguments={"record_id": 123}).arguments["record_id"] == 123
    with pytest.raises(ValidationError):
        native.GermanAction(operation="tool", tool_name="lookup", unexpected=True)


@pytest.mark.parametrize("raw", ['{"type":"state","type":"reset"}', '{"type":"state","data":NaN}',
                                 '{"type":"state","data":Infinity}'])
def test_actual_native_websocket_rejects_duplicate_or_nonfinite_json_without_state_change(raw):
    with TestClient(native.make_app(4)) as client:
        with client.websocket_connect("/ws") as socket:
            socket.send_text(raw)
            error = socket.receive_json()
            assert error["type"] == "error" and error["data"]["code"] == "INVALID_JSON"
            socket.send_json({"type": "state"})
            state = socket.receive_json()["data"]
            assert state["episode_id"] is None and state["step_count"] == 0


@pytest.mark.parametrize("raw", [b'{"task":NaN}', b'{"task":{},"task":{}}', b'{"task":"' + b'x' * native.MAX_PAYLOAD_BYTES + b'"}'])
def test_native_http_controller_json_is_bounded_finite_and_unambiguous(raw):
    with TestClient(native.make_app(4)) as client:
        response = client.post("/reset", content=raw, headers={"content-type": "application/json"})
        assert response.status_code == 400


def test_actual_isolated_native_child_ignores_foreign_python_overlay(tmp_path, monkeypatch):
    helper_spec = importlib.util.spec_from_file_location("hessian_openenv_launcher_test", ROOT / "examples/agent_gyms/openenv_smoke.py")
    helper = importlib.util.module_from_spec(helper_spec)
    sys.modules[helper_spec.name] = helper
    helper_spec.loader.exec_module(helper)
    overlay = tmp_path / "foreign-python311"
    overlay.mkdir()
    for name in ["json.py", "openenv.py"]:
        (overlay / name).write_text('raise RuntimeError("Foreign Python ABI overlay inherited")\n')
    monkeypatch.setenv("PYTHONPATH", str(overlay))
    monkeypatch.setenv("PYTHONHOME", str(overlay))
    with helper.launch_openenv_server(sys.executable, tmp_path / "native-service", max_sessions=2) as service:
        assert service["identity"]["runtime"]["python_version"] == native.runtime_identity()["python_version"]
        assert service["identity"]["upstream_revision"] == native.UPSTREAM_REVISION
    assert service["released"]