"""Multi-stream benchmark: many concurrent processors on one SDK.

Creates a single SDK instance and runs N audio processors concurrently (one per
thread), each processing the same input WAV. process_frame releases the GIL, so
the streams run in parallel. Reports success/failure counts, peak concurrency,
per-stream wall time, and (when psutil is installed) process memory growth over
the run.

Environment variables
---------------------
SANAS_API_KEY      (required) licence key
SANAS_STORAGE_DIR  SDK storage directory (data + logs)     (default: ./storage)

The model and input WAV are command-line arguments (--model / --input).

Usage::

    python multi_stream_example.py --streams 8 --concurrency 4 --rate 2 \
        --model VI_G_SE --input test_input.wav

Install ``psutil`` (``pip install psutil``) to also report memory usage, which
is useful for spotting leaks in long soak runs; it is optional and the
benchmark runs without it.
"""

import argparse
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)

import sanas  # noqa: E402
from helpers import (  # noqa: E402
    PipelineWaiter,
    feed_and_drain,
    load_samples,
)
from wav_utils import save_wav  # noqa: E402

# Memory reporting is optional: psutil gives cross-platform current RSS. When it
# is not installed the benchmark still runs; the memory lines are just omitted.
try:
    import psutil  # noqa: E402

    _PROC = psutil.Process()

    def _rss_mb():
        return _PROC.memory_info().rss / (1024 * 1024)

except ImportError:
    _PROC = None

    def _rss_mb():
        return None


class LiveStats:
    """Thread-safe counters aggregated across all streams as they run.

    Streams call ``stream_started`` when they begin and ``stream_finished`` when
    they end; the benchmark reads ``snapshot`` for the final report. ``active``
    and ``peak_concurrent`` reveal how many streams actually overlapped (bounded
    by --concurrency), which the per-stream result dicts cannot show on their
    own.
    """

    def __init__(self):
        self._lock = threading.Lock()
        self.started = 0
        self.completed = 0
        self.failed = 0
        self.active = 0
        self.peak_concurrent = 0
        self.start_time = time.monotonic()

    def stream_started(self):
        with self._lock:
            self.started += 1
            self.active += 1
            if self.active > self.peak_concurrent:
                self.peak_concurrent = self.active

    def stream_finished(self, ok):
        with self._lock:
            self.active -= 1
            if ok:
                self.completed += 1
            else:
                self.failed += 1

    def snapshot(self):
        with self._lock:
            return {
                "started": self.started,
                "completed": self.completed,
                "failed": self.failed,
                "active": self.active,
                "peak_concurrent": self.peak_concurrent,
                "uptime": time.monotonic() - self.start_time,
            }


class MemorySampler(threading.Thread):
    """Background thread that records peak process RSS during the run.

    A no-op when psutil is unavailable. Sampling (rather than reading once at
    the end) captures the true high-water mark while streams overlap, so a leak
    or a transient spike is visible in the final report.
    """

    def __init__(self, interval=0.5):
        super().__init__(daemon=True)
        self._stop = threading.Event()
        self._interval = interval
        self.initial_mb = _rss_mb()
        self.peak_mb = self.initial_mb or 0.0

    def run(self):
        if _PROC is None:
            return
        while not self._stop.wait(self._interval):
            cur = _rss_mb()
            if cur is not None and cur > self.peak_mb:
                self.peak_mb = cur

    def stop(self):
        self._stop.set()


def _parse_args():
    p = argparse.ArgumentParser(description="Sanas SDK multi-stream benchmark")
    p.add_argument(
        "--streams",
        type=int,
        default=4,
        help="total number of streams to run (default: 4)",
    )
    p.add_argument(
        "--concurrency",
        type=int,
        default=0,
        help="max streams running at once (default: = --streams)",
    )
    p.add_argument(
        "--rate",
        type=float,
        default=0.0,
        help="stream start rate per second (0 = all at once)",
    )
    p.add_argument(
        "--model",
        default="VI_G_SE",
        help="ml model key (default: %(default)s)",
    )
    p.add_argument(
        "--input",
        default="test_input.wav",
        help="input WAV path (default: %(default)s)",
    )
    p.add_argument(
        "--max-throughput",
        action="store_true",
        help="feed as fast as possible (default: real-time pacing)",
    )
    p.add_argument(
        "--pcm16",
        action="store_true",
        help="negotiate L16 (raw 16-bit PCM) instead of the rate-default codec; "
        "cloud inference only (ignored for local)",
    )
    p.add_argument(
        "--save-wav-out-period",
        type=int,
        default=0,
        metavar="N",
        help="save every Nth stream's processed audio to stream_<idx>.wav for "
        "spot-checking (default: 0 = don't save)",
    )
    args = p.parse_args()
    if args.streams < 1:
        p.error("--streams must be >= 1")
    if args.concurrency < 0:
        p.error("--concurrency must be >= 0 (0 = same as --streams)")
    if args.rate < 0:
        p.error("--rate must be >= 0 (0 = start all at once)")
    if args.save_wav_out_period < 0:
        p.error("--save-wav-out-period must be >= 0 (0 = don't save)")
    return args


