"""
Unit tests for the DispatchAgent assignment-failure context gatherer.

Focus is the same as the ExceptionAgent tests: the gatherer must degrade
gracefully (no coords, a failing GEORADIUS, a failing counter) and never raise.
The real `_find_zone` is exercised through an in-memory fake Redis; the agent is
built with __new__ so NATS / the message bus are never touched.

Run:
    python -m unittest discover -s tests
"""
import unittest

import nats.js.errors

from agents.dispatch_agent import DispatchAgent


class FakeRedis:
    """decode_responses=True style. `geo` is the GEORADIUS result for every call.
    `available` / `on_break` are miler ids with a miler_status:<id> key saying so.
    Any geo member in neither set has NO status key at all — the common case in
    production, which must read as unknown presence, not off-duty."""
    def __init__(self, geo=None, counter_start=0, raise_on=(), available=None,
                 on_break=(), alerted=False):
        self._geo = list(geo) if geo is not None else []
        self._counter = counter_start
        self._raise_on = set(raise_on)
        self._available = set(available) if available is not None else set(self._geo)
        self._on_break = set(on_break)
        self._kv = {}
        if alerted:
            self._kv["__alerted__"] = "1"
        self.incr_calls = []
        self.expire_calls = []
        self.set_calls = []

    async def get(self, key):
        if key.startswith("miler_status:"):
            mid = key.split(":", 1)[1]
            if mid in self._available:
                return '{"userid": %s, "status": "Available"}' % (mid if mid.isdigit() else 0)
            if mid in self._on_break:
                return '{"userid": 0, "status": "Break"}'
            return None          # no status key — unknown, not off-duty
        return self._kv.get(key)

    async def exists(self, key):
        if key.startswith("assignment_alert_sent:"):
            return 1 if "__alerted__" in self._kv or key in self._kv else 0
        return 1 if key in self._kv else 0

    async def set(self, key, value, ex=None, nx=False):
        self.set_calls.append((key, ex, nx))
        if nx and key in self._kv:
            return None
        self._kv[key] = value
        return True

    async def georadius(self, key, lon, lat, radius, unit, sort=None, count=None):
        if "georadius" in self._raise_on:
            raise RuntimeError("geo down")
        return list(self._geo)

    async def hgetall(self, key):
        return {"hub_id": "H1", "zone_id": "z1", "avg_delivery_time": "60"}

    async def incr(self, key):
        if "incr" in self._raise_on:
            raise RuntimeError("redis down")
        self._counter += 1
        self.incr_calls.append(key)
        return self._counter

    async def expire(self, key, ttl):
        self.expire_calls.append((key, ttl))
        return True


def make_agent(redis):
    agent = DispatchAgent.__new__(DispatchAgent)  # skip heavy __init__ / message-bus registration
    agent._redis = redis
    return agent


class FakeJS:
    """Stand-in JetStream context for binding tests."""
    def __init__(self, stream_name=None, not_found=False):
        self._stream = stream_name
        self._not_found = not_found
        self.subscribe_calls = []
        self.add_stream_calls = []

    async def find_stream_name_by_subject(self, subject):
        if self._not_found:
            raise nats.js.errors.NotFoundError()
        return self._stream

    async def subscribe(self, subject, durable=None, stream=None, cb=None):
        self.subscribe_calls.append((subject, durable, stream))
        return object()  # stand-in subscription

    async def add_stream(self, **kwargs):
        self.add_stream_calls.append(kwargs)


async def _cb(msg):
    pass


def make_agent_js(js):
    agent = DispatchAgent.__new__(DispatchAgent)
    agent._nats_js = js
    return agent


class TestGatherAssignmentFacts(unittest.IsolatedAsyncioTestCase):
    async def test_rider_nearby(self):
        redis = FakeRedis(geo=["m5"], counter_start=0)   # found on the first (10km) sweep
        facts, count = await make_agent(redis)._gather_assignment_facts("hyderabad", 17.4, 78.4)
        self.assertEqual(facts["zone_id"], "hyderabad")
        self.assertTrue(facts["has_coordinates"])
        self.assertEqual(facts["nearest_miler_within_km"], 10)
        self.assertEqual(facts["nearest_miler_presence"], "available")
        self.assertEqual(facts["milers_in_geo_index_within_30km"], 1)
        self.assertFalse(facts["alert_already_sent_today"])
        self.assertEqual(facts["failures_today"], 1)
        self.assertEqual(count, 1)
        self.assertEqual(redis.expire_calls[0][1], 172800)

    async def test_no_rider_within_30km(self):
        redis = FakeRedis(geo=[], counter_start=2)   # every sweep empty
        facts, count = await make_agent(redis)._gather_assignment_facts("pune", 18.5, 73.8)
        self.assertEqual(facts["nearest_miler_within_km"], "none within 30km")
        self.assertEqual(facts["milers_in_geo_index_within_30km"], 0)
        self.assertEqual(count, 3)

    async def test_no_coordinates(self):
        redis = FakeRedis(geo=["m1"])
        facts, count = await make_agent(redis)._gather_assignment_facts("unknown", None, None)
        self.assertFalse(facts["has_coordinates"])
        self.assertEqual(facts["nearest_miler_within_km"], "unknown (no coordinates)")
        self.assertNotIn("milers_in_geo_index_within_30km", facts)
        self.assertEqual(count, 1)              # counter still incremented
        self.assertEqual(redis.incr_calls and 1, 1)

    async def test_georadius_error_is_swallowed(self):
        redis = FakeRedis(raise_on={"georadius"})
        facts, count = await make_agent(redis)._gather_assignment_facts("z", 1.0, 2.0)  # must not raise
        self.assertEqual(facts["nearest_miler_within_km"], "none within 30km")
        self.assertEqual(facts["failures_today"], 1)

    async def test_counter_error_is_swallowed(self):
        redis = FakeRedis(geo=["m9"], raise_on={"incr"})
        facts, count = await make_agent(redis)._gather_assignment_facts("z", 1.0, 2.0)  # must not raise
        self.assertEqual(facts["nearest_miler_within_km"], 10)   # coverage still gathered
        self.assertEqual(facts["failures_today"], 0)             # counter degraded to 0
        self.assertEqual(count, 0)


