import unittest

from fgai.llm import _json_object_from_text, compact_dashboard_analysis


class LlmTests(unittest.TestCase):
    def test_json_object_from_text_accepts_markdown_wrapped_json(self):
        payload = _json_object_from_text(
            '```json\n{"profiles":[{"stream_id":"windows","entity_fields":["username"]}]}\n```'
        )

        self.assertEqual(payload["profiles"][0]["stream_id"], "windows")

    def test_json_object_from_text_extracts_object_from_extra_text(self):
        payload = _json_object_from_text(
            'Here is the profile:\n{"profiles":[{"stream_id":"firewall"}]}\nDone.'
        )

        self.assertEqual(payload["profiles"][0]["stream_id"], "firewall")

    def test_dashboard_compaction_removes_heavy_evidence_payloads(self):
        status = {
            "summary": {"total": 1000},
            "capabilities": {
                "graylog_mcp": {
                    "status": "connected",
                    "streams": [{"stream_id": f"stream-{index}", "events": index} for index in range(100)],
                    "aggregate_events": 1000,
                    "raw_events_fetched": 50,
                }
            },
            "event_context": {
                "source_profiles": [{"entity": f"10.0.0.{index}"} for index in range(20)],
                "related_activity": [{"entity": f"10.0.0.{index}"} for index in range(40)],
            },
            "field_deviations": {
                "alice": [
                    {
                        "stream_name": "Windows",
                        "detector": "new_relationship",
                        "field": "relationship:username->srcip",
                        "score": 90,
                        "reason": "new srcip value",
                        "sample_events": [{"message": "very large raw event", "graylog_query": "username:alice"}],
                    }
                ]
            },
        }

        compact = compact_dashboard_analysis(status)

        self.assertEqual(compact["capabilities"]["graylog_mcp"]["status"], "connected")
        self.assertNotIn("streams", compact["capabilities"]["graylog_mcp"])
        self.assertEqual(len(compact["event_context"]["source_profiles"]), 10)
        self.assertEqual(len(compact["event_context"]["related_activity"]), 20)
        self.assertEqual(compact["field_deviations"][0]["entity"], "alice")
        self.assertNotIn("sample_events", compact["field_deviations"][0])
        self.assertNotIn("graylog_query", str(compact))


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