def _run_stream(
    sdk,
    idx,
    samples,
    sample_rate,
    channels,
    model,
    realtime,
    use_pcm16,
    save_period,
    stats,
):
    """Run one stream; return a per-stream result dict."""
    result = {
        "idx": idx,
        "ok": False,
        "error": None,
        "wall_s": 0.0,
        "saved": None,
    }
    stats.stream_started()
    t0 = time.monotonic()
    try:
        cloud_params = sanas.CloudInferencingParams()
        cloud_params.use_pcm16 = use_pcm16
        waiter = PipelineWaiter()
        attrs = sanas.ProcessorAttributes(
            audio_attributes=sanas.AudioAttributes(
                sampling_rate=sample_rate,
                channels=channels,
                model_name=model,
                cloud_inferencing_params=cloud_params,
                audio_pipeline_state_notify=waiter.callback,
            )
        )
        with sdk.create_audio_processor(attrs) as proc:
            waiter.wait_until_running(timeout=30.0)
            drained = feed_and_drain(
                proc,
                samples,
                sample_rate,
                channels,
                realtime=realtime,
                collect_output=save_period > 0,
            )
        result["wall_s"] = time.monotonic() - t0
        # Spot-check output: save every Nth stream. Off by default. Note the
        # helper builds the output buffer regardless, so this only adds a disk
        # write and doesn't materially change the memory-growth metric.
        if save_period and idx % save_period == 0 and drained["output"]:
            out_path = f"stream_{idx}.wav"
            save_wav(out_path, drained["output"], sample_rate, channels)
            result["saved"] = out_path
        result["ok"] = True
    except Exception as e:  # noqa: BLE001 - report per-stream, keep others going
        result["error"] = f"{type(e).__name__}: {e}"
        result["wall_s"] = time.monotonic() - t0
    finally:
        stats.stream_finished(result["ok"])
    return result


def _report(results, stats, sampler):
    ok = [r for r in results if r["ok"]]
    failed = [r for r in results if not r["ok"]]
    snap = stats.snapshot()
    total_finished = snap["completed"] + snap["failed"]
    success_rate = 100.0 * snap["completed"] / total_finished if total_finished else 0.0

    print("\n" + "=" * 60)
    print("Benchmark summary")
    print("=" * 60)
    print(f"streams        : {len(results)}  (ok={len(ok)}, failed={len(failed)})")
    print(f"success rate   : {success_rate:.1f}%")
    print(f"peak concurrent: {snap['peak_concurrent']}")
    print(f"uptime         : {snap['uptime']:.2f}s")
    if sampler.initial_mb is not None:
        growth = sampler.peak_mb - sampler.initial_mb
        print(
            f"memory (RSS)   : start={sampler.initial_mb:.1f} MB  "
            f"peak={sampler.peak_mb:.1f} MB  growth={growth:+.1f} MB"
        )
    else:
        print("memory (RSS)   : n/a (install psutil to report memory)")
    for r in failed:
        print(f"  stream {r['idx']} FAILED: {r['error']}")


def main():
    args = _parse_args()
    if not os.environ.get("SANAS_API_KEY"):
        sys.exit("SANAS_API_KEY is not set.")
    concurrency = args.concurrency or args.streams
    realtime = not args.max_throughput

    sdk = sanas.create_sdk(
        sanas.InitParams(
            storage_dir=os.environ.get("SANAS_STORAGE_DIR", "./storage"),
        )
    )
    res = sdk.activate_api_key(os.environ["SANAS_API_KEY"])
    if not res.success:
        sys.exit(f"Activation failed: {res.message}")
    print(f"[activation] OK (SDK {sdk.version})")

    samples, sample_rate, channels = load_samples(args.input)
    audio_s = (len(samples) // channels) / sample_rate
    print(f"[input] {sample_rate} Hz, {channels} ch, {audio_s:.2f}s")
    print(
        f"[run] streams={args.streams} concurrency={concurrency} "
        f"rate={args.rate or 'all-at-once'} model={args.model} "
        f"{'realtime' if realtime else 'max-throughput'}"
    )

    stats = LiveStats()
    sampler = MemorySampler()
    sampler.start()

    results = []
    lock = threading.Lock()
    with ThreadPoolExecutor(max_workers=concurrency) as pool:
        futures = []
        for i in range(args.streams):
            futures.append(
                pool.submit(
                    _run_stream,
                    sdk,
                    i,
                    samples,
                    sample_rate,
                    channels,
                    args.model,
                    realtime,
                    args.pcm16,
                    args.save_wav_out_period,
                    stats,
                )
            )
            if args.rate > 0 and i + 1 < args.streams:
                time.sleep(1.0 / args.rate)  # stagger stream starts
        for fut in as_completed(futures):
            r = fut.result()
            with lock:
                results.append(r)
                tag = "ok" if r["ok"] else f"FAILED ({r['error']})"
                saved = f"  saved={r['saved']}" if r.get("saved") else ""
                print(f"[stream {r['idx']}] {tag}  wall={r['wall_s']:.2f}s{saved}")

    sampler.stop()
    results.sort(key=lambda r: r["idx"])
    _report(results, stats, sampler)


if __name__ == "__main__":
    main()
