import asyncio
from contextlib import asynccontextmanager, contextmanager
from datetime import datetime, timedelta, timezone
from types import SimpleNamespace
from unittest.mock import AsyncMock
from uuid import uuid4

import pytest
from fastapi import HTTPException
from pydantic import ValidationError

from app.mcp.catalog import TOOLS
from app.mcp.models import BoundContext, QueryResult
from app.mcp import extended_service, mutation_service, query_service, scenario_service, analysis_service


HANDLE = "x" * 43
START = "2026-10-21T08:00:00+08:00"


def args(tool_name, **values):
    return TOOLS[tool_name].model.model_validate({"context_handle": HANDLE, **values})


def binding(**values):
    return BoundContext(context_handle=HANDLE, project_id=uuid4(), subject="actor", client_id="host",
        native_session_id="session", purpose="ordinary_chat", tool_cap=list(TOOLS),
        permission_cap=["webgis.edit", "webgis.view", "scada.clean", "scada.view", "simulation.run"],
        expires_at=datetime.now(timezone.utc) + timedelta(minutes=5), **values)


@pytest.mark.parametrize("name,values", [
    ("query_network_extent", {"bounds": [120, 30, 122, 32]}),
    ("query_network_extent", {"bounds": [121, 31, 120, 30]}),
    ("query_network_extent", {"bounds": [float("nan"), 31, 122, 32]}),
    ("calculate_demand_allocation", {"target": "network", "demand": 10, "node_ids": ["J1"]}),
    ("edit_model_element", {"request_id": str(uuid4()), "operation": "delete", "element_type": "pipe", "element_id": "P1", "dry_run": False}),
    ("edit_model_element", {"request_id": str(uuid4()), "operation": "update", "element_type": "pipe", "element_id": "P1", "properties": {"id": "other"}}),
    ("clean_scada_readings", {"request_id": str(uuid4()), "start": START, "end": "2026-10-23T08:00:00+08:00", "device_ids": ["D1"]}),
    ("run_burst_simulation", {"start_time": START, "output_node_ids": ["J1"], "burst_pipe_ids": ["P1"], "orifice_areas_cm2": [1, 2]}),
])
def test_bounded_schemas_reject_unsafe_inputs(name, values):
    with pytest.raises(ValidationError):
        args(name, **values)


def test_saved_scenario_has_no_nested_identity_and_validates_timezone():
    values = dict(request_id=uuid4(), name="test", scenario="burst", parameters={
        "start_time": START, "output_node_ids": ["J1"], "burst_pipe_ids": ["P1"], "orifice_areas_cm2": [1],
    })
    valid = args("save_simulation_scenario", **values)
    assert "context_handle" not in valid.parameters.model_dump()
    values["parameters"]["start_time"] = "2026-10-21T08:00:00"
    with pytest.raises(ValidationError):
        args("save_simulation_scenario", **values)
    values["parameters"]["start_time"] = START
    values["parameters"]["network"] = "other-project"
    with pytest.raises(ValidationError):
        args("save_simulation_scenario", **values)


@pytest.mark.parametrize("purpose,role", [("task_reader", "member"), ("task_execution", "member"), ("task_collaboration", "member"), ("ordinary_chat", "viewer")])
def test_writes_never_enter_background_or_viewer_sessions(monkeypatch, purpose, role):
    monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"webgis.edit"})
    bound = binding()
    bound.purpose = purpose
    scope = SimpleNamespace(project=SimpleNamespace(project_role=role))
    with pytest.raises(HTTPException) as error:
        asyncio.run(query_service.authorize_query(scope, bound, TOOLS["edit_model_element"], object()))
    assert error.value.status_code == 403


def test_write_requires_current_and_captured_permissions(monkeypatch):
    scope = SimpleNamespace(project=SimpleNamespace(project_role="member"))
    bound = binding()
    monkeypatch.setattr(query_service, "permissions_for_context", lambda _: set())
    with pytest.raises(HTTPException):
        asyncio.run(query_service.authorize_query(scope, bound, TOOLS["edit_model_element"], object()))
    monkeypatch.setattr(query_service, "permissions_for_context", lambda _: {"webgis.edit"})
    bound.permission_cap = []
    with pytest.raises(HTTPException):
        asyncio.run(query_service.authorize_query(scope, bound, TOOLS["edit_model_element"], object()))


