"""Tests for Lab 01 — MCP server/client and the bank's tool estate.

    pytest test_lab.py -v
    LAB_MODULE=solution pytest test_lab.py -v    # the reference — must be green
"""

import importlib
import os

import pytest

lab = importlib.import_module(os.environ.get("LAB_MODULE", "lab"))


# ======================================================================================
# 1. JSON-RPC envelope
# ======================================================================================


def test_request_has_id_notification_does_not():
    assert "id" in lab.make_request(1, "ping")
    assert "id" not in lab.make_notification("notifications/initialized")


def test_params_are_omitted_when_absent():
    assert "params" not in lab.make_request(1, "ping")
    assert lab.make_request(1, "tools/call", {"name": "x"})["params"] == {"name": "x"}


@pytest.mark.parametrize("message", [
    {"jsonrpc": "1.0", "method": "ping"},
    {"method": "ping"},
    {"jsonrpc": "2.0"},
    {"jsonrpc": "2.0", "method": ""},
    {"jsonrpc": "2.0", "method": "ping", "params": []},
    {"jsonrpc": "2.0", "method": "ping", "id": 1.5},
])
def test_invalid_envelopes_are_rejected(message):
    with pytest.raises(lab.JsonRpcError) as exc:
        lab.validate_envelope(message)
    assert exc.value.code == lab.INVALID_REQUEST


def test_valid_envelope_passes():
    lab.validate_envelope({"jsonrpc": "2.0", "method": "ping", "id": 1, "params": {}})
    lab.validate_envelope({"jsonrpc": "2.0", "method": "ping", "id": "abc"})


# ======================================================================================
# 2. Schema validation
# ======================================================================================


SCHEMA = {
    "type": "object",
    "properties": {
        "reference": {"type": "string", "pattern": r"^PMT-\d{3,}$", "minLength": 7},
        "amount": {"type": "integer", "minimum": 1, "maximum": 1_000_000},
        "currency": {"type": "string", "enum": ["AED", "USD"]},
        "legs": {"type": "array", "items": {"type": "string"}, "maxItems": 3},
        "urgent": {"type": "boolean"},
    },
    "required": ["reference", "amount"],
    "additionalProperties": False,
}


def test_valid_object_produces_no_errors():
    assert lab.validate_schema({"reference": "PMT-771", "amount": 10}, SCHEMA) == []


def test_missing_required_is_reported_with_a_path():
    errors = lab.validate_schema({"amount": 1}, SCHEMA)
    assert [e.path for e in errors] == ["$.reference"]
    assert "required" in errors[0].message


def test_all_errors_are_returned_not_just_the_first():
    """A repair loop needs every problem in one turn."""
    errors = lab.validate_schema({"amount": "x", "currency": "GBP", "extra": 1}, SCHEMA)
    paths = {e.path for e in errors}
    assert {"$.reference", "$.amount", "$.currency", "$.extra"} <= paths


def test_errors_are_sorted_deterministically():
    errors = lab.validate_schema({"currency": "GBP", "amount": "x"}, SCHEMA)
    assert [(e.path, e.message) for e in errors] == sorted((e.path, e.message) for e in errors)


def test_boolean_is_not_an_integer():
    """bool subclasses int in Python; a schema saying integer must still reject True."""
    errors = lab.validate_schema({"reference": "PMT-771", "amount": True}, SCHEMA)
    assert any(e.path == "$.amount" for e in errors)


def test_integer_is_a_valid_number():
    assert lab.validate_schema(5, {"type": "number"}) == []
    assert lab.validate_schema(5.0, {"type": "integer"}) != []


def test_type_mismatch_suppresses_downstream_noise():
    """Once the type is wrong, range and pattern errors would be meaningless."""
    errors = lab.validate_schema({"reference": 12345, "amount": 1}, SCHEMA)
    assert len([e for e in errors if e.path == "$.reference"]) == 1


