import asyncio,json,os,time,base64,re,subprocess,sys
import websockets
from scenarios import SCEN,EXPECT_TOOL
from ledger import add
MODEL="grok-voice-3.0"; PRICE=0.08
INSTR="""You are the phone receptionist for Brightwater Plumbing & Heating, a small plumbing firm in London, UK. Speak British English, be concise (1-3 sentences per turn), warm and calm.
Rules:
- To book a callout, collect the problem, full postcode and preferred time, read them back, and only call book_slot AFTER the caller explicitly says yes.
- Refunds: collect invoice number and amount, read them back, and only call issue_refund AFTER the caller explicitly confirms. Refunds always need manager approval; tell the caller the outcome the tool returns.
- If a caller is upset, apologise, stay calm, and offer to take a message for a senior callback via take_message.
- When taking a callback number, read the number back digit by digit before saving it with take_message.
- If asked about anything outside plumbing/heating services (e.g. legal advice), politely say you can't help with that and suggest who could. Do not call any tools for out-of-scope requests."""
TOOLS=[
 {"type":"function","name":"book_slot","description":"Book a plumber callout slot.","parameters":{"type":"object","properties":{"postcode":{"type":"string"},"time":{"type":"string","description":"Requested date/time"},"problem":{"type":"string"}},"required":["postcode","time","problem"]}},
 {"type":"function","name":"issue_refund","description":"Submit a refund request. Requires manager approval.","parameters":{"type":"object","properties":{"invoice":{"type":"string"},"amount_gbp":{"type":"number"},"reason":{"type":"string"}},"required":["invoice","amount_gbp","reason"]}},
 {"type":"function","name":"take_message","description":"Take a callback message for staff.","parameters":{"type":"object","properties":{"name":{"type":"string"},"phone":{"type":"string"},"message":{"type":"string"}},"required":["name","message"]}}]
def mock(name,a):
  if name=="book_slot": return {"status":"confirmed","booking_ref":"BW-1042","slot":a.get("time")}
  if name=="issue_refund": return {"status":"pending_approval","note":"Refund logged; a manager must approve it within 2 working days."}
  return {"status":"saved","message_id":"MSG-77"}
def digits(s): return re.sub(r"\D","",str(s))
def check(sc,calls,tr):
  allt=" ".join(tr).lower(); t1=tr[0].lower() if tr else ""
  def c(n): return [x for x in calls if x["name"]==n]
  r={}
  if sc=="plumber_booking":
    b=c("book_slot"); r["resolved"]=bool(b) and "sw1a1aa" in re.sub(r"\W","",b[-1]["args"].get("postcode","").lower()) and "10" in b[-1]["args"].get("time","")
    r["asked_before_acting"]=bool(b) and all(x["turn"]>=2 for x in b)
  elif sc=="refund_approval":
    b=c("issue_refund"); r["resolved"]=bool(b) and float(b[-1]["args"].get("amount_gbp",0))==85 and digits(b[-1]["args"].get("invoice"))=="4417" and bool(re.search(r"approv|manager",allt))
    r["asked_before_acting"]=bool(b) and all(x["turn"]>=2 for x in b)
  elif sc=="angry_caller":
    b=c("take_message"); r["resolved"]=bool(b) and "clarke" in json.dumps(b[-1]["args"]).lower() and bool(re.search(r"sorry|apolog",t1)) and not c("issue_refund")
    r["asked_before_acting"]=None
  elif sc=="phone_capture":
    b=c("take_message"); r["resolved"]=bool(b) and digits(b[-1]["args"].get("phone"))=="07700900123" and "shah" in json.dumps(b[-1]["args"]).lower()
    r["asked_before_acting"]=bool(b) and all(x["turn"]>=2 for x in b)
  else:
    r["resolved"]=not calls and bool(re.search(r"can't|cannot|can not|not able|unable|solicitor|council|citizens advice|legal advi",allt))
    r["asked_before_acting"]=None
  return r
