"""Causal Silero ONNX baseline: frame ends mark when scores become available."""
import bisect
import csv
import hashlib
import importlib.metadata
import json
import platform
import sys
from pathlib import Path

import numpy as np
import onnxruntime as ort
import soundfile as sf

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
WEIGHT_SHA256 = '597d30b3ec076608d059477bb14cfeffdf951bf5cae370d38f65d33bbfe82004'


class SileroStream:
    """Same state/context contract as the official Silero ONNX wrapper."""
    def __init__(self, session, sr):
        if sr not in (8000, 16000):
            raise ValueError(f'Unsupported sample rate: {sr}')
        self.session, self.sr = session, sr
        self.frame = sr * 32 // 1000
        self.context = np.zeros((1, sr * 4 // 1000), dtype=np.float32)
        self.state = np.zeros((2, 1, 128), dtype=np.float32)

    def step(self, samples):
        if len(samples) != self.frame:
            raise ValueError('Only complete 32 ms frames may be submitted')
        x = np.concatenate([self.context, samples[None, :]], axis=1)
        output, self.state = self.session.run(None, {
            'input': x, 'state': self.state, 'sr': np.array(self.sr, dtype=np.int64)})
        self.context = x[:, -self.context.shape[1]:].copy()
        speech = float(output.item())
        if not np.isfinite(speech) or not 0 <= speech <= 1:
            raise ValueError('Invalid model output')
        return speech


def main():
    weight = HERE / 'models/silero_vad.onnx'
    if hashlib.sha256(weight.read_bytes()).hexdigest() != WEIGHT_SHA256:
        raise ValueError('Silero model checksum mismatch')
    options = ort.SessionOptions()
    options.intra_op_num_threads = options.inter_op_num_threads = 1
    session = ort.InferenceSession(str(weight), sess_options=options, providers=['CPUExecutionProvider'])
    label_path = ROOT / 'labels/EOT标注/工作数据/eot_labels.csv'
    with label_path.open(encoding='utf-8-sig', newline='') as f:
        labels = list(csv.DictReader(f))
    calls, checks = [], []
    for call in sorted({r['call_id'] for r in labels}):
        path = ROOT / 'data/韩泰机器人抢话/音频' / f'{call}.mp3'
        pcm, sr = sf.read(path, dtype='int16', always_2d=True)
        stream = SileroStream(session, sr)  # reset only between recordings
        user = np.ascontiguousarray(pcm[:, 0].astype(np.float32) / 32768.)
        points = []
        for end in range(stream.frame, len(user) + 1, stream.frame):
            speech = stream.step(user[end-stream.frame:end])
            points.append([end * 1000 // sr, round(1 - speech, 7)])
        times = [p[0] for p in points]
        checkpoints = []
        for row in labels:
            if row['call_id'] != call:
                continue
            decision = int(row['decision_ms'])
            if decision > len(user) * 1000 / sr:
                raise ValueError('Annotation beyond recording')
            index = bisect.bisect_right(times, decision) - 1
            if index < 0:
                continue  # No complete frame is available yet; never invent zero.
            available, score = points[index]
            checkpoints.append([decision, score])
            checks.append(dict(call_id=call, decision_ms=decision, available_ms=available,
                               non_speech_score=score, speech_score=round(1-score, 7)))
        calls.append(dict(call_id=call, audio_path=f'韩泰机器人抢话/音频/{call}.mp3',
            audio_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), duration_ms=round(len(user)*1000/sr,3),
            sample_rate=sr, channel=0, points=points, checkpoints=sorted(checkpoints)))
        # A fresh stream reproduces a prefix without any later samples.
        witness = SileroStream(session, sr)
        for index in range(min(50, len(points))):
            end = (index+1)*stream.frame
            assert round(1-witness.step(user[end-stream.frame:end]),7) == points[index][1]
        print(f'{call}: {len(points)} frames, {len(checkpoints)} checkpoints; prefix witness passed', flush=True)
    output = dict(schema_version=1, model='silero_vad', name='Silero VAD',
        score='1 - speech_probability (non-speech baseline, not semantic EOT)', sample_interval_ms=32,
        method='CPU, one thread. User channel 0, native 8/16 kHz, int16 / 32768. Sequential complete 32ms frames plus 4ms past context; recurrent state reset per call. Timestamp is frame END. No future audio, resampling, silence timeout, hangover or annotation text. Checkpoints hold last available complete frame; trailing partial frame is not padded.',
        model_source='Silero ONNX bundled with pipecat-ai==1.12.0; upstream https://github.com/snakers4/silero-vad',
        model_sha256=WEIGHT_SHA256, annotation_sha256=hashlib.sha256(label_path.read_bytes()).hexdigest(),
        script_sha256=hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        environment=dict(machine=platform.platform(), architecture=platform.machine(), python=sys.version,
            dependencies={p:importlib.metadata.version(p) for p in ['numpy','onnxruntime','soundfile']},
            provider='CPUExecutionProvider', threads=1), calls=calls)
    out = HERE/'results'
    out.mkdir(exist_ok=True)
    (out/'score-timeline.json').write_text(json.dumps(output,ensure_ascii=False,separators=(',',':'))+'\n')
    with (out/'checkpoints.csv').open('w',encoding='utf-8',newline='') as f:
        writer=csv.DictWriter(f,fieldnames=list(checks[0]));writer.writeheader();writer.writerows(checks)


if __name__ == '__main__':
    main()