"""JSONL telemetry for HAKO's five mandated streams.

Streams (one append-only JSONL file each):
    learning      -- losses, effective learning rate, Lyapunov value
    routing       -- router action probs/entropy, MoE gates, load balance
    orchestration -- cyclic-thinking potential P, cycles, GHSOM growth, N
    diffusion     -- game level, difficulty, reward, denoiser loss
    thinking      -- request decomposition/feedback/completeness trace

Proposition P-TEL (lossless auditability). Each record is a self-contained
JSON object with monotone field 't' inside its stream; because the writer is
append-only and atomic per line, any prefix of the file is a valid partial
trace of the run (crash-consistent audit).
"""
from __future__ import annotations

import json
import os
import time
from pathlib import Path
from typing import Any, Dict

_STREAMS = ("learning", "routing", "orchestration", "diffusion", "thinking")


class Telemetry:
    def __init__(self, cfg) -> None:
        self.dir = Path(cfg.telemetry_dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        self.paths = {s: self.dir / f"{s}.jsonl" for s in _STREAMS}
        self.counters = {s: 0 for s in _STREAMS}
        self.enabled = True

    def log(self, stream: str, record: Dict[str, Any]) -> None:
        if not self.enabled or stream not in self.paths:
            return
        rec = {"t": self.counters[stream], "ts": time.time(), **record}
        self.counters[stream] += 1
        line = json.dumps(rec, ensure_ascii=True, default=float)
        with open(self.paths[stream], "a", encoding="utf-8") as fh:
            fh.write(line + "\n")
            fh.flush()
            os.fsync(fh.fileno())

    # convenience wrappers -------------------------------------------------
    def learning(self, **kw):  self.log("learning", kw)
    def routing(self, **kw):   self.log("routing", kw)
    def orchestration(self, **kw): self.log("orchestration", kw)
    def diffusion(self, **kw): self.log("diffusion", kw)
    def thinking(self, **kw):  self.log("thinking", kw)

    def summary(self) -> Dict[str, int]:
        return {s: self.counters[s] for s in _STREAMS}