discord_bot.py

# -*- coding: utf-8 -*-
"""
Text2Receipt — Discord Bot  (+5% bonus)
========================================
Run this cell in Google Colab (T4) AFTER the fine-tuned model is loaded.
The bot wraps the same pipeline as the Gradio Space.

Usage in any Discord channel:
  !receipt <Hebrew income note>

Example:
  !receipt קיבלתי 1200 שקל ממשה כהן על ייעוץ עסקי

Setup:
  1. Go to https://discord.com/developers/applications → New Application
  2. Bot tab → Add Bot → copy token
  3. OAuth2 → URL Generator: scope=bot, permissions=Send Messages
  4. Paste DISCORD_TOKEN below (or set as Colab env var)
  5. Run this cell while the model is already loaded in the session
"""

import os, json, random, datetime as _dt, asyncio, threading
import discord
from discord.ext import commands

# ── same pipeline helpers as app.py ──────────────────────────────────────────
import t2r_core as core

DEMO_ISSUER = {
    "name":   "מים שקטים",
    "tax_id": "962569844",
    "status": "authorized_dealer",
}

_RNG = random.Random(99)

# These are defined here so the bot works standalone (copied from app.py)
INSTRUCTION = (
    "אתה ממיר הערת הכנסה חופשית בעברית למבנה JSON. "
    "חלץ אך ורק את מה שכתוב בהערה: שם הלקוח (client_name), "
    "האם הלקוח עסק (client_is_business), ורשימת פריטים (items) "
    "כאשר לכל פריט תיאור (description), מחיר ליחידה (unit_price) וכמות (quantity). "
    "החזר JSON תקין בלבד, ללא טקסט נוסף."
)

def _build_prompt(raw_text: str) -> str:
    return f"{INSTRUCTION}\n\nהערה: {raw_text}\n\nJSON:"

def _extract_json(text: str):
    s = text.find("{")
    if s < 0:
        return None
    depth = 0
    for i in range(s, len(text)):
        if text[i] == "{":
            depth += 1
        elif text[i] == "}":
            depth -= 1
            if depth == 0:
                try:
                    return json.loads(text[s:i+1])
                except Exception:
                    return None
    return None

def _apply_defaults(parse: dict) -> dict:
    p = dict(parse)
    p.setdefault("doc_type",           "receipt")
    p.setdefault("date",               _dt.date.today().isoformat())
    p.setdefault("payment_method",     "bank_transfer")
    p.setdefault("amount_basis",       "net")
    p.setdefault("currency",           "ILS")
    p.setdefault("client_tax_id",      None)
    p.setdefault("client_is_business", False)
    p.setdefault("items",              [])
    return p

def _format_document(completed: dict) -> str:
    """Format completed document as Discord-friendly Markdown."""
    c       = completed
    issuer  = c.get("issuer", {})
    client  = c.get("client", {})
    lines   = c.get("lines",  [])
    vat_pct = int(round(c.get("vat_rate", 0) * 100))

    line_rows = "\n".join(
        f"  • {ln['description']} × {ln['quantity']} → ₪{ln['line_total']:,.2f}"
        for ln in lines
    )
    alloc = (
        f"\n🔖 **מספר הקצאה:** `{c.get('allocation_number')}`"
        if c.get("allocation_required") else ""
    )

    return (
        f"```\n"
        f"{'━'*38}\n"
        f"  {c.get('doc_type_he','מסמך')}   |   מס׳ {c.get('serial_number','')}\n"
        f"  {c.get('issue_date','')}\n"
        f"{'━'*38}\n"
        f"  מנפיק : {issuer.get('name','')}  (ח.פ. {issuer.get('tax_id','')})\n"
        f"  לקוח   : {client.get('name','')}"
        f"{'  [עסק]' if client.get('is_business') else ''}\n"
        f"{'━'*38}\n"
        f"{line_rows}\n"
        f"{'━'*38}\n"
        f"  סכום נטו  : ₪{c.get('subtotal',0):>10,.2f}\n"
        f"  מע\"מ {vat_pct:2d}%  : ₪{c.get('vat_amount',0):>10,.2f}\n"
        f"  **סה\"כ**   : ₪{c.get('total',0):>10,.2f}\n"
        f"  תשלום : {core.PAYMENT_HE.get(c.get('payment_method',''),'—')}\n"
        f"{'━'*38}\n"
        f"```"
        f"{alloc}"
    )