def test_numeric_bounds():
    assert lab.validate_schema({"reference": "PMT-771", "amount": 0}, SCHEMA)
    assert lab.validate_schema({"reference": "PMT-771", "amount": 1_000_001}, SCHEMA)
    assert lab.validate_schema({"reference": "PMT-771", "amount": 1}, SCHEMA) == []
    assert lab.validate_schema({"reference": "PMT-771", "amount": 1_000_000}, SCHEMA) == []


def test_exclusive_bounds():
    assert lab.validate_schema(5, {"type": "integer", "exclusiveMinimum": 5})
    assert lab.validate_schema(6, {"type": "integer", "exclusiveMinimum": 5}) == []
    assert lab.validate_schema(5, {"type": "integer", "exclusiveMaximum": 5})


def test_enum_and_const():
    assert lab.validate_schema("GBP", {"enum": ["AED", "USD"]})
    assert lab.validate_schema("AED", {"enum": ["AED", "USD"]}) == []
    assert lab.validate_schema("b", {"const": "a"})


def test_string_length_and_pattern():
    assert lab.validate_schema({"reference": "PMT-77", "amount": 1}, SCHEMA)
    assert lab.validate_schema({"reference": "PMT-777", "amount": 1}, SCHEMA) == []


def test_array_items_are_validated_with_indexed_paths():
    errors = lab.validate_schema(
        {"reference": "PMT-771", "amount": 1, "legs": ["a", 2, "c"]}, SCHEMA)
    assert [e.path for e in errors] == ["$.legs[1]"]


def test_array_bounds_and_uniqueness():
    assert lab.validate_schema({"reference": "PMT-771", "amount": 1,
                                "legs": ["a", "b", "c", "d"]}, SCHEMA)
    assert lab.validate_schema([1, 1], {"type": "array", "uniqueItems": True})
    assert lab.validate_schema([1, 2], {"type": "array", "uniqueItems": True}) == []


def test_nested_objects():
    schema = {"type": "object",
              "properties": {"a": {"type": "object",
                                   "properties": {"b": {"type": "integer"}},
                                   "required": ["b"]}}}
    assert [e.path for e in lab.validate_schema({"a": {}}, schema)] == ["$.a.b"]


def test_unsupported_type_is_a_programming_error():
    with pytest.raises(ValueError):
        lab.validate_schema(1, {"type": "quaternion"})


def test_empty_schema_accepts_anything():
    assert lab.validate_schema({"anything": [1, 2]}, {}) == []


# ======================================================================================
# 3. Versioning
# ======================================================================================


def test_version_parse_and_str():
    assert str(lab.Version.parse("2.13.4")) == "2.13.4"
    with pytest.raises(ValueError):
        lab.Version.parse("2.13")
    with pytest.raises(ValueError):
        lab.Version.parse("v2.13.4")


def test_version_ordering():
    assert lab.Version(1, 2, 3) < lab.Version(1, 10, 0)
    assert lab.Version(2, 0, 0) > lab.Version(1, 99, 99)


@pytest.mark.parametrize("version,constraint,expected", [
    ("1.2.3", "*", True),
    ("1.2.3", "1.2.3", True),
    ("1.2.4", "1.2.3", False),
    ("1.9.0", "^1.2.3", True),
    ("2.0.0", "^1.2.3", False),
    ("1.2.0", "^1.2.3", False),
    ("1.2.9", "~1.2.3", True),
    ("1.3.0", "~1.2.3", False),
])
def test_constraints(version, constraint, expected):
    assert lab.Version.parse(version).satisfies(constraint) is expected


BASE = {"type": "object",
        "properties": {"a": {"type": "string"}, "b": {"type": "integer", "minimum": 0}},
        "required": ["a"]}


def test_patch_when_nothing_structural_changed():
    same = {"type": "object",
            "properties": {"a": {"type": "string"}, "b": {"type": "integer", "minimum": 0}},
            "required": ["a"]}
    assert lab.classify_schema_change(BASE, same) is lab.SchemaChange.PATCH


def test_adding_an_optional_property_is_minor():
    new = dict(BASE, properties=dict(BASE["properties"], c={"type": "string"}))
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MINOR


