"""Utilities for loading and classifying recorder event streams."""

from __future__ import annotations

import gzip
import json
import zlib
from gzip import BadGzipFile
from io import StringIO
from pathlib import Path
from typing import Any, TextIO


Event = dict[str, Any]
_GZIP_MAGIC = b"\x1f\x8b"


def _parse_jsonl(stream: TextIO) -> tuple[Event, ...]:
    events: list[Event] = []
    for line_number, raw_line in enumerate(stream, start=1):
        line = raw_line.strip()
        if not line:
            continue
        try:
            payload = json.loads(line)
        except json.JSONDecodeError as exc:
            raise ValueError(f"Invalid JSON on line {line_number}: {exc.msg}") from exc
        if not isinstance(payload, dict):
            raise ValueError(f"Line {line_number} must decode to a JSON object")
        events.append(payload)

    if not events:
        raise ValueError("JSONL file is empty")

    return tuple(events)


def _salvage_gzip_text(raw: bytes) -> str | None:
    """Best-effort recovery for truncated or partially corrupted gzip data."""
    chunks: list[bytes] = []
    remaining = raw

    while remaining:
        decompressor = zlib.decompressobj(16 + zlib.MAX_WBITS)
        try:
            chunk = decompressor.decompress(remaining) + decompressor.flush()
        except zlib.error:
            break

        if chunk:
            chunks.append(chunk)

        if not decompressor.unused_data:
            break
        if decompressor.unused_data == remaining:
            break

        remaining = decompressor.unused_data

    if not chunks:
        return None

    return b"".join(chunks).decode("utf-8", errors="replace")


def _decode_jsonl_bytes(raw: bytes, *, prefer_gzip: bool) -> tuple[Event, ...]:
    gzip_error: Exception | None = None
    text_error: Exception | None = None

    if prefer_gzip:
        try:
            text = gzip.decompress(raw).decode("utf-8")
            return _parse_jsonl(StringIO(text))
        except (
            BadGzipFile,
            EOFError,
            OSError,
            UnicodeDecodeError,
            ValueError,
        ) as exc:
            gzip_error = exc

        salvaged_text = _salvage_gzip_text(raw)
        if salvaged_text is not None:
            try:
                return _parse_jsonl(StringIO(salvaged_text))
            except ValueError as exc:
                gzip_error = exc

    try:
        text = raw.decode("utf-8")
        return _parse_jsonl(StringIO(text))
    except (UnicodeDecodeError, ValueError) as exc:
        text_error = exc

    if gzip_error is not None and text_error is not None:
        raise IOError(
            f"Unable to decode recording as gzip ({gzip_error}) "
            f"or plain text ({text_error})"
        ) from gzip_error
    if gzip_error is not None:
        raise IOError(f"Unable to decode recording as gzip: {gzip_error}") from gzip_error

    assert text_error is not None
    raise IOError(f"Unable to decode recording as plain text: {text_error}") from text_error


def load_jsonl(file: Path) -> tuple[Event, ...]:
    """Load a plain or gzipped JSONL recorder file."""
    if not file.exists():
        raise FileNotFoundError(f"File not found: {file}")
    if not file.is_file():
        raise ValueError(f"Path is not a file: {file}")

    try:
        raw = file.read_bytes()
    except OSError as exc:
        raise IOError(f"Error reading file: {exc}") from exc

    if not raw:
        raise ValueError("JSONL file is empty")

    prefer_gzip = raw.startswith(_GZIP_MAGIC) or file.suffix == ".gz"
    return _decode_jsonl_bytes(raw, prefer_gzip=prefer_gzip)


def is_edit_event(event: Event) -> bool:
    """Return True for recorder edit events, including legacy untyped entries."""
    return event.get("type") in (None, "edit")


def filter_edit_events(events: tuple[Event, ...]) -> tuple[Event, ...]:
    """Keep edit events and prefer the typed stream when mixed streams appear."""
    edit_events = tuple(event for event in events if is_edit_event(event))
    if not edit_events:
        return ()

    first_typed_idx = next(
        (index for index, event in enumerate(edit_events) if event.get("type") == "edit"),
        None,
    )
    first_legacy_idx = next(
        (index for index, event in enumerate(edit_events) if event.get("type") is None),
        None,
    )

    if (
        first_typed_idx is not None
        and first_legacy_idx is not None
        and first_typed_idx < first_legacy_idx
    ):
        return tuple(event for event in edit_events if event.get("type") == "edit")

    return edit_events


def get_focus_events(events: tuple[Event, ...]) -> tuple[Event, ...]:
    """Return focus status events from a recording."""
    return tuple(event for event in events if event.get("type") == "focusStatus")


def get_editors(events: tuple[Event, ...]) -> set[str]:
    """Return the set of editors observed in the recording."""
    editors = {str(event["editor"]) for event in events if event.get("editor")}
    return editors or {"jetbrains"}
