WhatsApp Sales Copilot

import base64
import io
import json
import re
import zipfile

# How many passages to retrieve per question.
RAG_TOP_K = [[[RAG_TOP_K|6]]]
# Chat messages per indexed sliding window.
CHAT_WINDOW = [[[CHAT_WINDOW|12]]]
# Cap on images OCR'd per export (each OCR pass costs seconds on CPU).
MAX_MEDIA = [[[MAX_MEDIA|40]]]

_STATE = {"messages": [], "participants": [], "artifacts": [], "skipped": []}
_PENDING_AUDIO = {}

# One WhatsApp message line, tolerant across exports:
#   iOS:     [01/07/2026, 09:12:03] Ana Souza: text
#   Android: 01/07/2026, 09:12 - Ana Souza: text     (also 12h "9:12 PM")
_MSG_RE = re.compile(
    r"^\[?([0-9]{1,2}[./-][0-9]{1,2}[./-][0-9]{2,4}),?\s+"
    r"([0-9]{1,2}:[0-9]{2}(?::[0-9]{2})?(?:\s?[APap][Mm])?)\]?"
    r"\s*(?:-\s*)?([^:]{1,40}):\s(.*)$"
)

# Artifact typing (filename + extracted text), first match wins — EN + PT-BR.
_KIND_RULES = [
    ("boarding pass", ("boarding pass", "cartão de embarque", "cartao de embarque",
                       "boarding", "gate", "assento", "seat")),
    ("payment proof", ("comprovante", "payment", "pagamento", "paid", "pago",
                       "pix", "transfer", "transferência", "transferencia", "receipt")),
    ("invoice", ("invoice", "fatura", "nota fiscal")),
    ("quote", ("quote", "quotation", "orçamento", "orcamento", "cotação", "cotacao")),
    ("reservation", ("reservation", "reserva", "booking", "locator", "localizador",
                     "pnr", "e-ticket", "eticket", "confirmation")),
    ("accommodation", ("hotel", "hospedagem", "hostel", "check-in", "checkin",
                       "accommodation")),
    ("identity document", ("passport", "passaporte", "identity", "identidade",
                           "cpf", "carteira de motorista", "driver's license")),
]

# Stage signals read from conversation text — stricter than _KIND_RULES
# (word-bounded) so ordinary chat words can't fake a stage.
_CHAT_STAGE_RULES = [
    ("boarding pass", re.compile(r"boarding pass|cart[ãa]o de embarque")),
    ("payment proof", re.compile(r"\bpaid\b|\bpago\b|comprovante|\bpix\b|receipt")),
    ("reservation", re.compile(r"\bbooked\b|reservation|reserva|locator|localizador|\bpnr\b")),
    ("quote", re.compile(r"\bquote\b|or[çc]amento|cota[çc][ãa]o")),
]

_STAGES = ["inquiry", "quoted", "booked", "paid", "ticketed"]
_NEXT_ACTION = {
    "inquiry": "Send a quote — no quote has gone out yet.",
    "quoted": "Follow up on the quote and confirm the booking.",
    "booked": "Request payment for the confirmed reservation.",
    "paid": "Issue the ticket and send it to the client.",
    "ticketed": "Send a pre-flight reminder with a documents checklist.",
}

_AMOUNT_RE = re.compile(
    r"(?:[$€£]|R\$|\bUSD\b|\bEUR\b|\bBRL\b)\s?[0-9][0-9.,]*"
    r"|\b[0-9]+[.,][0-9]{2}\b"
)

_IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp")
_AUDIO_EXTS = (".opus", ".ogg", ".m4a", ".mp3", ".wav", ".aac")

_DEMO_CHAT = """\
[01/07/2026, 09:12:03] Ana Souza: Hi! Do you have fares to Lisbon around Aug 14?
[01/07/2026, 09:15:40] Viva Agency: Good morning Ana! Yes — I will send a quote today.
[01/07/2026, 11:02:19] Viva Agency: Quote: GRU-LIS on 14 Aug, flight TP88, R$ 3.480,00 round trip. Valid for 48h.
[01/07/2026, 11:30:05] Ana Souza: Great, please book it!
[02/07/2026, 10:05:55] Viva Agency: Reservation confirmed — locator XKC9QD. Payment by pix, ok?
[02/07/2026, 10:20:31] Ana Souza: Paid! Just sent the receipt, R$ 3.480,00.
[02/07/2026, 10:44:12] Viva Agency: Received. I will issue the ticket and send it here.
"""


def _extract_pdf_text(pdf_bytes):
    """(JS-callable) Extract plain text from a PDF's raw bytes via pypdf."""
    from pypdf import PdfReader

    reader = PdfReader(io.BytesIO(pdf_bytes))
    return "\n\n".join((page.extract_text() or "") for page in reader.pages)


