#!/usr/bin/env python3
"""
helix_cdc/runtime/telemetry.py

WO-ECHO-TELEM-01 - Runtime Telemetry + Receipts

Provides instrumentation for runtime metrics:
- Performance: tok/s, step times, latencies
- Memory: RSS, cache sizes, eviction counts
- Regen: calls, bytes, time, hit-rate
- Paging: faults, bytes, io time
- Routing: path choices, fallback counts
- Se: optional correlation with runtime health

This enables training a policy model to steer paging/caching/regen/routing.

Usage:
    from helix_cdc.runtime.telemetry import Telemetry, TelemetrySinkJSONL

    # Create sink and telemetry
    sink = TelemetrySinkJSONL("receipts/runtime/run_001.jsonl")
    telem = Telemetry(sink=sink, snapshot_every_tokens=100)

    # Instrument your inference loop
    for token_idx, token in enumerate(tokens):
        with telem.step_timer():
            output = model(token)

        telem.record_token()
        telem.record_cache_hit()  # or cache_miss()
        telem.record_regen(bytes=1024, time_ms=2.3)

        # Periodic snapshot
        telem.maybe_snapshot()

    # Final summary
    telem.write_summary("receipts/runtime/run_001.summary.json")
"""

from __future__ import annotations

import hashlib
import json
import os
import subprocess
import sys
import time
from contextlib import contextmanager
from dataclasses import dataclass, field, asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Generator, List, Optional, Union
import statistics


# ============================================================================
# Provenance Helpers
# ============================================================================

def sha256_file(path: str) -> str:
    """Compute SHA256 hash of a file."""
    if not os.path.exists(path):
        return "file_not_found"
    h = hashlib.sha256()
    try:
        with open(path, 'rb') as f:
            for chunk in iter(lambda: f.read(8192), b''):
                h.update(chunk)
        return h.hexdigest()
    except Exception:
        return "read_error"


def sha256_string(s: str) -> str:
    """Compute SHA256 of a string."""
    return hashlib.sha256(s.encode('utf-8')).hexdigest()


def get_git_info() -> Dict[str, str]:
    """Get git commit and dirty status."""
    try:
        commit = subprocess.run(
            ['git', 'rev-parse', 'HEAD'],
            capture_output=True, text=True, timeout=5
        )
        status = subprocess.run(
            ['git', 'status', '--porcelain'],
            capture_output=True, text=True, timeout=5
        )
        return {
            'commit': commit.stdout.strip() if commit.returncode == 0 else 'unknown',
            'dirty': len(status.stdout.strip()) > 0 if status.returncode == 0 else True,
        }
    except Exception:
        return {'commit': 'unknown', 'dirty': True}


def get_rss_mb() -> Optional[float]:
    """Get current RSS memory in MB (Linux only)."""
    try:
        with open('/proc/self/statm', 'r') as f:
            # statm: size resident shared text lib data dt (all in pages)
            parts = f.read().split()
            resident_pages = int(parts[1])
            page_size = os.sysconf('SC_PAGE_SIZE')
            return (resident_pages * page_size) / (1024 * 1024)
    except Exception:
        return None


# ============================================================================
# Counters and Accumulators
# ============================================================================

@dataclass
class StepStats:
    """Statistics for a single step/token."""
    step_time_ms: float = 0.0
    cache_hit: bool = False
    cache_miss: bool = False
    regen_called: bool = False
    regen_bytes: int = 0
    regen_time_ms: float = 0.0
    page_fault: bool = False
    page_bytes: int = 0
    page_time_ms: float = 0.0
    route_choice: str = ""
    fallback: bool = False
    se: Optional[float] = None


@dataclass
class Counters:
    """Accumulated counters since last snapshot."""
    tokens: int = 0
    steps: int = 0
    cache_hits: int = 0
    cache_misses: int = 0
    regen_calls: int = 0
    regen_bytes: int = 0
    regen_time_ms: float = 0.0
    page_faults: int = 0
    page_bytes: int = 0
    page_time_ms: float = 0.0
    fallbacks: int = 0
    step_times_ms: List[float] = field(default_factory=list)
    se_values: List[float] = field(default_factory=list)
    route_choices: Dict[str, int] = field(default_factory=dict)

    def reset(self):
        """Reset all counters."""
        self.tokens = 0
        self.steps = 0
        self.cache_hits = 0
        self.cache_misses = 0
        self.regen_calls = 0
        self.regen_bytes = 0
        self.regen_time_ms = 0.0
        self.page_faults = 0
        self.page_bytes = 0
        self.page_time_ms = 0.0
        self.fallbacks = 0
        self.step_times_ms = []
        self.se_values = []
        self.route_choices = {}


