#!/usr/bin/env python3
"""Mock AI Agent — FreeSWITCH/ESL 版.

方案: FreeSWITCH 无 ARI externalMedia 式外挂媒体 API, 标准做法是本 agent 自身充当
SIP UAS 端点 (等价于 asterisk 版的 externalMedia 腿), 控制面走 ESL:
  0. ESL TCP 连接 fs:8021 (auth)
  1. 本侧 SIP UAS 监听 (5062) + RTP 监听 (40001)
  2. ESL originate 腿1: sofia/external/agent@<本容器>:5062 &park()   ← 本侧 UAS 应答
  3. ESL originate 腿2: sofia/gateway/mock-trunk/<被叫> &park()      ← mock-provider 应答
  4. ESL uuid_bridge 腿1 腿2 → 桥接, RTP 经 FreeSWITCH 双向
  5. RTP: RX=被叫音频(模拟 ASR 输入, RMS 统计); TX=440Hz 间歇音(模拟 LLM/TTS 输出)
  6. ESL uuid_kill 双腿 → BYE, 输出双向统计
纯 stdlib (asyncio + socket)。换真实 LLM/ASR: 替换 TX 音源与 RX 消费即可。
用法: mock-agent.py --esl-host fs1 --sip-host <本容器名>:5062 \
      --number +15105550123 [--duration 15]
"""
import argparse, asyncio, math, os, random, socket, struct, time

def log(*a): print(f"[agent {time.strftime('%H:%M:%S')}]", *a, flush=True)

# ---------- G.711 编码 ----------
def ulaw_encode(s: int) -> int:
    s = max(-32635, min(32635, s)); sign = 0x80 if s < 0 else 0
    if sign: s = -s
    s += 0x84
    e = 7
    for i in range(7, -1, -1):
        if s & (0x40 << i): e = i; break
    return (~(sign | (e << 4) | ((s >> (e + 3)) & 0x0F))) & 0xFF

def alaw_encode(s: int) -> int:
    s = max(-32635, min(32635, s)); sign = 0x80 if s < 0 else 0
    if sign: s = -s
    e = 7
    for i in range(7, -1, -1):
        if s & (0x40 << i): e = i; break
    m = (s >> (e + 3)) & 0x0F
    return ((sign | (e << 4) | m) ^ 0x55) & 0xFF

def ulaw_decode(u: int) -> int:
    u = ~u & 0xFF
    sign = u & 0x80; e = (u >> 4) & 0x07; m = u & 0x0F
    s = ((m << (e + 3)) + 0x84) << 2
    return -s if sign else s

def alaw_decode(a: int) -> int:
    a ^= 0x55
    t = (a & 0x0F) << 4
    seg = (a & 0x70) >> 4
    if seg == 0: t += 8
    else:
        t += 0x108
        if seg > 1: t <<= (seg - 1)
    return -t if a & 0x80 else t

def tone_payload(codec: str, ts: float) -> bytes:
    """20ms / 160 samples @8kHz, 440Hz 500ms-on 500ms-off (模拟 TTS)"""
    enc = ulaw_encode if codec == "PCMU" else alaw_encode
    out = bytearray()
    on = (ts % 1.0) < 0.5
    for i in range(160):
        t = ts + i / 8000.0
        v = int(6000 * math.sin(2 * math.pi * 440 * t)) if on else 0
        out.append(enc(v))
    return bytes(out)

# ---------- SIP 报文 (与 mock-provider 同源) ----------
def parse_msg(data: bytes):
    head, _, body = data.decode("utf-8", "replace").partition("\r\n\r\n")
    lines = head.split("\r\n")
    headers = []
    for ln in lines[1:]:
        n, _, v = ln.partition(":")
        headers.append((n.strip().lower(), v.strip()))
    return lines[0], headers, body

def hget(headers, name):
    for n, v in headers:
        if n == name: return v
    return None

def build_resp(first, headers, code, reason, tag=None, extra=None, body=""):
    out = [f"SIP/2.0 {code} {reason}"]
    for n, v in headers:
        if n == "via": out.append(f"Via: {v}")
    out.append(f"From: {hget(headers,'from')}")
    to = hget(headers, "to")
    if tag and "tag=" not in to: to = f"{to};tag={tag}"
    out.append(f"To: {to}")
    out.append(f"Call-ID: {hget(headers,'call-id')}")
    out.append(f"CSeq: {hget(headers,'cseq')}")
    if extra:
        out += extra
    if body:
        out.append("Content-Type: application/sdp")
    out.append(f"Content-Length: {len(body)}")
    return ("\r\n".join(out) + "\r\n\r\n" + body).encode()

def parse_sdp_offer(body: str):
    ip, port, pts = None, None, []
    for ln in body.splitlines():
        if ln.startswith("c=IN IP4 "): ip = ln.split()[2]
        elif ln.startswith("m=audio "):
            parts = ln.split(); port = int(parts[1]); pts = [p for p in parts[3:] if p.isdigit()]
    return ip, port, pts

