1"""Discovery helpers for os_log messages and os_signpost intervals.
2 
3These let an agent locate a focus window (e.g. "after the log saying X",
4"during signpost Y") before running the main lane analysis.
5"""
6from __future__ import annotations
7 
8from pathlib import Path
9from typing import Any
10 
11from . import xctrace, xml_utils
12 
13OS_LOG_SCHEMA = "os-log"
14OS_SIGNPOST_SCHEMA = "os-signpost"
15OS_SIGNPOST_INTERVAL_SCHEMA = "os-signpost-interval"
16 
17 
18def list_logs(
19    trace_path: Path,
20    toc_schemas: frozenset[str],
21    subsystem: str | None = None,
22    category: str | None = None,
23    message_contains: str | None = None,
24    message_type: str | None = None,
25    limit: int | None = None,
26    window_ns: tuple[int, int] | None = None,
27    run: int = 1,
28) -> list[dict[str, Any]]:
29    """Return os_log entries, optionally filtered. Case-insensitive contains.
30 
31    `limit` counts *post-filter* matches — including the window filter — so
32    the caller gets N matching logs inside the window rather than the first
33    N matching logs that might all fall outside it.
34    """
35    if OS_LOG_SCHEMA not in toc_schemas:
36        return []
37    xml_bytes = xctrace.export_schema(trace_path, OS_LOG_SCHEMA, run=run)
38    stream = xml_utils.RowStream(xml_bytes)
39    needle = message_contains.lower() if message_contains else None
40 
41    out: list[dict[str, Any]] = []
42    for row in stream:
43        time_el = row.get("time")
44        if time_el is None:
45            continue
46        time_ns = xml_utils.int_text(stream.resolve(time_el))
47        if time_ns is None:
48            continue
49        if not xml_utils.in_window(time_ns, window_ns):
50            continue
51 
52        sub = _str_of(row, stream, "subsystem")
53        cat = _str_of(row, stream, "category")
54        typ = _str_of(row, stream, "message-type")
55        fmt = _str_of(row, stream, "format-string")
56        msg = _str_of(row, stream, "message") or fmt
57 
58        if subsystem and (sub or "") != subsystem:
59            continue
60        if category and (cat or "") != category:
61            continue
62        if message_type and (typ or "") != message_type:
63            continue
64        if needle and needle not in (msg or "").lower() and needle not in (fmt or "").lower():
65            continue
66 
67        process_el = row.get("process")
68        process = (
69            xml_utils.extract_process(process_el, stream).get("name")
70            if process_el is not None else None
71        )
72 
73        out.append({
74            "time_ns": time_ns,
75            "time_ms": round(time_ns / 1_000_000, 3),
76            "type": typ,
77            "subsystem": sub,
78            "category": cat,
79            "process": process,
80            "message": msg,
81            "format_string": fmt,
82        })
83        if limit is not None and len(out) >= limit:
84            break
85 
86    out.sort(key=lambda e: e["time_ns"])
87    return out
88 
89 
90def list_signposts(
91    trace_path: Path,
92    toc_schemas: frozenset[str],
93    name_contains: str | None = None,
94    subsystem: str | None = None,
95    category: str | None = None,
96    window_ns: tuple[int, int] | None = None,
97    run: int = 1,
98) -> dict[str, list[dict[str, Any]]]:
99    """Return signpost intervals (paired begin/end) plus single-point events.
100 
101    Shape: { "intervals": [...], "events": [...] }. Intervals have
102    start_ms/end_ms/duration_ms; events have a single time_ms.
103 
104    Reads two complementary schemas:
105      * `os-signpost-interval`: already-paired intervals (this is where
106        user-emitted signposts like com.example.MyApp typically land).
107      * `os-signpost`: raw begin/end/event rows; we pair begins with ends
108        ourselves and fall back to point events for unpaired rows. Most
109        Apple-framework signposts (CloudKit, AppKit, …) live here.
110 
111    Filters are AND-combined. `name_contains` is a case-insensitive substring
112    match. `window_ns` keeps intervals that overlap the window (not strict
113    containment) and point events whose timestamp falls inside it.
114    """
115    # The two signpost schemas overlap: every paired begin/end in `os-signpost`
116    # also shows up as a row in `os-signpost-interval`. To avoid duplicates we
117    # prefer the pre-paired schema for intervals and only mine `os-signpost`
118    # for point events (and for begin/end pairing as a fallback when the
119    # interval schema is missing — older traces).
120    intervals: list[dict[str, Any]] = []
121    events: list[dict[str, Any]] = []
122 
123    has_intervals = OS_SIGNPOST_INTERVAL_SCHEMA in toc_schemas
124    if has_intervals:
125        intervals.extend(_read_interval_schema(trace_path, run=run))
126 
127    if OS_SIGNPOST_SCHEMA in toc_schemas:
128        more_intervals, more_events = _read_event_schema(trace_path, run=run)
129        if not has_intervals:
130            intervals.extend(more_intervals)
131        events.extend(more_events)
132 
133    intervals.sort(key=lambda i: i["start_ns"])
134    events.sort(key=lambda e: e["time_ns"])
135 
136    needle = name_contains.lower() if name_contains else None
137 
138    def _matches(entry: dict) -> bool:
139        if subsystem and (entry.get("subsystem") or "") != subsystem:
140            return False
141        if category and (entry.get("category") or "") != category:
142            return False
143        if needle and needle not in (entry.get("name") or "").lower():
144            return False
145        return True
146 
147    if subsystem or category or needle:
148        intervals = [i for i in intervals if _matches(i)]
149        events = [e for e in events if _matches(e)]
150 
151    if window_ns is not None:
152        s, e = window_ns
153        intervals = [
154            i for i in intervals
155            if not (i["end_ns"] < s or i["start_ns"] > e)
156        ]
157        events = [ev for ev in events if s <= ev["time_ns"] <= e]
158 
159    return {"intervals": intervals, "events": events}
160 
161 
162def _read_interval_schema(trace_path: Path, run: int = 1) -> list[dict[str, Any]]:
163    """Read the os-signpost-interval schema (pre-paired intervals)."""
164    xml_bytes = xctrace.export_schema(trace_path, OS_SIGNPOST_INTERVAL_SCHEMA, run=run)
165    stream = xml_utils.RowStream(xml_bytes)
166 
167    out: list[dict[str, Any]] = []
168    for row in stream:
169        start_el = xml_utils.first_present(row, "start", "time")
170        dur_el = row.get("duration")
171        if start_el is None or dur_el is None:
172            continue
173        start_ns = xml_utils.int_text(stream.resolve(start_el))
174        dur_ns = xml_utils.int_text(stream.resolve(dur_el))
175        if start_ns is None or dur_ns is None:
176            continue
177        end_ns = start_ns + dur_ns
178 
179        name = _str_of(row, stream, "name")
180        sub = _str_of(row, stream, "subsystem")
181        cat = _str_of(row, stream, "category")
182        signpost_id = _str_of(row, stream, "identifier") or _str_of(row, stream, "signpost-id")
183        process_el = row.get("process")
184        process = (
185            xml_utils.extract_process(process_el, stream).get("name")
186            if process_el is not None else None
187        )
188 
189        out.append({
190            "start_ns": start_ns,
191            "end_ns": end_ns,
192            "duration_ns": dur_ns,
193            "start_ms": round(start_ns / 1_000_000, 3),
194            "end_ms": round(end_ns / 1_000_000, 3),
195            "duration_ms": round(dur_ns / 1_000_000, 3),
196            "name": name,
197            "subsystem": sub,
198            "category": cat,
199            "process": process,
200            "signpost_id": signpost_id,
201        })
202    return out
203 
204 
205def _read_event_schema(
206    trace_path: Path,
207    run: int = 1,
208) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
209    """Read the os-signpost schema and pair begin/end rows into intervals."""
210    xml_bytes = xctrace.export_schema(trace_path, OS_SIGNPOST_SCHEMA, run=run)
211    stream = xml_utils.RowStream(xml_bytes)
212 
213    pending: dict[tuple, dict] = {}
214    intervals: list[dict[str, Any]] = []
215    events: list[dict[str, Any]] = []
216 
217    for row in stream:
218        time_el = xml_utils.first_present(row, "time", "start")
219        if time_el is None:
220            continue
221        time_ns = xml_utils.int_text(stream.resolve(time_el))
222        if time_ns is None:
223            continue
224 
225        name = _str_of(row, stream, "name")
226        sub = _str_of(row, stream, "subsystem")
227        cat = _str_of(row, stream, "category")
228        event_type = _str_of(row, stream, "event-type") or _str_of(row, stream, "message-type")
229        signpost_id = _str_of(row, stream, "signpost-id") or _str_of(row, stream, "identifier")
230        process_el = row.get("process")
231        process = (
232            xml_utils.extract_process(process_el, stream).get("name")
233            if process_el is not None else None
234        )
235 
236        key = (process, sub, cat, name, signpost_id)
237        etype = (event_type or "").lower()
238 
239        if etype in ("begin", "interval begin", "start"):
240            pending[key] = {"start_ns": time_ns, "name": name,
241                            "subsystem": sub, "category": cat,
242                            "process": process, "signpost_id": signpost_id}
243        elif etype in ("end", "interval end", "stop"):
244            start = pending.pop(key, None)
245            if start is not None:
246                dur_ns = time_ns - start["start_ns"]
247                intervals.append({
248                    **start,
249                    "end_ns": time_ns,
250                    "duration_ns": dur_ns,
251                    "start_ms": round(start["start_ns"] / 1_000_000, 3),
252                    "end_ms": round(time_ns / 1_000_000, 3),
253                    "duration_ms": round(dur_ns / 1_000_000, 3),
254                })
255            else:
256                events.append(_point_event(time_ns, name, sub, cat,
257                                            process, signpost_id, event_type))
258        else:
259            events.append(_point_event(time_ns, name, sub, cat,
260                                        process, signpost_id, event_type))
261 
262    # Unclosed begins are surfaced as point events so nothing is silently dropped.
263    for info in pending.values():
264        events.append(_point_event(info["start_ns"], info["name"],
265                                    info["subsystem"], info["category"],
266                                    info["process"], info["signpost_id"],
267                                    "Begin (unclosed)"))
268 
269    return intervals, events
270 
271 
272def _point_event(time_ns, name, subsystem, category, process, signpost_id, event_type):
273    return {
274        "time_ns": time_ns,
275        "time_ms": round(time_ns / 1_000_000, 3),
276        "name": name,
277        "subsystem": subsystem,
278        "category": category,
279        "process": process,
280        "signpost_id": signpost_id,
281        "event_type": event_type,
282    }
283 
284 
285def _str_of(row, stream, key):
286    el = row.get(key)
287    if el is None:
288        return None
289    resolved = stream.resolve(el)
290    txt = xml_utils.str_text(resolved) or resolved.get("fmt")
291    return txt
292