from __future__ import annotations

import json
import time
from collections.abc import Iterable

from .graylog_mcp import GraylogMcpClient
from .models import LogEvent
from .normalization import DEFAULT_SEARCH_FIELDS, normalize_fields


WINDOWS_DISCOVERY_FIELDS = (
    "winlog.channel",
    "winlog.provider_name",
    "winlog.record_id",
    "winlog.process.pid",
    "winlog.process.thread.id",
    "winlog.task",
    "winlog.opcode",
    "winlog.keywords",
    "winlog.event_data.logontype",
    "winlog.event_data.logon_type",
    "winlog.event_data.authenticationpackagename",
    "winlog.event_data.logonprocessname",
    "winlog.event_data.failurereason",
    "winlog.event_data.status",
    "winlog.event_data.substatus",
    "process.name",
    "process.executable",
    "process.pid",
    "process.command_line",
    "service.name",
)

DEFAULT_FIELDS = list(dict.fromkeys([*DEFAULT_SEARCH_FIELDS, *WINDOWS_DISCOVERY_FIELDS, "source", "eventid", "event_id", "winlog_event_id", "policyid", "policy_id", "policy", "rule", "rule_name", "sentbyte", "rcvdbyte", "bytes_in", "bytes_out", "hitcount", "elapsed"]))


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 = "", profile_fields: tuple[str, ...] = ()) -> None:
        self.client = client
        self.stream = stream
        self.query = query or "*"
        self.stream_label = stream_label or stream
        self.profile_fields = tuple(str(field).lower() for field in profile_fields if field)
        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, range_seconds: int = 300, probe_status: dict[str, object] | None = None, deadline_monotonic: float | None = None) -> tuple[list[LogEvent], dict[str, object]]:
        if probe_status is not None:
            status = dict(probe_status)
        else:
            try:
                status = self.client.probe()
            except RuntimeError as exc:
                return [], {
                    "status": "error",
                    "source": "graylog_mcp",
                    "events_fetched": 0,
                    "pages": 0,
                    "partial": True,
                    "error": f"probe_error: {exc}",
                    "truncated": False,
                    "latest_event_time": "",
                }
        mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
        arguments: dict[str, object] = {
            "query": self.query,
            "range_seconds": max(1, int(range_seconds)),
            "fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields, *self.profile_fields])),
        }
        if self.stream:
            arguments["streams"] = [self.stream]
        events: list[LogEvent] = []
        page_size = min(5_000, max(1, int(max_events)))
        pages = 0
        partial_error = ""
        while len(events) < max_events:
            if deadline_monotonic is not None and time.monotonic() >= deadline_monotonic:
                partial_error = "poll_budget_exceeded"
                break
            request_limit = min(page_size, max_events - len(events))
            try:
                result = self.client.call_tool("search_messages", {**arguments, "limit": request_limit, "offset": len(events)})
            except RuntimeError as exc:
                partial_error = str(exc)
                break
            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")
                partial_error = f"graylog_search_error: {detail}"
                break
            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) < request_limit:
                break
        latest = max((event.fields.get("eventtime", "") for event in events), default="")
        status.update({
            "source": "graylog_mcp",
            "events_fetched": len(events),
            "pages": pages,
            "partial": bool(partial_error),
            "error": partial_error,
            "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
        fields["fgai_stream_id"] = self.stream
        return LogEvent(raw=json.dumps(record, sort_keys=True), fields=normalize_fields(fields, self.mapping))