def _parse_chat(text):
    """Parse a WhatsApp chat export into message dicts (continuations merged)."""
    msgs = []
    for raw in text.splitlines():
        line = raw.replace("\u200e", "").strip()  # WhatsApp sprinkles LTR marks
        if not line:
            continue
        m = _MSG_RE.match(line)
        if m:
            date, time, sender, body = m.groups()
            msgs.append(
                {"date": date, "time": time, "sender": sender.strip(),
                 "text": body.strip()}
            )
        elif msgs:
            msgs[-1]["text"] += " " + line
    return msgs


def _classify_artifact(filename, text):
    """Type one media file from its name + extracted text (keyword rules)."""
    hay = (filename + " " + (text or "")[:2000]).lower()
    for kind, needles in _KIND_RULES:
        if any(n in hay for n in needles):
            return kind
    return "uncategorized"


def _stage():
    """Sales-pipeline stage from artifacts, plus conversation-text signals."""
    kinds = {a["kind"] for a in _STATE["artifacts"]}
    hay = " ".join(m["text"] for m in _STATE["messages"]).lower()
    for kind, rx in _CHAT_STAGE_RULES:
        if rx.search(hay):
            kinds.add(kind)
    for stage, kind in (
        ("ticketed", "boarding pass"),
        ("paid", "payment proof"),
        ("booked", "reservation"),
        ("quoted", "quote"),
    ):
        if kind in kinds:
            return stage
    return "inquiry"


async def _add_artifact(name, kind, text, indexed=True):
    entry = {"name": name, "kind": kind}
    m = _AMOUNT_RE.search(text or "")
    if m and kind in ("payment proof", "quote", "invoice"):
        entry["amount"] = m.group(0)
    _STATE["artifacts"].append(entry)
    if indexed and (text or "").strip():
        await agentop_rag.add_document(f"{kind} — {name}", text)


async def _index_chat(msgs, source):
    """Index the conversation as sliding windows of CHAT_WINDOW messages."""
    for start in range(0, len(msgs), CHAT_WINDOW):
        window = msgs[start : start + CHAT_WINDOW]
        lines = [
            f"{m['sender']} ({m['date']} {m['time']}): {m['text']}" for m in window
        ]
        label = f"chat {source} — {window[0]['date']} {window[0]['time']}"
        await agentop_rag.add_document(label, "\n".join(lines))


async def _ingest_chat_text(source, text):
    """(JS-callable) Parse + index one chat .txt. Returns the message count."""
    msgs = _parse_chat(text)
    if msgs:
        _STATE["messages"].extend(msgs)
        for m in msgs:
            if m["sender"] not in _STATE["participants"]:
                _STATE["participants"].append(m["sender"])
        await _index_chat(msgs, source)
    elif text.strip():
        # Unrecognized format — still make it searchable.
        await agentop_rag.add_document(f"chat {source}", text)
    return len(msgs)


async def _ocr_image(lowname, data):
    if lowname.endswith(".png"):
        mime = "image/png"
    elif lowname.endswith(".webp"):
        mime = "image/webp"
    else:
        mime = "image/jpeg"
    url = f"data:{mime};base64,{base64.b64encode(data).decode()}"
    return (await agentop_ml.ocr(url)) or ""


async def _ingest_export(name, data):
    """(JS-callable) Ingest one WhatsApp export .zip (chat + media).

    Returns JSON {"pending_audio": [names]} — voice notes must be decoded by
    the browser (AudioContext) before Python can transcribe them.
    """
    zf = zipfile.ZipFile(io.BytesIO(bytes(data)))
    media_done = 0
    for info in zf.infolist():
        if info.is_dir():
            continue
        base = info.filename.rsplit("/", 1)[-1]
        low = base.lower()
        try:
            if low.endswith(".txt"):
                await _ingest_chat_text(base, zf.read(info).decode("utf-8", "replace"))
            elif low.endswith(".pdf"):
                text = _extract_pdf_text(zf.read(info))
                kind = _classify_artifact(base, text)
                # Identity documents are counted, never indexed (privacy).
                await _add_artifact(
                    base, kind, text, indexed=(kind != "identity document")
                )
            elif low.endswith(_IMAGE_EXTS):
                if media_done >= MAX_MEDIA:
                    _STATE["skipped"].append(base)
                    continue
                media_done += 1
                text = await _ocr_image(low, zf.read(info))
                kind = _classify_artifact(base, text)
                await _add_artifact(
                    base, kind, text, indexed=(kind != "identity document")
                )
            elif low.endswith(_AUDIO_EXTS):
                _PENDING_AUDIO[base] = zf.read(info)
            else:
                _STATE["skipped"].append(base)
        except Exception as exc:  # one bad file must not kill the export
            _STATE["skipped"].append(f"{base} ({exc})")
    return json.dumps({"pending_audio": list(_PENDING_AUDIO)})