def test_extent_runs_real_parameterized_project_query_and_pages(monkeypatch):
    read = AsyncMock(return_value=[{"id": "A"}, {"id": "B"}])
    monkeypatch.setattr(extended_service, "read_rows", read)
    scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="own"))
    result = asyncio.run(extended_service.execute_extended(scope, object(), TOOLS["query_network_extent"],
        args("query_network_extent", bounds=[121.0, 31.0, 121.1, 31.1], limit=1)))
    assert result.page.has_more and result.truncated
    assert result.data["items"] == [{"id": "A"}]
    call = read.call_args.args
    assert "ST_Intersects" in str(call[1])
    assert call[2] == (121.0, 31.0, 121.1, 31.1, 2, 0)
    assert call[0] is scope.connection


def test_missing_gis_propagates_unavailable_without_cross_project_retry(monkeypatch):
    read = AsyncMock(side_effect=ValueError("GIS unavailable"))
    monkeypatch.setattr(extended_service, "read_rows", read)
    scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="no-gis"))
    with pytest.raises(ValueError):
        asyncio.run(extended_service.execute_extended(scope, object(), TOOLS["query_network_extent"],
            args("query_network_extent", bounds=[121.0, 31.0, 121.1, 31.1])))
    assert read.await_count == 1


def test_demand_allocation_is_paged_and_never_sets_demands(monkeypatch):
    from app.services import demand_allocation
    monkeypatch.setattr(demand_allocation, "calculate_demand_to_nodes", lambda *a: {"B": 6, "A": 4})
    values = args("calculate_demand_allocation", target="nodes", node_ids=["A", "B"], demand=10, limit=1)
    result = asyncio.run(analysis_service.calculate("calculate_demand_allocation", values, network="own", username="actor"))
    assert result == {"items": [{"node_id": "A", "demand": 4}], "total": 2, "limit": 1, "offset": 0, "has_more": True, "persisted": False}


def test_burst_preview_uses_isolated_nonpersistent_adapter(monkeypatch):
    from app.services import simulation_scenarios
    calls = []
    monkeypatch.setattr(simulation_scenarios, "burst_analysis", lambda **kw: calls.append(kw) or {"node_results": [{"node": "J1"}], "link_results": []})
    values = args("run_burst_simulation", start_time=START, output_node_ids=["J1"], burst_pipe_ids=["P1"], orifice_areas_cm2=[2])
    result = scenario_service.calculate_scenario("run_burst_simulation", values.model_dump(exclude={"context_handle", "timeout_seconds"}), network="own", username="actor")
    assert calls[0]["persist_results"] is False and calls[0]["scheme_name"] is None
    assert calls[0]["burst_size"] == [2]
    assert result["persisted_baseline"] is False


class Connection:
    @asynccontextmanager
    async def transaction(self):
        yield

    async def commit(self):
        pass


class Journal:
    records = {}

    def __init__(self, *args):
        pass

    async def lock_receipt(self, *args):
        pass

    async def receipt(self, scope, key):
        return self.records.get((scope, key))

    async def save_receipt(self, scope, key, parameters, result):
        self.records[(scope, key)] = {"parameters": parameters, "result": result}

    async def finish_receipt(self, scope, key, result):
        self.records[(scope, key)]["result"] = result


@pytest.fixture
def mutation(monkeypatch):
    Journal.records = {}
    monkeypatch.setattr(mutation_service, "McpQueryRepository", Journal)
    preview = AsyncMock(return_value={"before": {"diameter": 100}, "properties": {"diameter": 200}})
    apply = AsyncMock(return_value={"changed": True})
    monkeypatch.setattr(mutation_service, "preview_change", preview)
    monkeypatch.setattr(mutation_service, "apply_change", apply)
    scope = SimpleNamespace(connection=Connection(), project=SimpleNamespace(project_id=uuid4()))
    values = dict(request_id=uuid4(), operation="update", element_type="pipe", element_id="P1", properties={"diameter": 200})
    return scope, binding(), values, preview, apply


def test_preview_execution_stale_state_and_idempotent_replay(mutation):
    scope, bound, values, preview, apply = mutation
    tool = TOOLS["edit_model_element"]
    async def run():
        result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
        assert result.data["dry_run"] and not apply.called
        execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
        preview.return_value = {"before": {"diameter": 150}}
        with pytest.raises(HTTPException) as error:
            await mutation_service.execute_mutation(scope, bound, tool, execution)
        assert error.value.status_code == 409 and not apply.called
        assert Journal.records == {}
        preview.return_value = {"before": {"diameter": 100}, "properties": {"diameter": 200}}
        result = await mutation_service.execute_mutation(scope, bound, tool, execution)
        replay = await mutation_service.execute_mutation(scope, bound, tool, execution)
        assert replay == result
        assert apply.await_count == 1
    asyncio.run(run())