# ============================================================================
# Telemetry Sinks
# ============================================================================

class TelemetrySink:
    """Base class for telemetry output sinks."""

    def write_snapshot(self, snapshot: Dict[str, Any]) -> None:
        """Write a single snapshot."""
        raise NotImplementedError

    def close(self) -> None:
        """Close the sink."""
        pass


class TelemetrySinkJSONL(TelemetrySink):
    """JSONL file sink - one JSON object per line."""

    def __init__(self, path: str, flush_every: int = 1):
        self.path = path
        self.flush_every = flush_every
        self._write_count = 0

        # Ensure directory exists
        Path(path).parent.mkdir(parents=True, exist_ok=True)

        # Open file for append
        self._file = open(path, 'a')

    def write_snapshot(self, snapshot: Dict[str, Any]) -> None:
        """Write snapshot as single JSON line."""
        self._file.write(json.dumps(snapshot, separators=(',', ':')) + '\n')
        self._write_count += 1

        if self._write_count % self.flush_every == 0:
            self._file.flush()

    def close(self) -> None:
        """Close the file."""
        self._file.close()


class TelemetrySinkMemory(TelemetrySink):
    """In-memory sink for testing."""

    def __init__(self):
        self.snapshots: List[Dict[str, Any]] = []

    def write_snapshot(self, snapshot: Dict[str, Any]) -> None:
        self.snapshots.append(snapshot)


class TelemetrySinkMulti(TelemetrySink):
    """Fan-out to multiple sinks."""

    def __init__(self, sinks: List[TelemetrySink]):
        self.sinks = sinks

    def write_snapshot(self, snapshot: Dict[str, Any]) -> None:
        for sink in self.sinks:
            sink.write_snapshot(snapshot)

    def close(self) -> None:
        for sink in self.sinks:
            sink.close()


# ============================================================================
# Main Telemetry Class
# ============================================================================