def test_adding_a_required_property_is_major():
    new = dict(BASE, required=["a", "b"])
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MAJOR


def test_removing_a_required_property_is_only_minor():
    """Relaxation never breaks a caller."""
    new = dict(BASE, required=[])
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MINOR


def test_removing_a_property_is_major():
    new = dict(BASE, properties={"a": {"type": "string"}})
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MAJOR


def test_changing_a_type_is_major():
    new = dict(BASE, properties=dict(BASE["properties"], a={"type": "integer"}))
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MAJOR


def test_narrowing_an_enum_is_major_widening_is_minor():
    old = {"type": "object", "properties": {"c": {"type": "string", "enum": ["A", "B"]}}}
    narrow = {"type": "object", "properties": {"c": {"type": "string", "enum": ["A"]}}}
    wide = {"type": "object", "properties": {"c": {"type": "string", "enum": ["A", "B", "C"]}}}
    assert lab.classify_schema_change(old, narrow) is lab.SchemaChange.MAJOR
    assert lab.classify_schema_change(old, wide) is lab.SchemaChange.MINOR


def test_tightening_a_bound_is_major():
    new = dict(BASE, properties=dict(BASE["properties"],
                                     b={"type": "integer", "minimum": 5}))
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MAJOR


def test_closing_additional_properties_is_major():
    new = dict(BASE, additionalProperties=False)
    assert lab.classify_schema_change(BASE, new) is lab.SchemaChange.MAJOR


# ======================================================================================
# 4. The registry
# ======================================================================================


def spec(name="t", version="1.0.0", schema=None, **kwargs):
    kwargs.setdefault("side_effect", lab.SideEffect.READ)
    return lab.ToolSpec(
        name=name, version=lab.Version.parse(version), title=name,
        description=f"does {name}", input_schema=schema if schema is not None else BASE,
        **kwargs)


def test_publish_and_latest():
    reg = lab.ToolRegistry()
    reg.publish(spec())
    assert reg.latest("t").key == "t@1.0.0"
    assert reg.latest("missing") is None


def test_republishing_a_version_is_refused():
    reg = lab.ToolRegistry()
    reg.publish(spec())
    with pytest.raises(ValueError):
        reg.publish(spec())


def test_publishing_an_older_version_is_refused():
    reg = lab.ToolRegistry()
    reg.publish(spec(version="2.0.0"))
    with pytest.raises(ValueError):
        reg.publish(spec(version="1.5.0"))


def test_registry_enforces_the_version_bump_the_schema_change_requires():
    reg = lab.ToolRegistry()
    reg.publish(spec())
    breaking = dict(BASE, required=["a", "b"])
    with pytest.raises(ValueError, match="major"):
        reg.publish(spec(version="1.1.0", schema=breaking))
    reg.publish(spec(version="2.0.0", schema=breaking))
    assert reg.latest("t").key == "t@2.0.0"


def test_registry_enforces_a_minor_bump_for_a_relaxation():
    reg = lab.ToolRegistry()
    reg.publish(spec())
    relaxed = dict(BASE, properties=dict(BASE["properties"], c={"type": "string"}))
    with pytest.raises(ValueError):
        reg.publish(spec(version="1.0.1", schema=relaxed))   # patch is not enough
    reg.publish(spec(version="1.1.0", schema=relaxed))


def test_deprecated_is_not_latest_but_is_still_resolvable():
    reg = lab.ToolRegistry()
    reg.publish(spec())
    reg.publish(spec(version="2.0.0", schema=dict(BASE, required=["a", "b"])))
    reg.deprecate("t", "1.0.0", superseded_by="t@2.0.0")
    assert reg.latest("t").key == "t@2.0.0"
    assert reg.resolve("t", "^1.0.0").key == "t@1.0.0"
    assert reg.versions("t")[0].superseded_by == "t@2.0.0"


def test_retired_is_not_resolvable():
    reg = lab.ToolRegistry()
    reg.publish(spec())
    reg.retire("t", "1.0.0")
    assert reg.resolve("t", "*") is None
    assert reg.latest("t") is None


