"""The identity resolver's second run, after the tool, for identity the tool creates.

Modelled on blender-mcp: an addon calls register_client, and only after that
call can the resolver find the bus client (bus_client_uuid, bus_role) for the
caller's session. The first run, before the tool, already knows the user.
"""

import asyncio
import contextlib
import time

from fastmcp import Client, FastMCP
from fastmcp.exceptions import ToolError

from fastmcp_feedback.instrumentation import REDACTED, MemorySink, instrument


class Bus:
    """Who is registered, as the tool leaves it; read by the resolvers."""

    def __init__(self):
        self.client = None


class BlenderResolver:
    """blender-mcp's identify(): user and role always, bus keys once registered."""

    def __init__(self, bus):
        self.bus = bus
        self.calls = 0

    def __call__(self, context):
        self.calls += 1
        out = {"user_sub": "user-1", "caller_kind": "addon", "client_id": "dcr-9"}
        if self.bus.client is not None:
            out["bus_client_uuid"] = self.bus.client
            out["bus_role"] = "interactive"
        return out


class EmptyUntilRegistered:
    """Nothing before registration, every field after."""

    def __init__(self, bus):
        self.bus = bus
        self.calls = 0

    def __call__(self, context):
        self.calls += 1
        if self.bus.client is None:
            return {}
        return {
            "user_sub": "user-1",
            "caller_kind": "addon",
            "client_id": "dcr-9",
            "bus_client_uuid": self.bus.client,
        }


def make_app(resolver_cls, *, late=(), sample=None, **kwargs):
    bus = Bus()
    resolver = resolver_cls(bus)
    sink = MemorySink()
    app = FastMCP("late-identity")
    holder = {}

    @app.tool
    async def register_client(client_uuid: str) -> dict:
        holder["mw"].record_event("bus.registering", key=client_uuid)
        bus.client = client_uuid
        return {"status": "ok"}

    @app.tool
    async def register_then_fail(client_uuid: str) -> dict:
        bus.client = client_uuid
        raise RuntimeError("registered, then broke")

    @app.tool
    async def quiet() -> str:
        return "ok"

    @app.tool
    async def ping() -> str:
        holder["mw"].record_event("bus.ping")
        return "pong"

    kwargs.setdefault("identity_resolver", resolver)
    if sample is not None:
        kwargs.update(
            sample_rates={"ping": sample, "quiet": sample}, random=lambda: 0.99
        )
    mw = instrument(app, [sink], late_identity_keys=late, **kwargs)
    holder["mw"] = mw
    return app, sink, mw, resolver


async def call(app, mw, *calls):
    async with Client(app) as client:
        for tool, args in calls:
            with contextlib.suppress(ToolError):
                await client.call_tool(tool, args)
    await mw.flush()


def columns(record):
    return (record.user_sub, record.caller_kind, record.client_id)


class TestRegistration:
    async def test_blender_shaped_registration_row_gets_bus_identity(self):
        app, sink, mw, who = make_app(BlenderResolver, late=("bus_client_uuid",))
        await call(app, mw, ("register_client", {"client_uuid": "c-1"}))
        [rec] = sink.calls
        assert columns(rec) == ("user-1", "addon", "dcr-9")
        assert rec.identity == {"bus_client_uuid": "c-1", "bus_role": "interactive"}
        assert who.calls == 2
        # The event recorded during the call saw the first answer only.
        [ev] = sink.events
        assert ev.call_id == rec.id and columns(ev) == ("user-1", "addon", "dcr-9")

    async def test_blender_shaped_without_late_keys_runs_once(self):
        # The three columns are all there before the tool, so nothing marks
        # the answer incomplete: the bus keys need late_identity_keys.
        app, sink, mw, who = make_app(BlenderResolver)
        await call(app, mw, ("register_client", {"client_uuid": "c-1"}))
        [rec] = sink.calls
        assert rec.identity == {} and who.calls == 1

    async def test_registered_caller_runs_once(self):
        app, sink, mw, who = make_app(BlenderResolver, late=("bus_client_uuid",))
        await call(
            app,
            mw,
            ("register_client", {"client_uuid": "c-1"}),
            ("ping", {}),
            ("ping", {}),
        )
        assert who.calls == 2 + 1 + 1
        assert [r.identity.get("bus_client_uuid") for r in sink.calls] == [
            "c-1",
            "c-1",
            "c-1",
        ]

    async def test_empty_first_answer_filled_after(self):
        app, sink, mw, who = make_app(EmptyUntilRegistered)
        await call(app, mw, ("register_client", {"client_uuid": "c-2"}))
        [rec] = sink.calls
        [ev] = sink.events
        assert columns(rec) == ("user-1", "addon", "dcr-9")
        assert rec.identity == {"bus_client_uuid": "c-2"}
        assert columns(ev) == (None, None, None)  # not rewritten
        assert who.calls == 2

    async def test_tool_that_raises_gets_second_run_too(self):
        app, sink, mw, who = make_app(EmptyUntilRegistered)
        await call(app, mw, ("register_then_fail", {"client_uuid": "c-3"}))
        [rec] = sink.calls
        assert rec.outcome == "error"
        assert columns(rec) == ("user-1", "addon", "dcr-9")
        assert who.calls == 2


