"""Benchmark the actual Pipecat Smart Turn endpoint call on the 44 inputs."""
import csv
import importlib.metadata
import json
import os
from pathlib import Path
import time

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
os.environ.setdefault("HF_HOME", str(HERE / ".cache/huggingface"))
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
os.environ.setdefault("VECLIB_MAXIMUM_THREADS", "1")
import numpy as np
import soundfile as sf
from loguru import logger
import pipecat.audio.turn.smart_turn.local_smart_turn_v3 as local


def stats(values):
    v = np.asarray(values, dtype=float)
    return dict(n=len(v), mean_ms=float(v.mean()), median_ms=float(np.median(v)),
                p95_ms=float(np.percentile(v, 95)), p99_ms=float(np.percentile(v, 99)),
                min_ms=float(v.min()), max_ms=float(v.max()))


class TimedSession:
    def __init__(self, session, current):
        self.session, self.current = session, current

    def run(self, *args, **kwargs):
        begin = time.perf_counter()
        output = self.session.run(*args, **kwargs)
        self.current["onnx_ms"] = (time.perf_counter()-begin)*1000
        return output


def main():
    logger.remove()
    with (HERE / "results/work-data-can-reply/predictions.csv").open(encoding="utf-8-sig") as f:
        labels = list(csv.DictReader(f))
    audio_cache, inputs = {}, []
    for row in labels:
        call = row["call_id"]
        if call not in audio_cache:
            pcm, sr = sf.read(ROOT / "data/韩泰机器人抢话/音频" / f"{call}.mp3", dtype="int16", always_2d=True)
            audio_cache[call] = pcm[:, 0].astype(np.float32)/32768.0, sr
        pcm, sr = audio_cache[call]
        start = round(float(row["audio_start_ms"])*sr/1000)
        end = int(row["decision_ms"])*sr//1000
        inputs.append((row, np.ascontiguousarray(pcm[start:end]), sr))
    report = dict(method="Warm cached model endpoint timing; same 44 Smart Turn evaluation crops, 5 repetitions. Includes resampling, left padding, Whisper features and ONNX inference. Excludes audio decoding, import, model loading and silence/ASR waits.",
                  pipecat_version=importlib.metadata.version("pipecat-ai"),
                  onnxruntime_version=importlib.metadata.version("onnxruntime"),
                  repetitions=5, cases={})
    original_feature = local.compute_whisper_log_mel_features
    for threads in (1, 4):
        began = time.perf_counter()
        analyzer = local.LocalSmartTurnAnalyzerV3(smart_turn_model_path=str(HERE / "smart-turn-v3.2-cpu.onnx"), cpu_count=threads)
        load_ms = (time.perf_counter()-began)*1000
        analyzer.set_sample_rate(inputs[0][2])
        for _ in range(5):
            analyzer._predict_endpoint(inputs[0][1])
        current = {}
        original_resample = analyzer._resample_to_model_rate

        def resample(audio):
            began = time.perf_counter()
            output = original_resample(audio)
            current["resample_ms"] = (time.perf_counter()-began)*1000
            return output

        def features(*args, **kwargs):
            began = time.perf_counter()
            output = original_feature(*args, **kwargs)
            current["feature_ms"] = (time.perf_counter()-began)*1000
            return output

        analyzer._resample_to_model_rate = resample
        analyzer._session = TimedSession(analyzer._session, current)
        local.compute_whisper_log_mel_features = features
        measured = []
        try:
            for repeat in range(5):
                # Rotate order to avoid always timing the same input first.
                order = inputs[repeat:] + inputs[:repeat]
                for label, sample, sr in order:
                    analyzer.set_sample_rate(sr)
                    current.clear()
                    began = time.perf_counter()
                    result = analyzer._predict_endpoint(sample)
                    total = (time.perf_counter()-began)*1000
                    assert abs(result["probability"]-float(label["probability"])) < 1e-4, (label["call_id"], label["decision_ms"], result["probability"], label["probability"])
                    measured.append(dict(call_id=label["call_id"], decision_ms=int(label["decision_ms"]),
                                         repetition=repeat, total_ms=total, **current))
        finally:
            local.compute_whisper_log_mel_features = original_feature
        report["cases"][str(threads)] = dict(cpu_threads=threads, model_load_ms=load_ms,
            total=stats([r["total_ms"] for r in measured]),
            resample=stats([r["resample_ms"] for r in measured]),
            features=stats([r["feature_ms"] for r in measured]),
            onnx=stats([r["onnx_ms"] for r in measured]), samples=measured)
        print(threads, json.dumps({k:v for k,v in report["cases"][str(threads)].items() if k != "samples"}), flush=True)
    destination = HERE / "results/latency.json"
    destination.write_text(json.dumps(report, indent=2)+"\n")
    print(f"Saved {destination}")


if __name__ == "__main__":
    main()