from __future__ import annotations

from collections import defaultdict

from .entities import entity_display
from .entities import event_entities, sample_timeline
from .logs import THREAT_ACTIONS, is_utm_event
from .models import LogEvent


def correlate_source_ips(events: list[LogEvent], *, limit: int = 20) -> list[dict[str, object]]:
    """Correlate IPs, users, and hosts across independently configured Graylog streams."""
    grouped: dict[tuple[str, str], list[LogEvent]] = defaultdict(list)
    for event in events:
        for identity in event_entities(event):
            grouped[(identity["entity_type"], identity["entity"])].append(event)
    correlations = []
    for (kind, entity), source_events in grouped.items():
        streams = sorted({event.fields.get("fgai_stream", "local_syslog") for event in source_events})
        if len(streams) < 2:
            continue
        threat_events = sum(event.action in THREAT_ACTIONS or is_utm_event(event) for event in source_events)
        samples = sample_timeline(source_events)
        correlations.append({
            "entity": entity,
            "entity_type": kind,
            **entity_display(kind, entity, source_events),
            "source_ip": entity if kind == "ip" else "",
            "streams": streams,
            "events": len(source_events),
            "security_events": threat_events,
            "samples": samples,
            "first_seen": samples[0]["timestamp"] if samples else "",
            "last_seen": samples[-1]["timestamp"] if samples else "",
        })
    return sorted(correlations, key=lambda item: (int(item["security_events"]), int(item["events"])), reverse=True)[:limit]
