import array
import threading
import time
import sanas
from examples.wav_utils import read_wav, save_wav  # or copy these helpers
 
# 1. Create + activate (activation is synchronous and returns an SdkResult).
sdk = sanas.create_sdk(sanas.InitParams(storage_dir="./storage"))
res = sdk.activate_api_key("YOUR_API_KEY")
if not res.success:
    raise RuntimeError(f"activation failed ({res.error_type}): {res.message}")
 
# 2. Read input audio as interleaved float32.
float32_bytes, sample_rate, channels = read_wav("input.wav")
samples = array.array("f"); samples.frombytes(float32_bytes)
 
# 3. Get notified when the pipeline is ready. The state callback fires on a
#    background thread, so publish each state under a Condition and let the
#    main thread wait for a terminal state (RUNNING = go, NOT_RUNNING = failed).
cond = threading.Condition()
state = {"value": None}
def on_state(s):
    with cond:
        state["value"] = s
        cond.notify_all()
 
attrs = sanas.ProcessorAttributes(
    audio_attributes=sanas.AudioAttributes(
        sampling_rate=sample_rate,
        channels=channels,
        model_name="<your-model-key>",  # e.g. "AT5.2", "SE2.2", "VI_G_SE"
        audio_pipeline_state_notify=on_state,
    )
)
 
# 4. Feed 20 ms frames at real time; process_frame returns the processed frame.
frame_len = sample_rate * channels // 50  # samples in 20 ms
frame_period = 0.020  # seconds per frame
out_bytes = []
_terminal = (sanas.PipelineState.RUNNING, sanas.PipelineState.NOT_RUNNING)
with sdk.create_audio_processor(attrs) as proc:
    with cond:  # wait until the pipeline settles
        cond.wait_for(lambda: state["value"] in _terminal, timeout=30)
        if state["value"] != sanas.PipelineState.RUNNING:
            raise RuntimeError(f"pipeline did not start (state: {state['value']})")
 
    # `deadline` is the wall-clock time the NEXT frame should be sent.
    deadline = time.monotonic()
    for i in range(0, len(samples) - frame_len + 1, frame_len):
        frame = proc.process_frame(sanas.AudioFrame(samples=samples[i:i + frame_len]))
        if frame.frame_count:
            out_bytes.append(bytes(frame.samples))  # memoryview is transient; copy
        deadline += frame_period
        time.sleep(max(0.0, deadline - time.monotonic()))
 
    # Drain the buffered tail by pushing silence (still paced at real time).
    silence = array.array("f", bytes(frame_len * 4))
    for _ in range(10):
        frame = proc.process_frame(sanas.AudioFrame(samples=silence))
        if frame.frame_count:
            out_bytes.append(bytes(frame.samples))
        deadline += frame_period
        time.sleep(max(0.0, deadline - time.monotonic()))
 
save_wav("output.wav", b"".join(out_bytes), sample_rate, channels)