class TestPresence(unittest.IsolatedAsyncioTestCase):
    """Presence is three-state. The prod incident had 26 riders in the geo index;
    only 2 of 34 milers had a status key at all, so "no key" must read as unknown
    — reporting it as off-duty would be a confident wrong answer."""

    async def test_confirmed_available_rider_is_preferred(self):
        redis = FakeRedis(geo=["4", "7", "12"], available=["12"], on_break=["4", "7"])
        found = await make_agent(redis)._find_zone(11.0, 76.9, radius_km=10)
        self.assertEqual(found["miler_id"], "12")
        self.assertEqual(found["presence"], "available")

    async def test_no_status_key_reads_as_unknown_not_unavailable(self):
        redis = FakeRedis(geo=["4", "7"], available=[], on_break=[])   # no status keys at all
        agent = make_agent(redis)
        self.assertEqual(await agent._miler_presence("4"), "unknown")
        found = await agent._find_zone(11.0, 76.9, radius_km=10)
        self.assertIsNotNone(found)                       # still reported — may be assignable
        self.assertEqual(found["presence"], "unknown")

    async def test_all_confirmed_off_duty_returns_none(self):
        redis = FakeRedis(geo=["4", "7"], available=[], on_break=["4", "7"])
        self.assertIsNone(await make_agent(redis)._find_zone(11.0, 76.9, radius_km=10))

    async def test_facts_report_presence_unknown_with_rider_nearby(self):
        redis = FakeRedis(geo=["4", "7", "12"], available=[], on_break=[])
        facts, _ = await make_agent(redis)._gather_assignment_facts("coimbatore", 11.0, 76.9)
        self.assertEqual(facts["nearest_miler_within_km"], 10)
        self.assertEqual(facts["nearest_miler_presence"], "unknown")
        self.assertEqual(facts["milers_in_geo_index_within_30km"], 3)

    async def test_facts_report_none_found_when_all_off_duty(self):
        redis = FakeRedis(geo=["4"], available=[], on_break=["4"])
        facts, _ = await make_agent(redis)._gather_assignment_facts("coimbatore", 11.0, 76.9)
        self.assertEqual(facts["nearest_miler_within_km"], "none within 30km")
        self.assertEqual(facts["nearest_miler_presence"], "none found")
        self.assertEqual(facts["milers_in_geo_index_within_30km"], 1)   # index still knows them


class FakeMessages:
    def __init__(self):
        self.sent = []

    async def send_message(self, recipient, message_type, payload, correlation_id=None):
        self.sent.append((recipient, message_type, payload))


def make_handler_agent(redis, decision):
    """Agent with send_message captured and the LLM stubbed to a fixed decision (or None)."""
    from agents import dispatch_agent as mod
    agent = DispatchAgent.__new__(DispatchAgent)
    agent._redis = redis
    agent.agent_id = "DISPATCH_AGENT"
    box = FakeMessages()
    agent.send_message = box.send_message
    calls = []
    async def fake_decide(context):
        calls.append(context)
        return decision
    mod.decide_assignment_failure = fake_decide
    return agent, box, calls


class FakeMsg:
    def __init__(self, payload):
        import json
        self.data = json.dumps(payload).encode()
        self.acked = False

    async def ack(self):
        self.acked = True


