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"]), 12)
        self.assertEqual(compact["field_deviations"][0]["entity"], "alice")
        self.assertNotIn("sample_events", compact["field_deviations"][0])
        self.assertNotIn("graylog_query", str(compact))

    def test_dashboard_compaction_strips_incident_and_correlation_samples(self):
        status = {
            "incidents": [
                {
                    "entity": "10.0.0.5",
                    "score": 90,
                    "severity": "critical",
                    "correlated_streams": ["Firewall", "Proxy"],
                    "evidence": ["large burst of blocked traffic"],
                    "timeline": [{"message": "raw incident evidence", "graylog_query": "srcip:10.0.0.5"}],
                }
            ],
            "cross_source_correlations": [
                {
                    "entity": "10.0.0.5",
                    "streams": ["Firewall", "Proxy"],
                    "events": 20,
                    "samples": [{"message": "raw correlation evidence", "graylog_query": "srcip:10.0.0.5"}],
                }
            ],
            "anomalies": [
                {
                    "subject": "10.0.0.5",
                    "score": 90,
                    "severity": "critical",
                    "reasons": ["burst"],
                    "evidence": {"timed_events": 100, "sample_events": [{"message": "raw anomaly"}]},
                }
            ],
        }

        compact = compact_dashboard_analysis(status)
        compact_text = str(compact)

        self.assertNotIn("timeline", compact_text)
        self.assertNotIn("samples", compact_text)
        self.assertNotIn("sample_events", compact_text)
        self.assertNotIn("graylog_query", compact_text)
        self.assertNotIn("raw incident evidence", compact_text)
        self.assertEqual(compact["incidents"][0]["entity"], "10.0.0.5")


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