def test_lifecycle_change_on_a_missing_version_raises():
    reg = lab.ToolRegistry()
    with pytest.raises(KeyError):
        reg.deprecate("t", "1.0.0")


# ======================================================================================
# 5. Authorization-aware discovery
# ======================================================================================


def estate():
    reg = lab.ToolRegistry()
    reg.publish(spec(name="read.public", data_classification="public"))
    reg.publish(spec(name="read.scoped", required_scopes=("s1",)))
    reg.publish(spec(name="read.restricted", data_classification="restricted",
                     required_scopes=("s1",)))
    reg.publish(spec(name="wholesale.only", tenants=("wholesale",)))
    return reg


def test_discovery_filters_by_scope():
    reg = estate()
    p = lab.Principal("a", "retail", scopes=(), max_classification="restricted")
    assert [s.name for s in reg.discover(p)] == ["read.public"]


def test_discovery_filters_by_classification():
    reg = estate()
    p = lab.Principal("a", "retail", scopes=("s1",), max_classification="internal")
    assert "read.restricted" not in [s.name for s in reg.discover(p)]
    p2 = lab.Principal("a", "retail", scopes=("s1",), max_classification="restricted")
    assert "read.restricted" in [s.name for s in reg.discover(p2)]


def test_discovery_filters_by_tenant():
    reg = estate()
    retail = lab.Principal("a", "retail", scopes=("s1",), max_classification="restricted")
    wholesale = lab.Principal("a", "wholesale", scopes=("s1",), max_classification="restricted")
    assert "wholesale.only" not in [s.name for s in reg.discover(retail)]
    assert "wholesale.only" in [s.name for s in reg.discover(wholesale)]


def test_discovery_returns_only_the_newest_visible_version():
    reg = lab.ToolRegistry()
    reg.publish(spec(name="t"))
    reg.publish(spec(name="t", version="2.0.0", schema=dict(BASE, required=["a", "b"])))
    p = lab.Principal("a", "retail", max_classification="restricted")
    found = reg.discover(p)
    assert len(found) == 1
    assert found[0].version == lab.Version(2, 0, 0)


def test_discovery_excludes_deprecated_unless_asked():
    reg = lab.ToolRegistry()
    reg.publish(spec(name="t"))
    reg.deprecate("t", "1.0.0")
    p = lab.Principal("a", "retail", max_classification="restricted")
    assert reg.discover(p) == []
    assert [s.key for s in reg.discover(p, include_deprecated=True)] == ["t@1.0.0"]


def test_discovery_is_deterministic():
    reg = estate()
    p = lab.Principal("a", "wholesale", scopes=("s1",), max_classification="restricted")
    once = [s.key for s in reg.discover(p)]
    assert once == sorted(once)
    assert once == [s.key for s in reg.discover(p)]


def test_unknown_classification_is_an_error():
    with pytest.raises(ValueError):
        lab.classification_rank("cosmic")


# ======================================================================================
# 6. Server and client
# ======================================================================================


def build(principal=None, handlers=None, registry=None, protocol_versions=None):
    registry = registry if registry is not None else estate()
    principal = principal or lab.Principal("agent", "wholesale", scopes=("s1",),
                                           max_classification="restricted")
    handlers = handlers if handlers is not None else {
        name: (lambda a, n=name: lab.ToolCallResult(f"ran {n}"))
        for name in ("read.public", "read.scoped", "read.restricted", "wholesale.only")
    }
    server = lab.MCPServer(
        name="test-server", version="1.0.0", registry=registry, handlers=handlers,
        principal=principal,
        resources=[lab.Resource("doc://a", "A", "text/plain", "alpha"),
                   lab.Resource("doc://w", "W", "text/plain", "wholesale only",
                                tenants=("wholesale",))],
        prompts=[lab.Prompt("greet", "greeting", ("who",), "Hello {who}")],
        protocol_versions=protocol_versions or (lab.PROTOCOL_VERSION,),
    )
    return server, lab.MCPClient(server)


