#!/usr/bin/env python3
"""WebSocket Groq Whisper proxy — chunked text frames (DPI-proof)."""
import asyncio, json, websockets, requests, uuid, sys, time, base64

GROQ_BASE = "https://api.groq.com"

def log(msg):
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", file=sys.stderr, flush=True)

async def handler(ws):
    chunks = {}   # index -> base64_part
    total = None
    token = ""
    mime = "audio/webm"
    try:
        log(f"conn from {ws.remote_address}")
        async for raw in ws:
            msg = json.loads(raw)
            ci = msg.get("c", 0)

            if ci == -1:
                # Final chunk — reassemble
                if not chunks:
                    await ws.send(json.dumps({"error": "no chunks"}))
                    return
                audio_b64 = "".join(chunks[i] for i in sorted(chunks))
                log(f"assembled {len(chunks)} chunks → {len(audio_b64)} chars, decoding...")
                audio_data = base64.b64decode(audio_b64)
                log(f"decoded {len(audio_data)}B")

                ext = "webm"
                if "ogg" in mime: ext = "ogg"
                elif "mp4" in mime or "aac" in mime: ext = "mp4"

                boundary = "----WKFBoundary" + uuid.uuid4().hex[:16]
                payload = b"".join([
                    f'--{boundary}\r\nContent-Disposition: form-data; name="file"; filename="audio.{ext}"\r\nContent-Type: {mime}\r\n\r\n'.encode(),
                    audio_data,
                    f'\r\n--{boundary}\r\nContent-Disposition: form-data; name="model"\r\n\r\nwhisper-large-v3'.encode(),
                    f'\r\n--{boundary}\r\nContent-Disposition: form-data; name="language"\r\n\r\nit'.encode(),
                    f'\r\n--{boundary}--\r\n'.encode(),
                ])

                t0 = time.time()
                loop = asyncio.get_event_loop()
                r = await loop.run_in_executor(None, lambda: requests.post(
                    f"{GROQ_BASE}/openai/v1/audio/transcriptions",
                    headers={
                        "Authorization": f"Bearer {token}",
                        "Content-Type": f"multipart/form-data; boundary={boundary}",
                    },
                    data=payload, timeout=60
                ))
                log(f"Groq {r.status_code} in {time.time()-t0:.2f}s: {r.text[:120]}")
                await ws.send(json.dumps({"text": r.json().get("text", ""), "status": r.status_code}))
                return
            else:
                # Data chunk
                chunks[ci] = msg.get("b", "")
                token = msg.get("t", token)
                mime = msg.get("m", mime)
                total = msg.get("n", total)
                # log(f"chunk {ci+1}/{total}")

    except Exception as e:
        log(f"ERR: {e}")
        try: await ws.send(json.dumps({"error": str(e)}))
        except: pass

async def main():
    async with websockets.serve(handler, "127.0.0.1", 8766):
        log("groq-ws :8766 (chunked)")
        await asyncio.Future()

asyncio.run(main())