class Telemetry:
    """
    Runtime telemetry accumulator and snapshot writer.

    Collects metrics during inference and periodically writes snapshots
    to a sink (JSONL file, memory, etc.).

    Usage:
        sink = TelemetrySinkJSONL("telemetry.jsonl")
        telem = Telemetry(sink=sink, snapshot_every_tokens=100)

        for token in tokens:
            with telem.step_timer():
                output = model(token)
            telem.record_token()
            telem.maybe_snapshot()

        telem.write_summary("summary.json")
    """

    def __init__(
        self,
        sink: Optional[TelemetrySink] = None,
        snapshot_every_tokens: int = 100,
        run_id: Optional[str] = None,
        provenance: Optional[Dict[str, Any]] = None,
    ):
        self.sink = sink
        self.snapshot_every_tokens = snapshot_every_tokens
        self.run_id = run_id or datetime.utcnow().strftime("%Y%m%d_%H%M%S")

        # Provenance (court-grade)
        self.provenance = provenance or {}

        # Counters
        self._counters = Counters()
        self._global_counters = Counters()  # Never reset, for summary

        # Timing
        self._start_time = time.perf_counter()
        self._last_snapshot_time = self._start_time
        self._last_snapshot_tokens = 0
        self._snapshot_count = 0

        # Current step timing
        self._step_start: Optional[float] = None

        # Policy events
        self._policy_events: List[Dict[str, Any]] = []

    # ========================================================================
    # Step Timing
    # ========================================================================

    @contextmanager
    def step_timer(self) -> Generator[None, None, None]:
        """Context manager to time a single step."""
        start = time.perf_counter()
        try:
            yield
        finally:
            elapsed_ms = (time.perf_counter() - start) * 1000
            self._counters.step_times_ms.append(elapsed_ms)
            self._global_counters.step_times_ms.append(elapsed_ms)
            self._counters.steps += 1
            self._global_counters.steps += 1

    def start_step(self) -> None:
        """Manual step start (alternative to context manager)."""
        self._step_start = time.perf_counter()

    def end_step(self) -> float:
        """Manual step end. Returns elapsed ms."""
        if self._step_start is None:
            return 0.0
        elapsed_ms = (time.perf_counter() - self._step_start) * 1000
        self._counters.step_times_ms.append(elapsed_ms)
        self._global_counters.step_times_ms.append(elapsed_ms)
        self._counters.steps += 1
        self._global_counters.steps += 1
        self._step_start = None
        return elapsed_ms

    # ========================================================================
    # Recording Methods
    # ========================================================================

    def record_token(self, count: int = 1) -> None:
        """Record token(s) generated."""
        self._counters.tokens += count
        self._global_counters.tokens += count

    def record_cache_hit(self, count: int = 1) -> None:
        """Record cache hit(s)."""
        self._counters.cache_hits += count
        self._global_counters.cache_hits += count

    def record_cache_miss(self, count: int = 1) -> None:
        """Record cache miss(es)."""
        self._counters.cache_misses += count
        self._global_counters.cache_misses += count

    def record_regen(self, bytes: int = 0, time_ms: float = 0.0) -> None:
        """Record regeneration event."""
        self._counters.regen_calls += 1
        self._counters.regen_bytes += bytes
        self._counters.regen_time_ms += time_ms
        self._global_counters.regen_calls += 1
        self._global_counters.regen_bytes += bytes
        self._global_counters.regen_time_ms += time_ms

    def record_page_fault(self, bytes: int = 0, time_ms: float = 0.0) -> None:
        """Record page fault / paging event."""
        self._counters.page_faults += 1
        self._counters.page_bytes += bytes
        self._counters.page_time_ms += time_ms
        self._global_counters.page_faults += 1
        self._global_counters.page_bytes += bytes
        self._global_counters.page_time_ms += time_ms

    def record_route(self, choice: str, fallback: bool = False) -> None:
        """Record routing decision."""
        self._counters.route_choices[choice] = self._counters.route_choices.get(choice, 0) + 1
        self._global_counters.route_choices[choice] = self._global_counters.route_choices.get(choice, 0) + 1
        if fallback:
            self._counters.fallbacks += 1
            self._global_counters.fallbacks += 1

    def record_se(self, se: float) -> None:
        """Record Se value for correlation."""
        self._counters.se_values.append(se)
        self._global_counters.se_values.append(se)

    def record_policy_event(
        self,
        diff: Dict[str, Any],
        source: str = "manual",
    ) -> None:
        """Record a policy/knob change event."""
        event = {
            "event": "policy_update",
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "token_index": self._global_counters.tokens,
            "diff": diff,
            "source": source,
        }
        self._policy_events.append(event)

        # Also write to sink if available
        if self.sink:
            self.sink.write_snapshot(event)

    def record_policy_action(self, action_dict: Dict[str, Any]) -> None:
        """
        Record policy action for training data collection.

        WO-ECHO-POLICY-HARNESS-01: Logs the full action dict from PolicyHarness
        for later experience replay during training.

        Args:
            action_dict: Dict from PolicyAction.to_dict() containing:
                - verify_policy: "always" | "sampled" | "trust_cached"
                - cpp_kernel: bool
                - fused_matmul: bool
                - fibpi_mode: "round_robin" | "se_bucket"
                - sample_rate: float
                - cache_size: int
                - block_rows: int
        """
        event = {
            "event": "policy_action",
            "timestamp": datetime.utcnow().isoformat() + "Z",
            "token_index": self._global_counters.tokens,
            "action": action_dict,
        }

        # Track in policy events list
        self._policy_events.append(event)

        # Write to sink for training data collection
        if self.sink:
            self.sink.write_snapshot(event)

    # ========================================================================
    # Snapshots
    # ========================================================================

    def maybe_snapshot(self, force: bool = False) -> bool:
        """
        Write a snapshot if enough tokens have been processed.

        Returns True if snapshot was written.
        """
        tokens_since = self._counters.tokens
        if not force and tokens_since < self.snapshot_every_tokens:
            return False

        self._write_snapshot()
        return True

    def _write_snapshot(self) -> None:
        """Write a telemetry snapshot."""
        now = time.perf_counter()
        elapsed_since_snapshot = now - self._last_snapshot_time
        elapsed_total = now - self._start_time

        # Compute stats
        tokens = self._counters.tokens
        tok_per_sec = tokens / elapsed_since_snapshot if elapsed_since_snapshot > 0 else 0

        step_times = self._counters.step_times_ms
        mean_step = statistics.mean(step_times) if step_times else 0.0
        p95_step = (
            statistics.quantiles(step_times, n=20)[18]
            if len(step_times) >= 20 else max(step_times) if step_times else 0.0
        )

        cache_total = self._counters.cache_hits + self._counters.cache_misses
        cache_hit_rate = self._counters.cache_hits / cache_total if cache_total > 0 else 0.0

        regen_hit_rate = (
            1.0 - (self._counters.regen_calls / self._counters.steps)
            if self._counters.steps > 0 else 1.0
        )

        se_mean = statistics.mean(self._counters.se_values) if self._counters.se_values else None

        # Get memory
        rss_mb = get_rss_mb()

        snapshot = {
            "schema": "echo.runtime.telemetry.v1",
            "run_id": self.run_id,
            "snapshot_index": self._snapshot_count,
            "timestamp": datetime.utcnow().isoformat() + "Z",

            # Token/step counts
            "token_index": self._global_counters.tokens,
            "tokens_in_window": tokens,
            "steps_in_window": self._counters.steps,

            # Performance
            "elapsed_sec": elapsed_since_snapshot,
            "elapsed_total_sec": elapsed_total,
            "tok_per_sec": tok_per_sec,
            "mean_step_ms": mean_step,
            "p95_step_ms": p95_step,

            # Memory
            "rss_mb": rss_mb,

            # Cache
            "cache_hits": self._counters.cache_hits,
            "cache_misses": self._counters.cache_misses,
            "cache_hit_rate": cache_hit_rate,

            # Regen
            "regen_calls": self._counters.regen_calls,
            "regen_bytes": self._counters.regen_bytes,
            "regen_time_ms": self._counters.regen_time_ms,
            "regen_hit_rate": regen_hit_rate,

            # Paging
            "page_faults": self._counters.page_faults,
            "page_bytes": self._counters.page_bytes,
            "page_time_ms": self._counters.page_time_ms,

            # Routing
            "route_choices": dict(self._counters.route_choices),
            "fallbacks": self._counters.fallbacks,

            # Se
            "se_mean": se_mean,
        }

        if self.sink:
            self.sink.write_snapshot(snapshot)

        # Reset window counters
        self._counters.reset()
        self._last_snapshot_time = now
        self._last_snapshot_tokens = self._global_counters.tokens
        self._snapshot_count += 1

    # ========================================================================
    # Summary
    # ========================================================================

    def write_summary(self, path: str, wall_time_ms: Optional[float] = None) -> Dict[str, Any]:
        """
        Write final summary receipt with aggregates and provenance.

        Args:
            path: Output path for summary JSON
            wall_time_ms: Optional wall-clock time override (e.g., from request context)

        Returns the summary dict.
        """
        now = time.perf_counter()
        elapsed_total = now - self._start_time

        # Allow wall_time_ms override for more accurate request timing
        if wall_time_ms is not None:
            elapsed_total = wall_time_ms / 1000.0

        gc = self._global_counters
        tokens = gc.tokens
        tok_per_sec = tokens / elapsed_total if elapsed_total > 0 else 0

        step_times = gc.step_times_ms
        mean_step = statistics.mean(step_times) if step_times else 0.0
        p50_step = statistics.median(step_times) if step_times else 0.0
        p95_step = (
            statistics.quantiles(step_times, n=20)[18]
            if len(step_times) >= 20 else max(step_times) if step_times else 0.0
        )
        p99_step = (
            statistics.quantiles(step_times, n=100)[98]
            if len(step_times) >= 100 else max(step_times) if step_times else 0.0
        )

        cache_total = gc.cache_hits + gc.cache_misses
        cache_hit_rate = gc.cache_hits / cache_total if cache_total > 0 else 0.0

        regen_hit_rate = (
            1.0 - (gc.regen_calls / gc.steps) if gc.steps > 0 else 1.0
        )

        se_mean = statistics.mean(gc.se_values) if gc.se_values else None
        se_std = statistics.stdev(gc.se_values) if len(gc.se_values) > 1 else None

        rss_mb = get_rss_mb()
        git_info = get_git_info()

        summary = {
            "schema": "echo.runtime.summary.v1",
            "work_order": "WO-ECHO-TELEM-01",
            "run_id": self.run_id,
            "timestamp": datetime.utcnow().isoformat() + "Z",

            # Provenance
            "provenance": {
                **self.provenance,
                "git_commit": git_info["commit"],
                "git_dirty": git_info["dirty"],
            },

            # Totals
            "total_tokens": tokens,
            "total_steps": gc.steps,
            "total_elapsed_sec": elapsed_total,
            "snapshots_written": self._snapshot_count,

            # Performance aggregates
            "performance": {
                "tok_per_sec": tok_per_sec,
                "mean_step_ms": mean_step,
                "p50_step_ms": p50_step,
                "p95_step_ms": p95_step,
                "p99_step_ms": p99_step,
            },

            # Memory
            "memory": {
                "final_rss_mb": rss_mb,
            },

            # Cache aggregates
            "cache": {
                "total_hits": gc.cache_hits,
                "total_misses": gc.cache_misses,
                "hit_rate": cache_hit_rate,
            },

            # Regen aggregates
            "regen": {
                "total_calls": gc.regen_calls,
                "total_bytes": gc.regen_bytes,
                "total_time_ms": gc.regen_time_ms,
                "hit_rate": regen_hit_rate,
                "pct_of_total_time": (gc.regen_time_ms / (elapsed_total * 1000) * 100) if elapsed_total > 0 else 0,
            },

            # Paging aggregates
            "paging": {
                "total_faults": gc.page_faults,
                "total_bytes": gc.page_bytes,
                "total_time_ms": gc.page_time_ms,
                "pct_of_total_time": (gc.page_time_ms / (elapsed_total * 1000) * 100) if elapsed_total > 0 else 0,
            },

            # Routing
            "routing": {
                "choices": dict(gc.route_choices),
                "total_fallbacks": gc.fallbacks,
            },

            # Se correlation
            "se": {
                "mean": se_mean,
                "std": se_std,
                "samples": len(gc.se_values),
            },

            # Policy events
            "policy_events_count": len(self._policy_events),
        }

        # Write to file
        Path(path).parent.mkdir(parents=True, exist_ok=True)
        with open(path, 'w') as f:
            json.dump(summary, f, indent=2)

        return summary

    def close(self) -> None:
        """Close the sink."""
        if self.sink:
            self.sink.close()