def pick_codec(pts):
    if "0" in pts: return "0", "PCMU"
    if "8" in pts: return "8", "PCMA"
    return (pts[0], "PCMU") if pts else ("0", "PCMU")

# ---------- SIP UAS (被 FreeSWITCH originate 的"externalMedia 腿") ----------
class SipUas(asyncio.DatagramProtocol):
    """应答 FS 的 INVITE (200+SDP), 处理 ACK/BYE/OPTIONS/re-INVITE/重传."""
    def __init__(self, my_ip, sip_port, rtp_port):
        self.my_ip, self.sip_port, self.rtp_port = my_ip, sip_port, rtp_port
        self.transport = None
        self.tag = os.urandom(5).hex()
        self.cid = None
        self.fs_rtp = None            # FS 侧 RTP 地址 (INVITE SDP)
        self.pt, self.codec = "0", "PCMU"
        self.confirmed = False
        self.ended = asyncio.Event()

    def connection_made(self, transport):
        self.transport = transport

    def datagram_received(self, data, addr):
        first, headers, body = parse_msg(data)
        method = first.split()[0] if first else "?"
        if method == "INVITE":
            cid = hget(headers, "call-id")
            new_call = self.cid is None
            if new_call:
                self.cid = cid
                log(f"INVITE from FS {addr[0]}:{addr[1]} call_id={cid}")
            ip, port, pts = parse_sdp_offer(body)
            if ip and port:
                self.fs_rtp = (ip, port)   # re-INVITE 时同步更新
            if new_call:
                self.pt, self.codec = pick_codec(pts)
            # 重传与 re-INVITE 统一: 按当前请求头重建 200 OK (CSeq 正确)
            sdp = (f"v=0\r\no=agent {random.randint(1,10**9)} {random.randint(1,10**9)} IN IP4 {self.my_ip}\r\n"
                   f"s=agent\r\nc=IN IP4 {self.my_ip}\r\nt=0 0\r\n"
                   f"m=audio {self.rtp_port} RTP/AVP {self.pt}\r\n"
                   f"a=rtpmap:{self.pt} {self.codec}/8000\r\na=sendrecv\r\n")
            self.transport.sendto(build_resp(
                first, headers, 200, "OK", tag=self.tag,
                extra=[f"Contact: <sip:agent@{self.my_ip}:{self.sip_port}>",
                       "User-Agent: MockAgent/1.0"], body=sdp), addr)
        elif method == "ACK":
            if not self.confirmed:
                self.confirmed = True
                log(f"leg answered (ACK), fs rtp={self.fs_rtp} codec={self.codec}")
        elif method == "BYE":
            self.transport.sendto(build_resp(first, headers, 200, "OK", tag=self.tag), addr)
            self.ended.set()
        elif method == "OPTIONS":
            self.transport.sendto(build_resp(first, headers, 200, "OK", tag=self.tag), addr)

# ---------- RTP 媒体 ----------
class Rx(asyncio.DatagramProtocol):
    def __init__(self, stats):
        self.stats = stats

    def datagram_received(self, data, addr):
        if len(data) < 12: return
        pt = data[1] & 0x7F
        self.stats["rx"] += 1
        payload = data[12:]
        dec = [(ulaw_decode(b) if pt == 0 else alaw_decode(b)) for b in payload[:40]]
        self.stats["amp_sum"] += sum(abs(x) for x in dec) / max(1, len(dec))
        self.stats["tx_to"] = addr   # 始终跟随最新来源 (自适应媒体路径)
        if self.stats["rx"] % 250 == 0:
            log(f"RX frames={self.stats['rx']} avg_rms={self.stats['amp_sum']/self.stats['rx']:.0f}")

# ---------- ESL 客户端 (极简: auth + api, 同步命令式) ----------
class Esl:
    def __init__(self, host, port, password):
        self.host, self.port, self.password = host, port, password

    async def connect(self):
        self.r, self.w = await asyncio.open_connection(self.host, self.port)
        m = await self._read()
        if m.get("content-type") != "auth/request":
            raise ConnectionError(f"unexpected esl greeting: {m}")
        self.w.write(f"auth {self.password}\n\n".encode())
        await self.w.drain()
        m = await self._read()
        if "+OK" not in m.get("reply-text", ""):
            raise ConnectionError(f"esl auth failed: {m}")

    async def _read(self):
        headers = {}
        while True:
            line = await self.r.readline()
            if not line: raise ConnectionError("esl closed")
            line = line.decode("utf-8", "replace").rstrip("\r\n")
            if not line: break
            k, _, v = line.partition(":")
            headers[k.strip().lower()] = v.strip()
        n = int(headers.get("content-length", 0) or 0)
        headers["_body"] = (await self.r.readexactly(n)).decode("utf-8", "replace") if n else ""
        return headers

    async def api(self, cmd):
        self.w.write(f"api {cmd}\n\n".encode())
        await self.w.drain()
        return (await self._read())["_body"].strip()

    async def close(self):
        try: self.w.close()
        except Exception: pass