def test_failed_or_timed_out_write_cannot_be_automatically_reexecuted(mutation):
    scope, bound, values, preview, apply = mutation
    tool = TOOLS["edit_model_element"]
    async def run():
        result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
        execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
        apply.side_effect = TimeoutError()
        with pytest.raises(TimeoutError):
            await mutation_service.execute_mutation(scope, bound, tool, execution)
        with pytest.raises(HTTPException) as error:
            await mutation_service.execute_mutation(scope, bound, tool, execution)
        assert error.value.status_code == 409
        assert apply.await_count == 1
    asyncio.run(run())


def test_capacity_rejection_is_retryable_under_the_same_request_after_slots_free(mutation, monkeypatch):
    from threading import BoundedSemaphore

    scope, bound, values, preview, apply = mutation
    tool = TOOLS["edit_model_element"]
    slots = BoundedSemaphore(2)
    monkeypatch.setattr(analysis_service, "_slots", slots)
    writes = []

    async def bounded_write(*unused):
        return await analysis_service.bounded_calculation(lambda: writes.append("written") or {"changed": True})

    apply.side_effect = bounded_write

    async def run():
        result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
        execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
        assert slots.acquire(blocking=False) and slots.acquire(blocking=False)
        try:
            with pytest.raises(HTTPException) as error:
                await mutation_service.execute_mutation(scope, bound, tool, execution)
            assert error.value.status_code == 503
            record = next(iter(Journal.records.values()))
            assert record["result"]["status"] == "rejected"
            assert record["result"]["retryable"] is True
            assert record["result"]["error"]["code"] == "CALCULATION_CAPACITY_EXCEEDED"
            assert writes == []
        finally:
            slots.release()
            slots.release()
        saved = await mutation_service.execute_mutation(scope, bound, tool, execution)
        replay = await mutation_service.execute_mutation(scope, bound, tool, execution)
        assert replay == saved and writes == ["written"]
        assert next(iter(Journal.records.values()))["result"]["status"] == "completed"
    asyncio.run(run())


@pytest.mark.parametrize("failure", [TimeoutError(), asyncio.CancelledError(), RuntimeError("native failure"), HTTPException(503, "unrelated source failure")])
def test_ambiguous_write_failure_remains_started_and_cannot_be_retried(mutation, failure):
    scope, bound, values, preview, apply = mutation
    tool = TOOLS["edit_model_element"]

    async def run():
        result = await mutation_service.execute_mutation(scope, bound, tool, args(tool.name, **values))
        execution = args(tool.name, **values, dry_run=False, expected_digest=result.data["expected_digest"])
        apply.side_effect = failure
        with pytest.raises(type(failure)):
            await mutation_service.execute_mutation(scope, bound, tool, execution)
        assert next(iter(Journal.records.values()))["result"]["status"] == "started"
        with pytest.raises(HTTPException) as error:
            await mutation_service.execute_mutation(scope, bound, tool, execution)
        assert error.value.status_code == 409 and apply.await_count == 1
    asyncio.run(run())


def test_native_model_lock_revalidates_preview_and_only_writes_allowed_fields(monkeypatch):
    from app.services import tjnetwork
    from app.native.wndb.core import database
    state = {"id": "P1", "diameter": 100}
    monkeypatch.setattr(tjnetwork, "get_pipe", lambda *a: dict(state))
    monkeypatch.setattr(tjnetwork, "get_pipe_schema", lambda *a: {"id": {"readonly": True}, "diameter": {"type": "float", "readonly": False}})
    @contextmanager
    def lock(*a):
        yield
    monkeypatch.setattr(database, "model_mutation_transaction", lock)
    writes = []
    monkeypatch.setattr(tjnetwork, "set_pipe", lambda network, changes: writes.append(changes.operations) or changes)
    values = dict(request_id=uuid4(), operation="update", element_type="pipe", element_id="P1", properties={"diameter": 200})
    preview = mutation_service.model_change("own", "edit_model_element", args("edit_model_element", **values), execute=False)
    execution = args("edit_model_element", **values, dry_run=False, expected_digest=mutation_service.digest(preview))
    state["diameter"] = 110
    with pytest.raises(HTTPException):
        mutation_service.model_change("own", "edit_model_element", execution, execute=True)
    assert not writes
    state["diameter"] = 100
    mutation_service.model_change("own", "edit_model_element", execution, execute=True)
    assert writes == [[{"id": "P1", "diameter": 200}]]


def test_artifacts_are_integrity_checked_bounded_binary_envelopes():
    result = extended_service.artifact("model.zip", b"zip-data", "application/zip")
    assert result["file"]["size"] == 8
    assert len(result["file"]["sha256"]) == 64
    with pytest.raises(HTTPException) as error:
        extended_service.artifact("model.zip", bytes(5 * 1024 * 1024 + 1), "application/zip")
    assert error.value.status_code == 413