class TestRateLimitAndSinks(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        from agents import dispatch_agent as mod
        self._orig = mod.decide_assignment_failure

    def tearDown(self):
        from agents import dispatch_agent as mod
        mod.decide_assignment_failure = self._orig

    async def test_ops_alert_goes_to_jarvis_as_exception_and_marks_zone(self):
        from core.types import MessageType
        from core.llm import AssignmentDecision
        redis = FakeRedis(geo=[], counter_start=2)
        agent, box, calls = make_handler_agent(redis, AssignmentDecision("ops_alert", "gap", 0.9))
        msg = FakeMsg({"booking_id": 1, "lat": 11.0, "lon": 76.9})
        await agent._on_nats_booking_assignment_failed(msg)
        self.assertTrue(msg.acked)
        self.assertEqual(len(calls), 1)                                  # LLM consulted once
        recipient, mtype, payload = box.sent[0]
        self.assertEqual(recipient, "JARVIS")                            # not CUSTOMER_AGENT
        self.assertEqual(mtype, MessageType.EXCEPTION_DETECTED)          # the path JARVIS handles
        self.assertEqual(payload["exception_type"], "coverage_gap")
        self.assertTrue(any(k.startswith("assignment_alert_sent:") for k, _, _ in redis.set_calls))

    async def test_burst_after_alert_skips_llm_and_realerts_every_n(self):
        from agents import dispatch_agent as mod
        from core.llm import AssignmentDecision
        redis = FakeRedis(geo=[], counter_start=0, alerted=True)          # zone already alerted today
        agent, box, calls = make_handler_agent(redis, AssignmentDecision("ops_alert", "gap", 0.9))
        old = mod.DISPATCH_REALERT_EVERY
        mod.DISPATCH_REALERT_EVERY = 100
        try:
            for i in range(1, 251):
                await agent._on_nats_booking_assignment_failed(FakeMsg({"booking_id": i, "lat": 11.0, "lon": 76.9}))
        finally:
            mod.DISPATCH_REALERT_EVERY = old
        self.assertEqual(calls, [])                 # zero LLM calls during the burst
        self.assertEqual(len(box.sent), 2)          # re-alerts at 100 and 200 only
        self.assertIn("200 failed assignments", box.sent[1][2]["reasoning"])

    async def test_llm_unavailable_falls_back_to_heuristic(self):
        redis = FakeRedis(geo=[], counter_start=2)
        agent, box, calls = make_handler_agent(redis, None)               # LLM down
        await agent._on_nats_booking_assignment_failed(FakeMsg({"booking_id": 5, "lat": 1.0, "lon": 2.0}))
        self.assertEqual(len(box.sent), 1)          # count reached 3 → heuristic alert
        self.assertEqual(box.sent[0][0], "JARVIS")

    async def test_escalate_goes_to_jarvis_as_exception(self):
        from core.types import MessageType
        from core.llm import AssignmentDecision
        redis = FakeRedis(geo=["4"], available=["4"], counter_start=4)
        agent, box, calls = make_handler_agent(redis, AssignmentDecision("escalate", "riders near, still failing", 0.8))
        await agent._on_nats_booking_assignment_failed(FakeMsg({"booking_id": 9, "lat": 1.0, "lon": 2.0}))
        recipient, mtype, payload = box.sent[0]
        self.assertEqual((recipient, mtype), ("JARVIS", MessageType.EXCEPTION_DETECTED))
        self.assertEqual(payload["proposed_action"], "escalate")


class TestZoneKey(unittest.IsolatedAsyncioTestCase):
    """The backend sends no zone_id today; without a fallback every failure in
    the country shares one bucket and the daily alert limit becomes global."""

    async def test_real_zone_id_wins(self):
        from agents.dispatch_agent import zone_key
        self.assertEqual(zone_key("hyderabad", 17.4, 78.4), "hyderabad")

    async def test_missing_zone_falls_back_to_geo_cell(self):
        from agents.dispatch_agent import zone_key
        self.assertEqual(zone_key(None, 11.0182714, 76.9677744), "geo:11.0,77.0")
        self.assertEqual(zone_key("unknown", 11.0182714, 76.9677744), "geo:11.0,77.0")

    async def test_distant_pickups_get_different_buckets(self):
        from agents.dispatch_agent import zone_key
        self.assertNotEqual(zone_key(None, 11.01, 76.96), zone_key(None, 17.44, 78.39))

    async def test_no_zone_and_no_coords_is_unknown(self):
        from agents.dispatch_agent import zone_key
        self.assertEqual(zone_key(None, None, None), "unknown")
        self.assertEqual(zone_key("unknown", "bad", "coords"), "unknown")


class TestBindConsumer(unittest.IsolatedAsyncioTestCase):
    async def test_binds_to_discovered_stream_without_creating(self):
        js = FakeJS(stream_name="TRACKING")
        sub = await make_agent_js(js)._bind_consumer("booking.assigned", "d1", _cb)
        self.assertIsNotNone(sub)
        # Bound explicitly on the discovered stream...
        self.assertEqual(js.subscribe_calls, [("booking.assigned", "d1", "TRACKING")])
        # ...and never tried to create a stream (the old overlap bug).
        self.assertEqual(js.add_stream_calls, [])

    async def test_returns_none_when_no_stream_carries_subject(self):
        js = FakeJS(not_found=True)
        sub = await make_agent_js(js)._bind_consumer("booking.assigned", "d1", _cb)
        self.assertIsNone(sub)              # clear failure, not a raise
        self.assertEqual(js.subscribe_calls, [])  # did not attempt to subscribe


if __name__ == "__main__":
    unittest.main()