# ---------- 主流程 ----------
async def run(args):
    my_ip = socket.gethostbyname(socket.gethostname())
    sip_port = int(args.sip_host.rsplit(":", 1)[1])
    suffix = f"{int(time.time())}-{random.randint(100, 999)}"
    agent_uuid, callee_uuid = f"agent-{suffix}", f"callee-{suffix}"
    tmo = int(args.ring_timeout)

    stats = {"rx": 0, "amp_sum": 0.0, "tx": 0, "tx_to": None, "codec": "PCMU"}
    loop = asyncio.get_running_loop()
    uas_t, uas = await loop.create_datagram_endpoint(
        lambda: SipUas(my_ip, sip_port, args.rtp_port), local_addr=("0.0.0.0", sip_port))
    rtp_t, _ = await loop.create_datagram_endpoint(
        lambda: Rx(stats), local_addr=("0.0.0.0", args.rtp_port))
    log(f"agent up: sip=udp/{sip_port} rtp=udp/{args.rtp_port} ip={my_ip}")

    esl = Esl(args.esl_host, args.esl_port, args.esl_password)
    await esl.connect()
    log(f"esl connected: {args.esl_host}:{args.esl_port}")

    def fail(msg, detail=None):
        log(f"FAIL: {msg}", detail or "")
        raise SystemExit(1)

    try:
        # 腿1: 本 agent (SIP UAS, 等价 externalMedia): 应答后 park
        r = await esl.api(
            f"originate {{origination_uuid={agent_uuid},"
            f"origination_caller_id_number={args.caller_id},originate_timeout={tmo}}}"
            f"sofia/external/agent@{args.sip_host} &park()")
        if not r.startswith("+OK"): fail("originate agent leg", r)
        log(f"leg1 (agent UAS) up: {r}")

        # 腿2: 被叫经 mock-trunk 网关 (网关冷启动自动重试)
        for attempt in range(3):
            r = await esl.api(
                f"originate {{origination_uuid={callee_uuid},"
                f"origination_caller_id_number={args.caller_id},originate_timeout={tmo}}}"
                f"sofia/gateway/mock-trunk/{args.number} &park()")
            if r.startswith("+OK"): break
            log(f"leg2 originate retry {attempt + 1}: {r}")
            await asyncio.sleep(2)
        if not r.startswith("+OK"): fail("originate callee leg", r)
        log(f"leg2 (callee) answered: {r}")

        stats["codec"] = uas.codec
        if uas.fs_rtp and not stats["tx_to"]:
            stats["tx_to"] = uas.fs_rtp
            log(f"fs rtp peer (INVITE SDP): {uas.fs_rtp[0]}:{uas.fs_rtp[1]}")

        # 桥接双腿 → RTP 经 FS 双向
        r = await esl.api(f"uuid_bridge {agent_uuid} {callee_uuid}")
        if not r.startswith("+OK"): fail("uuid_bridge", r)
        log(f"bridged [{agent_uuid} + {callee_uuid}], streaming {args.duration}s "
            f"(RX=ASR mock, TX=440Hz LLM/TTS mock)")

        # 媒体收发
        async def tx():
            seq = random.randint(0, 65535); ts = random.randint(0, 10**6)
            ssrc = random.randint(1, 2**31)
            pt = 0 if stats["codec"] == "PCMU" else 8
            t0 = time.time()
            while time.time() - t0 < args.duration:
                if stats["tx_to"]:
                    hdr = struct.pack("!BBHII", 0x80, pt, seq & 0xFFFF,
                                      ts & 0xFFFFFFFF, ssrc)
                    rtp_t.sendto(hdr + tone_payload(stats["codec"], time.time() - t0),
                                 stats["tx_to"])
                    stats["tx"] += 1
                seq += 1; ts += 160
                await asyncio.sleep(0.02)
        await asyncio.wait_for(tx(), timeout=args.duration + 5)
    finally:
        for u in (callee_uuid, agent_uuid):
            try: await esl.api(f"uuid_kill {u}")
            except Exception: pass
        await esl.close()
        uas_t.close(); rtp_t.close()
        await asyncio.sleep(0.3)   # 等 BYE 应答落地

    n = stats["rx"]
    amp = stats["amp_sum"] / n if n else 0.0
    print(f"RESULT number={args.number} rx_frames={n} rx_avg_rms={amp:.0f} "
          f"tx_frames={stats['tx']} bidirectional={'YES' if n > 50 and amp > 100 else 'NO'}",
          flush=True)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--esl-host", default="fs1")
    ap.add_argument("--esl-port", type=int, default=8021)
    ap.add_argument("--esl-password", default="esl_mock_a95a037399153e9092b653fc")
    ap.add_argument("--sip-host", required=True, help="本侧 SIP 地址 host:port (供 FS INVITE)")
    ap.add_argument("--rtp-port", type=int, default=40001)
    ap.add_argument("--number", default="+15105550123")
    ap.add_argument("--caller-id", default="+15109990001")
    ap.add_argument("--duration", type=float, default=15)
    ap.add_argument("--ring-timeout", type=float, default=30)
    args = ap.parse_args()
    asyncio.run(run(args))

if __name__ == "__main__":
    main()