import unittest

from fgai.correlation import correlate_source_ips
from fgai.logs import parse_log_line


class CorrelationTests(unittest.TestCase):
    def test_correlates_same_source_across_streams(self):
        events = [
            parse_log_line("srcip=10.0.0.5 fgai_stream=Fortigate action=blocked"),
            parse_log_line("srcip=10.0.0.5 fgai_stream=DNS action=deny"),
        ]
        result = correlate_source_ips(events)
        self.assertEqual(result[0]["source_ip"], "10.0.0.5")
        self.assertEqual(result[0]["streams"], ["DNS", "Fortigate"])
        self.assertEqual(len(result[0]["samples"]), 2)

    def test_correlates_user_across_streams(self):
        events = [
            parse_log_line("username=alice fgai_stream=Windows action=login timestamp=2026-06-24T10:00:00Z"),
            parse_log_line("username=alice fgai_stream=VPN action=accept timestamp=2026-06-24T10:01:00Z"),
        ]
        result = correlate_source_ips(events)
        self.assertEqual(result[0]["entity"], "alice")
        self.assertEqual(result[0]["entity_type"], "user")

    def test_ip_entity_prefers_hostname_display_label(self):
        events = [
            parse_log_line("srcip=10.0.0.5 hostname=win01 fgai_stream=Windows action=login"),
            parse_log_line("srcip=10.0.0.5 hostname=win01 fgai_stream=Firewall action=blocked"),
        ]
        result = correlate_source_ips(events)
        self.assertEqual(result[0]["entity"], "10.0.0.5")
        self.assertEqual(result[0]["entity_display"], "win01")
        self.assertEqual(result[0]["entity_label"], "win01 (10.0.0.5)")
        self.assertEqual(result[0]["entity_detail"], "10.0.0.5")

    def test_ip_entity_falls_back_to_username_display_label(self):
        events = [
            parse_log_line("srcip=10.0.0.6 username=alice fgai_stream=Windows action=login"),
            parse_log_line("srcip=10.0.0.6 username=alice fgai_stream=VPN action=accept"),
        ]
        result = correlate_source_ips(events)
        self.assertEqual(result[0]["entity"], "10.0.0.6")
        self.assertEqual(result[0]["entity_label"], "alice (10.0.0.6)")

    def test_ip_entity_falls_back_to_service_display_label(self):
        events = [
            parse_log_line("srcip=10.0.0.7 process.name=spoolsv.exe fgai_stream=Windows action=started"),
            parse_log_line("srcip=10.0.0.7 process.name=spoolsv.exe fgai_stream=Sysmon action=connect"),
        ]
        result = correlate_source_ips(events)
        self.assertEqual(result[0]["entity"], "10.0.0.7")
        self.assertEqual(result[0]["entity_label"], "spoolsv.exe (10.0.0.7)")