from __future__ import annotations

from collections import defaultdict

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]]:
    grouped: dict[str, list[LogEvent]] = defaultdict(list)
    for event in events:
        if event.src_ip:
            grouped[event.src_ip].append(event)
    correlations = []
    for source_ip, 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 = [
            {
                "stream": event.fields.get("fgai_stream", "local_syslog"), "timestamp": event.fields.get("eventtime", event.fields.get("timestamp", "")),
                "type": event.fields.get("type", ""), "subtype": event.subtype, "action": event.action,
                "severity": event.severity, "destination": event.dst_ip or "", "service": event.fields.get("service", ""),
                "context": event.fields.get("query_domain", event.fields.get("qh", event.fields.get("message", ""))),
            }
            for event in source_events[:20]
        ]
        correlations.append({"source_ip": source_ip, "streams": streams, "events": len(source_events), "security_events": threat_events, "samples": samples})
    return sorted(correlations, key=lambda item: (int(item["security_events"]), int(item["events"])), reverse=True)[:limit]