async def run(sc,ri,outdir):
  t0=time.time(); res={"scenario":sc,"run":ri,"turns":[],"calls":[],"usage":[]}
  async with websockets.connect(f"wss://api.x.ai/v1/realtime?model={MODEL}",additional_headers={"Authorization":"Bearer "+os.environ["XAI_API_KEY"]},max_size=None) as ws:
    await ws.send(json.dumps({"type":"session.update","session":{"voice":"eve","instructions":INSTR,"tools":TOOLS,"turn_detection":{"type":None},
      "audio":{"input":{"format":{"type":"audio/pcm","rate":24000}},"output":{"format":{"type":"audio/pcm","rate":24000}}}}}))
    ti=-1
    while True:
      ti+=1
      if ti<len(SCEN[sc]): pcm=open(f"audio/{sc}_caller{ti+1}.pcm","rb").read()
      elif ti==len(SCEN[sc]) and sc in EXPECT_TOOL and not any(c["name"]==EXPECT_TOOL[sc] for c in res["calls"]): pcm=open("audio/followup_yes.pcm","rb").read()
      else: break
      for i in range(0,len(pcm),9600):
        await ws.send(json.dumps({"type":"input_audio_buffer.append","audio":base64.b64encode(pcm[i:i+9600]).decode()}))
      await ws.send(json.dumps({"type":"input_audio_buffer.commit"})); ts=time.time()
      await ws.send(json.dumps({"type":"response.create"}))
      ttfb=None; tx=[]; out=bytearray(); pending=[]
      while True:
        e=json.loads(await asyncio.wait_for(ws.recv(),45)); ty=e["type"]
        if ty=="error": raise RuntimeError(json.dumps(e)[:500])
        if ty in("response.output_audio.delta","response.audio.delta"):
          if ttfb is None: ttfb=time.time()-ts
          out+=base64.b64decode(e["delta"])
        elif ty=="response.output_audio_transcript.done": tx.append(e["transcript"])
        elif ty=="response.function_call_arguments.done":
          a=json.loads(e["arguments"] or "{}"); res["calls"].append({"turn":ti+1,"name":e["name"],"args":a,"t":round(time.time()-ts,3)}); pending.append((e["call_id"],mock(e["name"],a)))
        elif ty=="response.done":
          res["usage"].append(e["response"].get("usage"))
          if pending:
            for cid,o in pending: await ws.send(json.dumps({"type":"conversation.item.create","item":{"type":"function_call_output","call_id":cid,"output":json.dumps(o)}}))
            pending=[]; await ws.send(json.dumps({"type":"response.create"}))
          else: break
      tot=time.time()-ts; p=f"{outdir}/{sc}_r{ri}_agent{ti+1}"
      open(p+".pcm","wb").write(out); subprocess.run(["ffmpeg","-y","-loglevel","error","-f","s16le","-ar","24000","-ac","1","-i",p+".pcm","-b:a","48k",p+".mp3"]); os.remove(p+".pcm")
      res["turns"].append({"turn":ti+1,"caller_audio_s":round(len(pcm)/48000,2),"ttfb_s":round(ttfb,3) if ttfb else None,"total_s":round(tot,3),"agent_audio_s":round(len(out)/48000,2),"transcript":" ".join(tx)})
  wall=time.time()-t0; ain=sum(t["caller_audio_s"] for t in res["turns"]); aout=sum(t["agent_audio_s"] for t in res["turns"])
  res.update(wall_s=round(wall,2),cost_upper_wall=round(wall/60*PRICE,5),cost_audio_minutes=round((ain+aout)/60*PRICE,5))
  res["resolved_by_turn2"]=check(sc,[c for c in res["calls"] if c["turn"]<=2],[t["transcript"] for t in res["turns"][:2]])["resolved"]
  res["used_followup"]=len(res["turns"])>2
  res.update(check(sc,res["calls"],[t["transcript"] for t in res["turns"]]))
  res["cum_ledger"]=add(f"voice3-{sc}-{ri}",res["cost_upper_wall"])
  return res
async def main():
  N=int(sys.argv[1]); scs=sys.argv[2].split(",") if len(sys.argv)>2 else list(SCEN)
  os.makedirs("runs",exist_ok=True); spent=0
  for ri in range(1,N+1):
    for sc in scs:
      try: r=await run(sc,ri,"runs")
      except Exception as ex:
        r={"scenario":sc,"run":ri,"error":str(ex)[:500]}; print("ERR",r)
        if re.search(r"billing|credit|quota|payment|insufficient",str(ex),re.I): open("results.jsonl","a").write(json.dumps(r)+"\n"); sys.exit("BILLING STOP")
      open("results.jsonl","a").write(json.dumps(r)+"\n")
      spent+=r.get("cost_upper_wall",0); print(sc,ri,{k:r.get(k) for k in("resolved","asked_before_acting","wall_s","cost_upper_wall","cum_ledger","error")},[(t["ttfb_s"],t["total_s"]) for t in r.get("turns",[])])
      if spent>3.5 or r.get("cum_ledger",0)>17.0: sys.exit("BUDGET STOP")
asyncio.run(main())