# ============================================================================
# Convenience: Instrumented Timer
# ============================================================================

class Timer:
    """Simple timer for measuring durations."""

    def __init__(self):
        self._start: Optional[float] = None
        self._elapsed_ms: float = 0.0

    def start(self) -> "Timer":
        self._start = time.perf_counter()
        return self

    def stop(self) -> float:
        if self._start is not None:
            self._elapsed_ms = (time.perf_counter() - self._start) * 1000
            self._start = None
        return self._elapsed_ms

    @property
    def elapsed_ms(self) -> float:
        if self._start is not None:
            return (time.perf_counter() - self._start) * 1000
        return self._elapsed_ms

    @contextmanager
    def measure(self) -> Generator[None, None, None]:
        self.start()
        try:
            yield
        finally:
            self.stop()


# ============================================================================
# Self-test
# ============================================================================

if __name__ == "__main__":
    import random

    print("=" * 60)
    print("WO-ECHO-TELEM-01 - Telemetry Self-Test")
    print("=" * 60)
    print()

    # Create in-memory sink for testing
    sink = TelemetrySinkMemory()
    telem = Telemetry(
        sink=sink,
        snapshot_every_tokens=10,
        run_id="selftest_001",
        provenance={"test": True, "script": __file__},
    )

    # Simulate 50 tokens
    print("Simulating 50 tokens...")
    for i in range(50):
        with telem.step_timer():
            # Simulate work
            time.sleep(random.uniform(0.001, 0.005))

        telem.record_token()

        # Random cache hit/miss
        if random.random() > 0.3:
            telem.record_cache_hit()
        else:
            telem.record_cache_miss()

        # Occasional regen
        if random.random() > 0.8:
            telem.record_regen(bytes=random.randint(1000, 5000), time_ms=random.uniform(1, 5))

        # Record Se
        telem.record_se(random.uniform(0.2, 0.4))

        # Snapshot check
        telem.maybe_snapshot()

    print(f"  Snapshots written: {len(sink.snapshots)}")
    print()

    # Print a snapshot sample
    if sink.snapshots:
        print("Sample snapshot:")
        sample = sink.snapshots[0]
        for key in ['token_index', 'tok_per_sec', 'cache_hit_rate', 'regen_hit_rate', 'se_mean']:
            print(f"  {key}: {sample.get(key)}")
    print()

    # Write summary
    summary = telem.write_summary("/tmp/telemetry_selftest_summary.json")
    print("Summary written to /tmp/telemetry_selftest_summary.json")
    print(f"  Total tokens: {summary['total_tokens']}")
    print(f"  tok/s: {summary['performance']['tok_per_sec']:.1f}")
    print(f"  Cache hit rate: {summary['cache']['hit_rate']:.1%}")
    print(f"  Regen hit rate: {summary['regen']['hit_rate']:.1%}")
    print()

    print("Self-test complete.")