def test_initialize_negotiates_and_reports_capabilities():
    server, client = build()
    result = client.initialize()
    assert result["protocolVersion"] == lab.PROTOCOL_VERSION
    assert result["serverInfo"]["name"] == "test-server"
    assert client.supports("tools")
    assert server.initialized


def test_unsupported_protocol_version_falls_back_instead_of_failing():
    server, client = build(protocol_versions=("2024-11-05", "2025-06-18"))
    result = client.initialize("1999-01-01")
    assert result["protocolVersion"] == "2025-06-18"


def test_methods_before_initialize_are_refused():
    server, _ = build()
    response = server.handle(lab.make_request(1, "tools/list"))
    assert response["error"]["code"] == lab.INVALID_REQUEST


def test_unknown_method_is_method_not_found():
    server, client = build()
    client.initialize()
    response = server.handle(lab.make_request(2, "tools/teleport"))
    assert response["error"]["code"] == lab.METHOD_NOT_FOUND


def test_notifications_are_never_answered():
    server, _ = build()
    assert server.handle(lab.make_notification("notifications/initialized")) is None
    assert server.handle(lab.make_notification("notifications/tools/list_changed")) is None


def test_tools_list_returns_only_discoverable_tools():
    server, client = build(principal=lab.Principal("a", "retail", scopes=(),
                                                    max_classification="internal"))
    client.initialize()
    assert [t["name"] for t in client.list_tools()] == ["read.public"]


def test_tools_list_does_not_leak_platform_metadata():
    """Scopes, classification, owner and tenants are the platform's, not the model's."""
    server, client = build()
    client.initialize()
    tool = client.list_tools()[0]
    assert set(tool) == {"name", "title", "description", "inputSchema", "_meta"}
    assert "requiredScopes" not in tool
    assert "dataClassification" not in tool


def test_calling_an_undiscoverable_tool_is_indistinguishable_from_missing():
    server, client = build(principal=lab.Principal("a", "retail", scopes=(),
                                                    max_classification="internal"))
    client.initialize()
    with pytest.raises(lab.JsonRpcError) as unauthorized:
        client.call_tool("read.scoped", {"a": "x"})
    with pytest.raises(lab.JsonRpcError) as missing:
        client.call_tool("does.not.exist", {"a": "x"})
    assert unauthorized.value.code == missing.value.code == lab.METHOD_NOT_FOUND


def test_schema_violation_never_reaches_the_handler():
    calls = []
    server, client = build(handlers={"read.public": lambda a: calls.append(a) or
                                     lab.ToolCallResult("ok")})
    client.initialize()
    with pytest.raises(lab.JsonRpcError) as exc:
        client.call_tool("read.public", {"a": 123})
    assert exc.value.code == lab.INVALID_PARAMS
    assert exc.value.data["errors"]
    assert calls == []


def test_a_tool_failure_is_a_successful_response_with_iserror():
    server, client = build(handlers={"read.public":
                                     lambda a: lab.ToolCallResult("downstream 503", is_error=True)})
    client.initialize()
    result = client.call_tool("read.public", {"a": "x"})
    assert result["isError"] is True
    assert result["content"][0]["text"] == "downstream 503"


def test_successful_call_records_the_arguments():
    server, client = build()
    client.initialize()
    client.call_tool("read.public", {"a": "x"})
    assert server.call_log == [("read.public", {"a": "x"})]


def test_tools_call_validates_its_own_params():
    server, client = build()
    client.initialize()
    for params in ({}, {"name": 1}, {"name": "read.public", "arguments": []}):
        response = server.handle(lab.make_request(9, "tools/call", params))
        assert response["error"]["code"] == lab.INVALID_PARAMS


