import unittest
import time

from fgai.graylog_aggregate import GraylogAggregateSource


class _AggregateClient:
    def __init__(self, responses, schema=None):
        self.responses = list(responses)
        self.arguments = []
        self.schema = schema
        self.probes = 0

    def probe(self):
        self.probes += 1
        status = {"status": "connected"}
        if self.schema:
            status["tool_schemas"] = {"aggregate_messages": self.schema}
        return status

    def call_tool(self, _name, arguments):
        self.arguments.append(arguments)
        return self.responses.pop(0)


class _ProbeErrorClient:
    def probe(self):
        raise RuntimeError("connection refused")


class GraylogAggregateTests(unittest.TestCase):
    def test_reads_count_from_graylog_schema_rows(self):
        client = _AggregateClient([
            {"result": {"content": [{"type": "text", "text": '{"schema":[{"name":"metric: count()"}],"datarows":[[12345]]}'}]}}
        ])

        status = GraylogAggregateSource(client, "firewall").fetch_count(range_seconds=300)

        self.assertEqual(status["aggregate_status"], "ok")
        self.assertEqual(status["aggregate_events"], 12345)
        self.assertEqual(client.arguments[0]["streams"], ["firewall"])
        self.assertEqual(client.arguments[0]["metrics"], ["count()"])

    def test_returns_error_status_instead_of_raising_on_probe_error(self):
        status = GraylogAggregateSource(_ProbeErrorClient(), "firewall").fetch_count()
        self.assertEqual(status["aggregate_status"], "error")
        self.assertIn("probe_error", status["aggregate_error"])

    def test_uses_existing_probe_status_without_reprobing(self):
        client = _AggregateClient([
            {"result": {"content": [{"type": "text", "text": '{"events": 9}'}]}}
        ])
        status = GraylogAggregateSource(client, "firewall").fetch_count(probe_status={"status": "connected"})
        self.assertEqual(status["aggregate_events"], 9)
        self.assertEqual(client.probes, 0)

    def test_tries_fallback_argument_shape_after_tool_error(self):
        client = _AggregateClient([
            {"result": {"isError": True, "content": [{"type": "text", "text": "bad metrics"}]}},
            {"result": {"content": [{"type": "text", "text": '{"events": 42}'}]}},
        ])

        status = GraylogAggregateSource(client, "firewall").fetch_count()

        self.assertEqual(status["aggregate_status"], "ok")
        self.assertEqual(status["aggregate_events"], 42)
        self.assertEqual(len(client.arguments), 2)

    def test_fallbacks_only_use_graylog_supported_aggregate_fields(self):
        client = _AggregateClient([
            {"result": {"isError": True, "content": [{"type": "text", "text": "bad count parens"}]}},
            {"result": {"isError": True, "content": [{"type": "text", "text": "bad count"}]}},
            {"result": {"isError": True, "content": [{"type": "text", "text": "bad object metric"}]}},
        ])

        status = GraylogAggregateSource(client, "firewall").fetch_count()

        self.assertEqual(status["aggregate_status"], "error")
        self.assertEqual(len(client.arguments), 3)
        for arguments in client.arguments:
            self.assertEqual(set(arguments), {"query", "streams", "range_seconds", "groupings", "metrics"})

    def test_uses_tool_schema_to_avoid_unsupported_fields(self):
        schema = {"properties": {"query": {}, "streams": {}, "range_seconds": {}, "metrics": {}, "groupings": {}}}
        client = _AggregateClient([
            {"result": {"content": [{"type": "text", "text": '{"schema":[{"name":"count"}],"datarows":[[7]]}'}]}}
        ], schema=schema)

        status = GraylogAggregateSource(client, "firewall").fetch_count()

        self.assertEqual(status["aggregate_events"], 7)
        self.assertIn("metrics", client.arguments[0])
        self.assertIn("groupings", client.arguments[0])
        self.assertNotIn("group_by", client.arguments[0])

    def test_deadline_stops_aggregate_attempts(self):
        client = _AggregateClient([
            {"result": {"content": [{"type": "text", "text": '{"events": 9}'}]}}
        ])

        status = GraylogAggregateSource(client, "firewall").fetch_count(deadline_monotonic=time.monotonic() - 1)

        self.assertEqual(status["aggregate_status"], "error")
        self.assertIn("poll_budget_exceeded", status["aggregate_error"])
        self.assertEqual(client.arguments, [])


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