def test_no_health_prediction_tool_and_every_write_requests_native_permission():
    assert not any("health" in name for name in TOOLS)
    assert len(TOOLS) == 67
    assert sum(t.mutating for t in TOOLS.values()) == 13
    for tool in TOOLS.values():
        if tool.mutating:
            assert not tool.protocol_tool().annotations.read_only_hint
            assert tool.protocol_tool().annotations.destructive_hint


@pytest.mark.parametrize("scenario,parameters,function", [
    ("burst", {"burst_pipe_ids": ["P1"], "orifice_areas_cm2": [1]}, "burst_analysis"),
    ("contaminant", {"source": "J2", "concentration": 1}, "contaminant_simulation"),
    ("valve_closure", {"valves": ["V1"]}, "valve_close_analysis"),
])
def test_formal_scenarios_bind_actor_and_persist_only_through_existing_services(monkeypatch, scenario, parameters, function):
    from app.services import simulation_scenarios
    calls = []
    monkeypatch.setattr(simulation_scenarios, function, lambda **kw: calls.append(kw))
    payload = {"start_time": START, "duration": 900, "output_node_ids": ["J1"], **parameters}
    values = args("save_simulation_scenario", request_id=uuid4(), name="formal-test", scenario=scenario, parameters=payload)
    mutation_service.saved_scenario("owned-project", "verified-user", values)
    assert calls[0]["name"] == "owned-project"
    assert calls[0]["username"] == "verified-user"
    assert calls[0]["persist_results"] is True
    assert calls[0]["scheme_name"] == "formal-test"


def test_mutation_status_is_scoped_to_bound_identity_session_and_project(monkeypatch):
    calls = []
    class Repository:
        def __init__(self, conn, project):
            calls.append(project)
        async def receipt(self, scope_key, request_id):
            calls.extend([scope_key, request_id])
            return {"parameters": {"tool": "edit_model_element"}, "result": {"status": "started"}}
    import app.infra.db.postgresql.mcp_queries as storage
    monkeypatch.setattr(storage, "McpQueryRepository", Repository)
    bound = binding()
    request_id = uuid4()
    scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="own", project_id=bound.project_id))
    result = asyncio.run(extended_service.execute_extended(scope, bound, TOOLS["get_mutation_status"], args("get_mutation_status", request_id=request_id)))
    assert calls == [bound.project_id, "mutation:actor:session", str(request_id)]
    assert result.data["status"] == "started"


def test_mutation_status_exposes_rejected_capacity_and_retryability(monkeypatch):
    class Repository:
        def __init__(self, *unused):
            pass
        async def receipt(self, *unused):
            return {"parameters": {"tool": "save_simulation_scenario"}, "result": {
                "status": "rejected", "retryable": True,
                "error": {"code": "CALCULATION_CAPACITY_EXCEEDED", "message": "calculation not started"},
            }}
    import app.infra.db.postgresql.mcp_queries as storage
    monkeypatch.setattr(storage, "McpQueryRepository", Repository)
    bound = binding()
    scope = SimpleNamespace(connection=object(), project=SimpleNamespace(project_code="own", project_id=bound.project_id))
    result = asyncio.run(extended_service.execute_extended(scope, bound, TOOLS["get_mutation_status"],
        args("get_mutation_status", request_id=uuid4())))
    assert result.data["status"] == "rejected" and result.data["retryable"] is True
    assert result.data["error"]["code"] == "CALCULATION_CAPACITY_EXCEEDED"


def test_scada_preview_and_journal_distinguish_omitted_fields_from_explicit_null():
    values = dict(request_id=uuid4(), operation="update", provenance="operator correction",
                  records=[{"device_id": "D1", "time": START, "monitored_value": 1}])
    original = args("write_scada_readings", **values)
    values["records"][0]["cleaned_value"] = None
    clearing = args("write_scada_readings", **values)
    for preview in [True, False]:
        left = mutation_service.mutation_arguments(original, preview=preview)
        right = mutation_service.mutation_arguments(clearing, preview=preview)
        assert "cleaned_value" not in left["records"][0]
        assert right["records"][0]["cleaned_value"] is None
        assert mutation_service.digest(left) != mutation_service.digest(right)


def test_nested_model_properties_cannot_hide_nonfinite_or_unknown_fields():
    schema = {"demands": {"type": "list", "readonly": False, "element": {"demand": {"type": "float", "readonly": False}}}}
    for value in [{"demand": float("nan")}, {"demand": 1, "url": "https://foreign"}]:
        with pytest.raises(HTTPException) as error:
            mutation_service.check_properties(schema, {"demands": [value]})
        assert error.value.status_code == 422