"""Offline duplex driver: runs scripted conversations (or a batch eval) through the SAME components as the live server,
on a simulated clock. Mic audio is fed in 32 ms frames; compute stages are timed on the real GPU and placed on the timeline,
so latency = (end-of-turn silence wait) + ASR + mix/LID + LLM/validator + first TTS chunk.
  python offline.py convo scenarios.json OUTDIR
  python offline.py eval  list.jsonl OUTDIR [N]"""
import os, sys, json, time, re, numpy as np, soundfile as sf, librosa, httpx
from turn import TurnDetector, FR, SR, FMS
from respond import Responder
from tts_engine import TTS

ASR_URL = os.environ.get("ASR_URL", "http://127.0.0.1:8002/asr")
_cli = httpx.Client(timeout=30)
def asr(y):
    t = time.perf_counter(); r = _cli.post(ASR_URL, content=np.ascontiguousarray(y, np.float32).tobytes()).json()
    return r["text"], time.perf_counter() - t

BC = {"eng": "<eng>Mm-hmm.</eng>", "pcm": "<pcm>Ehen.</pcm>", "ibo": "<ibo>Ehen.</ibo>", "yor": "<yor>Ehen.</yor>", "hau": "<hau>To.</hau>"}
def trim(w, sr, db=35):
    yt, idx = librosa.effects.trim(w, top_db=db); return w[max(0, idx[0] - int(0.03 * sr)): idx[1] + int(0.05 * sr)]

class Sim:
    def __init__(self, tts, resp, eot_ms=600):
        self.tts, self.resp, self.eot_ms = tts, resp, eot_ms
        self.osr = tts.sr
        self.bc = {k: trim(tts.synth(v, seed=3), self.osr) for k, v in BC.items()}
        import mix as M; M.estimate("warm up the GlotLID model", asr=True)
        import logging; logging.getLogger("httpx").setLevel(logging.WARNING); logging.getLogger("httpcore").setLevel(logging.WARNING)

    def run(self, turns, label=""):
        td = TurnDetector(eot_ms=self.eot_ms)
        T = 600; mic = np.zeros(int(T * SR), np.float32); agt = np.zeros(int(T * self.osr), np.float32)
        f = 0; busy_until = 0.0; reply_start = None; hist = []; log = []; t_next = 0.5; cur_chunks = []; style_lang = "eng"
        for k, tn in enumerate(turns):
            u = tn["audio"]
            if tn.get("barge_at") is not None and reply_start is not None:
                start = reply_start + tn["barge_at"]
            else:
                start = max(t_next, busy_until + tn.get("gap", 0.7))
            s0 = int(start * SR); mic[s0:s0 + len(u)] += u
            e = np.abs(u) > 0.02; onset = start + (np.argmax(e) / SR if e.any() else 0)
            rec = dict(turn=k, user_label=tn.get("label", ""), user_ref=tn.get("text", ""), user_start=round(start, 3), user_onset=round(onset, 3), backchannels=[])
            ev_eot = None
            while ev_eot is None:
                fr = mic[f * FR:(f + 1) * FR]; f += 1; tnow = f * FMS
                busy = tnow < busy_until
                for name, info in td.feed(fr, agent_busy=busy):
                    if name == "barge" and busy:
                        t0 = time.perf_counter(); stop_overhead = time.perf_counter() - t0 + 0.040   # + 40 ms output buffer flush
                        t_stop = tnow + stop_overhead
                        a0 = int(t_stop * self.osr); fade = int(0.02 * self.osr)
                        agt[a0:a0 + fade] *= np.linspace(1, 0, fade)[: len(agt[a0:a0 + fade])]; agt[a0 + fade:] = 0
                        played = [c for c in cur_chunks if c[1] < t_stop]
                        if hist:
                            hist[-1] = (hist[-1][0], " ".join(c[0] for c in played) + " [interrupted]")
                        rec["barge"] = dict(user_onset=round(onset, 3), detected=round(tnow, 3), agent_stopped=round(t_stop, 3), reaction_s=round(t_stop - onset, 3),
                                            agent_chunks_played=len(played), agent_chunks_total=len(cur_chunks))
                        busy_until = t_stop
                    elif name == "backchannel" and not busy:
                        w = self.bc.get(style_lang, self.bc["eng"]); a0 = int((tnow + 0.05) * self.osr); agt[a0:a0 + len(w)] += 0.8 * w
                        rec["backchannel_at"] = round(tnow, 3); rec["backchannels"].append(dict(t=round(tnow, 3), turn_s=round(info["turn_s"], 2)))
                    elif name == "eot":
                        ev_eot = (tnow, info)
            T_eot, info = ev_eot
            text, t_asr = asr(info["audio"])
            t_start = info["speech_end"] + (td.spec_f * FMS if td.spec_f else (T_eot - info["speech_end"]))   # speculative start at spec_ms of silence
            r = self.resp.reply(text, hist)
            style_lang = r["target"]["matrix"]
            t_cmp = t_asr + r["t_total"]
            prev_end = 0; cum = 0; first = None; cur_chunks = []; tts_times = []
            for c, w, ds in self.tts.stream(r["reply"], seed=k):
                cum += ds; tts_times.append(round(ds, 3))
                avail = max(T_eot, t_start + t_cmp + cum); ps = max(avail, prev_end)
                a0 = int(ps * self.osr); agt[a0:a0 + len(w)] += w
                prev_end = ps + len(w) / self.osr; cur_chunks.append((c, ps, prev_end))
                if first is None: first = ps
            reply_start = first; busy_until = prev_end; t_next = prev_end
            hist.append((text, r["reply"]))
            rec.update(asr=text, asr_s=round(t_asr, 3), eot_wait_s=round(T_eot - info["speech_end"], 3), speech_end=round(info["speech_end"], 3),
                       target=r["target"], target_desc=r["target_desc"], reply=r["reply"], reply_mix=r["reply_mix"], tvd=r["tvd"], how=r["how"],
                       first_cand_ok=r["first_cand_ok"], rounds=r["rounds"], mix_s=r["t_mix"], llm_s=r["t_llm"], tts_first_s=tts_times[0] if tts_times else None, spec_start_s=round(t_start - info["speech_end"], 3), compute_to_first_audio_s=round(t_cmp + (tts_times[0] if tts_times else 0), 3),
                       tts_breakdown=getattr(self.tts, "last", None),
                       tts_chunks_s=tts_times, first_audio=round(first, 3) if first else None, reply_end=round(prev_end, 3),
                       latency_s=round(first - info["speech_end"], 3) if first else None)
            log.append(rec); print(json.dumps({k2: rec[k2] for k2 in ("turn", "asr", "reply", "latency_s", "how")}, ensure_ascii=False), flush=True)
        end = max(busy_until, f * FMS) + 0.6
        mic24 = librosa.resample(mic[: int(end * SR)], orig_sr=SR, target_sr=self.osr)
        n = min(len(mic24), int(end * self.osr)); a = agt[:n]; m = mic24[:n]
        return dict(log=log, mono=np.clip(0.9 * (m + a), -1, 1), stereo=np.stack([m, a], 1), sr=self.osr)