def _audio_b64(name):
    """(JS-callable) Pending voice-note bytes as base64 for browser decoding."""
    return base64.b64encode(_PENDING_AUDIO[name]).decode()


async def _transcribe_voice_note(name, samples):
    """(JS-callable) Transcribe decoded 16 kHz mono samples; index the text."""
    _PENDING_AUDIO.pop(name, None)
    text = (await agentop_ml.transcribe(samples)).strip()
    await _add_artifact(name, "voice note", text)
    return text


async def _load_demo():
    """(JS-callable) Run a small fictional conversation through the pipeline."""
    return str(await _ingest_chat_text("demo-chat.txt", _DEMO_CHAT))


def _dashboard():
    """(JS-callable) Deterministic CRM snapshot as JSON — no LLM involved."""
    stage = _stage()
    return json.dumps(
        {
            "messages": len(_STATE["messages"]),
            "participants": _STATE["participants"],
            "artifacts": _STATE["artifacts"],
            "payments": [
                a["amount"]
                for a in _STATE["artifacts"]
                if a["kind"] == "payment proof" and a.get("amount")
            ],
            "stage": stage,
            "stages": _STAGES,
            "next_action": _NEXT_ACTION[stage],
            "skipped": _STATE["skipped"],
        }
    )


def _facts_sheet():
    stage = _stage()
    lines = [
        f"Pipeline stage: {stage}. Rule-suggested next action: {_NEXT_ACTION[stage]}",
        f"Participants: {', '.join(_STATE['participants']) or '(unknown)'}",
        f"Messages in export: {len(_STATE['messages'])}",
        "Artifacts:",
    ]
    for a in _STATE["artifacts"][:30]:
        amount = f" (amount: {a['amount']})" if a.get("amount") else ""
        lines.append(f"  - {a['kind']}: {a['name']}{amount}")
    if not _STATE["artifacts"]:
        lines.append("  - (none)")
    return "\n".join(lines)


async def _suggest_next_action():
    """(JS-callable) Stage rules + LLM: next move and a ready-to-send reply."""
    if not _STATE["messages"] and not _STATE["artifacts"]:
        return "Drop a WhatsApp chat export first (or load the demo)."
    recent = _STATE["messages"][-CHAT_WINDOW:]
    convo = "\n".join(f"{m['sender']}: {m['text']}"[:200] for m in recent)
    prompt = (
        "You are helping a travel agent decide their next move with this "
        "client.\n"
        f"CLIENT FACTS:\n{_facts_sheet()}\n\n"
        f"RECENT MESSAGES:\n{convo}\n\n"
        "Reply with exactly two sections:\n"
        "## Next action\n(one sentence, consistent with the pipeline stage)\n\n"
        "## Draft message\n(a short, friendly WhatsApp message the agent can "
        "send as-is, in the same language the client writes in)"
    )
    # TEMPLATE_SYSTEM_PROMPT only exists as a JS global; the query bridge falls
    # back to it when custom_prompt is empty (same idiom as the base template).
    return await process_user_query_wllama(
        prompt, globals().get("TEMPLATE_SYSTEM_PROMPT", "")
    )


async def process_user_query(query):
    """Hybrid grounding: the deterministic facts sheet ALWAYS accompanies the
    retrieved excerpts, so the model reads a small, precise context instead of
    the whole conversation.

    Overrides the default router so retrieval is deterministic (small local
    models are unreliable at deciding to call a search tool themselves).
    """
    if not _STATE["messages"] and not _STATE["artifacts"]:
        return (
            "I don't have a conversation yet — drop a WhatsApp chat export "
            "above (or load the demo), then ask me who paid, what's booked, "
            "or what to do next."
        )
    hits = await agentop_rag.search(query, RAG_TOP_K)
    excerpts = "\n\n".join(
        f"[{i + 1}] (from {h['source']}) {h['text']}" for i, h in enumerate(hits)
    ) or "(nothing retrieved)"
    grounded_prompt = (
        "You are a CRM copilot for a small travel agency that sells over "
        "WhatsApp. Answer the QUESTION using ONLY the CLIENT FACTS and "
        "CONVERSATION EXCERPTS below. Quote names, amounts, dates, flight "
        "numbers, and booking codes exactly as written; cite excerpts like "
        "[1]. If the answer is not there, say the export does not show it.\n\n"
        f"CLIENT FACTS:\n{_facts_sheet()}\n\n"
        f"CONVERSATION EXCERPTS:\n{excerpts}\n\nQUESTION: {query}"
    )
    return await process_user_query_wllama(
        grounded_prompt, globals().get("TEMPLATE_SYSTEM_PROMPT", "")
    )