def test_resources_are_tenant_scoped():
    server, client = build()
    client.initialize()
    assert client.read_resource("doc://w") == "wholesale only"

    server2, client2 = build(principal=lab.Principal("a", "retail",
                                                     max_classification="restricted"))
    client2.initialize()
    uris = [r["uri"] for r in client2._send("resources/list")["resources"]]
    assert uris == ["doc://a"]
    with pytest.raises(lab.JsonRpcError):
        client2.read_resource("doc://w")


def test_prompts_interpolate_and_validate_arguments():
    server, client = build()
    client.initialize()
    assert client.get_prompt("greet", {"who": "Layla"}) == "Hello Layla"
    with pytest.raises(lab.JsonRpcError):
        client.get_prompt("greet", {})
    with pytest.raises(lab.JsonRpcError):
        client.get_prompt("nope", {})


def test_client_caches_the_tool_list_until_told_otherwise():
    server, client = build()
    client.initialize()
    client.list_tools()
    client.list_tools()
    assert client.list_calls == 1
    client.on_notification(server.notify_tools_changed())
    client.list_tools()
    assert client.list_calls == 2


def test_list_changed_notification_is_recorded_on_the_server():
    server, client = build()
    client.initialize()
    message = server.notify_tools_changed()
    assert message["method"] == "notifications/tools/list_changed"
    assert "id" not in message
    assert server.outbound == [message]


def test_ping_works_after_initialize():
    server, client = build()
    client.initialize()
    assert client._send("ping") == {}


# ======================================================================================
# 7. The repair loop
# ======================================================================================


REPAIR_SCHEMA = {
    "type": "object",
    "properties": {
        "amount": {"type": "integer"},
        "rate": {"type": "number"},
        "reference": {"type": "string"},
        "urgent": {"type": "boolean"},
    },
    "required": ["amount", "reference"],
    "additionalProperties": False,
}


def test_repair_coerces_unambiguous_types():
    repaired, errors = lab.repair_arguments(
        {"amount": "42", "rate": "1.5", "reference": 771, "urgent": "true"}, REPAIR_SCHEMA)
    assert repaired == {"amount": 42, "rate": 1.5, "reference": "771", "urgent": True}
    assert errors == []


def test_repair_drops_disallowed_additional_properties():
    repaired, errors = lab.repair_arguments(
        {"amount": 1, "reference": "x", "nonsense": True}, REPAIR_SCHEMA)
    assert "nonsense" not in repaired
    assert errors == []


def test_repair_fills_required_fields_only_from_defaults():
    repaired, errors = lab.repair_arguments({"amount": 1}, REPAIR_SCHEMA,
                                            defaults={"reference": "PMT-000"})
    assert repaired["reference"] == "PMT-000"
    assert errors == []


def test_repair_does_not_invent_values():
    """Guessing a missing account number is fabrication, not repair."""
    repaired, errors = lab.repair_arguments({"amount": 1}, REPAIR_SCHEMA)
    assert "reference" not in repaired
    assert [e.path for e in errors] == ["$.reference"]


def test_repair_leaves_uncoercible_values_for_the_model():
    repaired, errors = lab.repair_arguments({"amount": "not a number", "reference": "x"},
                                            REPAIR_SCHEMA)
    assert repaired["amount"] == "not a number"
    assert [e.path for e in errors] == ["$.amount"]


def test_repair_is_idempotent():
    once, _ = lab.repair_arguments({"amount": "42", "reference": 771}, REPAIR_SCHEMA)
    twice, _ = lab.repair_arguments(once, REPAIR_SCHEMA)
    assert once == twice


# ======================================================================================
# 8. Side-effect classes
# ======================================================================================


def test_retryability_is_derived_from_the_side_effect_class():
    assert lab.RETRYABLE[lab.SideEffect.READ] is True
    assert lab.RETRYABLE[lab.SideEffect.WRITE_IDEMPOTENT] is True
    assert lab.RETRYABLE[lab.SideEffect.WRITE_NON_IDEMPOTENT] is False
    assert lab.RETRYABLE[lab.SideEffect.IRREVERSIBLE] is False


def test_every_side_effect_class_has_a_retry_policy():
    assert set(lab.RETRYABLE) == set(lab.SideEffect)