def user_audio(tts, tn):
    if tn.get("wav"):
        y, sr = sf.read(tn["wav"], dtype="float32"); y = y.mean(1) if y.ndim > 1 else y
        return librosa.resample(y, orig_sr=sr, target_sr=SR) if sr != SR else y
    parts = tn["text_parts"] if "text_parts" in tn else [tn["text"]]
    ws = []
    for i, p in enumerate(parts):
        w = trim(tts.synth(p, seed=11 + i, voice=tn.get("voice", "waxal_ibo_5_female")), tts.sr)
        ws += [w, np.zeros(int(tn.get("inner_pause", 0.35) * tts.sr), np.float32)]
    y = np.concatenate(ws[:-1]); return librosa.resample(y, orig_sr=tts.sr, target_sr=SR).astype(np.float32) * 0.9

def save(out, name, res, meta):
    os.makedirs(out, exist_ok=True)
    sf.write(f"{out}/{name}.wav", res["mono"], res["sr"]); sf.write(f"{out}/{name}_stereo_userL_agentR.wav", res["stereo"], res["sr"])
    json.dump(dict(meta, turns=res["log"]), open(f"{out}/{name}.json", "w"), ensure_ascii=False, indent=1, default=str)

if __name__ == "__main__":
    mode, src, out = sys.argv[1], sys.argv[2], sys.argv[3]
    tts = TTS(os.environ.get("AGENT_VOICE", "anchor_male")); resp = Responder(); sim = Sim(tts, resp)
    if mode == "convo":
        for sc in json.load(open(src)):
            resp.style = None
            turns = [dict(t, audio=user_audio(tts, t)) for t in sc["turns"]]
            res = sim.run(turns, sc["name"]); save(out, sc["name"], res, {k: v for k, v in sc.items() if k != "turns"} | dict(user_turns=[{k: v for k, v in t.items() if k != "audio"} for t in turns]))
    else:   # eval: single-turn items from the ASR eval list (real held-out audio), each followed by a barge-in at 1.0 s
        R = [json.loads(l) for l in open(src)]; N = int(sys.argv[4]) if len(sys.argv) > 4 else 40
        import random; random.seed(1); R = random.sample(R, min(N, len(R))); allr = []
        for i, r in enumerate(R):
            resp.style = None
            barger = R[(i + 1) % len(R)]
            turns = [dict(wav=r["wav"], label=r["id"], text=r["text"]), dict(wav=barger["wav"], label=barger["id"] + " (barge-in)", text=barger["text"], barge_at=1.0)]
            turns = [dict(t, audio=user_audio(tts, t)) for t in turns]
            res = sim.run(turns, r["id"]); allr.append(dict(id=r["id"], lang=r["lang"], ref=r["text"], turns=res["log"]))
            json.dump(allr, open(f"{out}/eval_turns.json", "w"), ensure_ascii=False, indent=1, default=str)
    print("OFFLINE_DONE")