def run_pipeline(raw_text: str, tok, model) -> str:
    """
    Full pipeline: raw Hebrew note → formatted fiscal document string.
    tok and model are the already-loaded objects from the Colab session.
    """
    import torch

    # 1. Parse
    msgs   = [{"role": "user", "content": _build_prompt(raw_text)}]
    prompt = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
    enc    = tok(prompt, return_tensors="pt",
                 add_special_tokens=False).to(model.device)
    with torch.no_grad():
        out = model.generate(**enc, max_new_tokens=200, do_sample=False,
                             pad_token_id=tok.pad_token_id)
    decoded = tok.decode(out[0, enc["input_ids"].shape[1]:],
                         skip_special_tokens=True)
    parse = _extract_json(decoded)

    if parse is None:
        return "❌ לא הצלחתי לחלץ את הפרטים. נסה לנסח מחדש."

    # 2. Apply defaults + complete
    final_parse = _apply_defaults(parse)
    try:
        completed = core.complete(DEMO_ISSUER, final_parse, _RNG)
    except Exception as e:
        return f"❌ שגיאה בעיבוד: {e}"

    # 3. Format
    return _format_document(completed)


def start_bot(tok, model, discord_token: str = None):
    """
    Start the Discord bot. Call this from Colab after loading tok + model.

    Args:
        tok:           HuggingFace tokenizer (already loaded)
        model:         fine-tuned model (already loaded, on GPU)
        discord_token: Bot token. Falls back to DISCORD_TOKEN env var.
    """
    token = discord_token or os.environ.get("DISCORD_TOKEN", "")
    if not token:
        raise ValueError(
            "No Discord token found. "
            "Pass discord_token= or set os.environ['DISCORD_TOKEN']"
        )

    intents = discord.Intents.default()
    intents.message_content = True
    bot = commands.Bot(command_prefix="!", intents=intents)

    @bot.event
    async def on_ready():
        print(f"✅ Discord bot ready: {bot.user} (id: {bot.user.id})")
        print("   Listening for: !receipt <Hebrew note>")

    @bot.command(name="receipt")
    async def receipt_cmd(ctx, *, note: str):
        """Convert a Hebrew income note to a fiscal document."""
        await ctx.message.add_reaction("⏳")
        try:
            # run blocking inference in executor to not block event loop
            loop   = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None, run_pipeline, note, tok, model
            )
            await ctx.reply(result)
        except Exception as e:
            await ctx.reply(f"❌ שגיאה: {e}")
        finally:
            await ctx.message.remove_reaction("⏳", bot.user)
            await ctx.message.add_reaction("✅")

    @bot.command(name="help_t2r")
    async def help_cmd(ctx):
        await ctx.reply(
            "**Text2Receipt Bot** 🧾\n"
            "המר הערת הכנסה עברית למסמך פיסקלי ישראלי.\n\n"
            "שימוש: `!receipt <הערה בעברית>`\n\n"
            "דוגמאות:\n"
            "• `!receipt קיבלתי 500 שקל ממשה על ייעוץ`\n"
            "• `!receipt מעגל בע\"מ שילמה 15,000 ש\"ח על פרויקט אתר`"
        )

    # Run in background thread so Colab cell doesn't block
    def _run():
        asyncio.run(bot.start(token))

    t = threading.Thread(target=_run, daemon=True)
    t.start()
    print("🤖 Bot thread started. Keep this cell running.")
    return bot


# ═════════════════════════════════════════════════════════════════════════════
# Colab usage example (paste into a cell after loading tok + model in nb03):
# ═════════════════════════════════════════════════════════════════════════════
# import os
# os.environ["DISCORD_TOKEN"] = "YOUR_BOT_TOKEN_HERE"
#
# from discord_bot import start_bot
# bot = start_bot(tok, ft_model)   # tok + ft_model from §5 of nb03
#
# # Test locally without Discord:
# from discord_bot import run_pipeline
# print(run_pipeline("קיבלתי 1200 שקל ממשה כהן על ייעוץ עסקי", tok, ft_model))