class TestMerge:
    async def test_first_answer_wins_on_conflict(self):
        answers = iter(
            [
                {"user_sub": "early", "caller_kind": None, "org": None},
                {
                    "user_sub": "late",
                    "caller_kind": "k",
                    "client_id": "c",
                    "org": "acme",
                    "team": "qa",
                },
            ]
        )
        app, sink, mw, _ = make_app(
            EmptyUntilRegistered, identity_resolver=lambda c: next(answers)
        )
        await call(app, mw, ("ping", {}))
        [rec] = sink.calls
        [ev] = sink.events
        assert columns(rec) == ("early", "k", "c")
        assert rec.identity == {"org": "acme", "team": "qa"}
        assert columns(ev) == ("early", None, None)

    async def test_second_answer_redacted(self):
        token = "bmcp_abcdefgh12345678"
        answers = iter([{}, {"user_sub": "u", "caller_kind": f"pat:{token}"}])
        app, sink, mw, _ = make_app(
            EmptyUntilRegistered, identity_resolver=lambda c: next(answers)
        )
        await call(app, mw, ("ping", {}))
        assert sink.calls[0].caller_kind == f"pat:{REDACTED}"


class TestFailures:
    async def test_second_run_raising_keeps_first_answer(self):
        runs = []

        def who(context):
            runs.append(1)
            if len(runs) == 1:
                return {"user_sub": "u1"}
            raise RuntimeError("resolver bug")

        app, sink, mw, _ = make_app(EmptyUntilRegistered, identity_resolver=who)
        await call(app, mw, ("ping", {}))
        [rec] = sink.calls
        assert rec.ok and rec.user_sub == "u1" and len(runs) == 2

    async def test_second_run_timing_out_keeps_first_answer(self):
        runs = []

        async def who(context):
            runs.append(1)
            if len(runs) == 1:
                return {"user_sub": "u1"}
            await asyncio.sleep(5)
            return {"user_sub": "late", "caller_kind": "k"}

        app, sink, mw, _ = make_app(
            EmptyUntilRegistered, identity_resolver=who, hook_timeout=0.05
        )
        await call(app, mw, ("ping", {}))
        [rec] = sink.calls
        assert rec.ok and rec.user_sub == "u1" and rec.caller_kind is None


class TestWhenItRuns:
    async def test_complete_answer_runs_once(self):
        app, sink, mw, who = make_app(BlenderResolver)
        await call(app, mw, ("ping", {}), ("ping", {}))
        assert who.calls == 2 and len(sink.calls) == 2

    async def test_sampled_out_call_makes_no_second_run(self):
        app, sink, mw, who = make_app(EmptyUntilRegistered, sample=0.1)
        await call(app, mw, ("quiet", {}))
        assert sink.calls == [] and who.calls == 1

    async def test_sampled_call_kept_for_its_event_gets_second_run(self):
        app, sink, mw, who = make_app(EmptyUntilRegistered, sample=0.1)
        await call(app, mw, ("ping", {}))
        assert len(sink.calls) == 1 and who.calls == 2

    async def test_no_resolver_no_runs(self):
        app, sink, mw, _ = make_app(EmptyUntilRegistered, identity_resolver=None)
        await call(app, mw, ("ping", {}))
        [rec] = sink.calls
        assert columns(rec) == (None, None, None)

    async def test_second_run_overlaps_enricher(self):
        # The record is built before the result reaches the client, so the
        # second run must not add its time to the enricher's.
        runs = []

        async def who(context):
            runs.append(1)
            if len(runs) > 1:
                await asyncio.sleep(0.4)
            return {}

        async def enrich(tool, args, result, context):
            await asyncio.sleep(0.4)
            return {}

        app, sink, mw, _ = make_app(
            EmptyUntilRegistered,
            identity_resolver=who,
            enricher=enrich,
            hook_timeout=2,
        )
        async with Client(app) as client:
            start = time.perf_counter()
            await client.call_tool("ping", {})
            elapsed = time.perf_counter() - start
        assert len(runs) == 2
        assert elapsed < 0.75, elapsed


async def test_mode_off_runs_nothing():
    app, sink, mw, who = make_app(EmptyUntilRegistered, mode="off")
    await call(app, mw, ("ping", {}))
    assert who.calls == 0