from __future__ import annotations

import json
from collections.abc import Iterable

from .graylog_mcp import GraylogMcpClient
from .models import LogEvent


DEFAULT_FIELD_MAP = {
    "srcip": ("srcip", "src_ip", "source_ip", "client_ip"),
    "dstip": ("dstip", "dst_ip", "destination_ip", "server_ip"),
    "dstport": ("dstport", "dst_port", "destination_port"),
    "srcport": ("srcport", "src_port", "source_port"),
    "eventtime": ("eventtime", "timestamp", "time"),
    "severity": ("severity", "level"),
    "action": ("action", "event_action", "disposition"),
}

DEFAULT_FIELDS = ["timestamp", "source", "srcip", "dstip", "srcport", "dstport", "service", "action", "severity", "policyid", "subtype", "type", "sentbyte", "rcvdbyte", "hitcount"]


def _records(value: object) -> Iterable[dict[str, object]]:
    if isinstance(value, dict):
        schema = value.get("schema")
        datarows = value.get("datarows")
        if isinstance(schema, list) and isinstance(datarows, list):
            fields = [str(column.get("field", "")) for column in schema if isinstance(column, dict)]
            for row in datarows:
                if isinstance(row, list):
                    yield {field: row[index] for index, field in enumerate(fields) if field and index < len(row)}
            return
        for key in ("messages", "results", "events", "data"):
            if isinstance(value.get(key), list):
                yield from (item for item in value[key] if isinstance(item, dict))
                return
        if value:
            yield value
    elif isinstance(value, list):
        yield from (item for item in value if isinstance(item, dict))


class GraylogStreamSource:
    def __init__(self, client: GraylogMcpClient, stream: str, query: str = "*", field_mapping: str = "", stream_label: str = "") -> None:
        self.client = client
        self.stream = stream
        self.query = query or "*"
        self.stream_label = stream_label or stream
        try:
            self.mapping = json.loads(field_mapping) if field_mapping else {}
        except json.JSONDecodeError as exc:
            raise RuntimeError("invalid_graylog_field_mapping") from exc
        if not isinstance(self.mapping, dict):
            raise RuntimeError("invalid_graylog_field_mapping")

    def fetch(self, *, max_events: int = 5_000) -> tuple[list[LogEvent], dict[str, object]]:
        status = self.client.probe()
        mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
        arguments: dict[str, object] = {
            "query": self.query,
            "range_seconds": 300,
            "fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields])),
        }
        if self.stream:
            arguments["streams"] = [self.stream]
        events: list[LogEvent] = []
        page_size = 1_000
        pages = 0
        while len(events) < max_events:
            result = self.client.call_tool("search_messages", {**arguments, "limit": page_size, "offset": len(events)})
            content = result.get("result", {}).get("content", []) if isinstance(result.get("result"), dict) else []
            if isinstance(result.get("result"), dict) and result["result"].get("isError"):
                detail = next((str(item.get("text")) for item in content if isinstance(item, dict) and item.get("type") == "text"), "Graylog search failed")
                raise RuntimeError(f"graylog_search_error: {detail}")
            records: list[dict[str, object]] = []
            for item in content if isinstance(content, list) else []:
                if isinstance(item, dict) and item.get("type") == "text":
                    try:
                        records.extend(_records(json.loads(str(item.get("text", "")))))
                    except json.JSONDecodeError:
                        continue
            events.extend(self._event(record) for record in records)
            pages += 1
            if len(records) < page_size:
                break
        latest = max((event.fields.get("eventtime", "") for event in events), default="")
        status.update({"source": "graylog_mcp", "events_fetched": len(events), "pages": pages, "truncated": len(events) >= max_events, "latest_event_time": latest})
        return events, status

    def _event(self, record: dict[str, object]) -> LogEvent:
        fields = {str(key).lower(): str(value) for key, value in record.items() if value is not None}
        fields["fgai_stream"] = self.stream_label
        for canonical, candidates in DEFAULT_FIELD_MAP.items():
            mapped = self.mapping.get(canonical)
            candidates = (str(mapped),) if mapped else candidates
            for candidate in candidates:
                if candidate.lower() in fields:
                    fields[canonical] = fields[candidate.lower()]
                    break
        return LogEvent(raw=json.dumps(record, sort_keys=True), fields=fields)
