#!/usr/bin/python3
"""Assemble the monthly utilities message for the rental and push it to ntfy.

Run monthly from rent-utilities.timer. There are two ways this produces a
number for a bill:

  fetch mode     -- [imap].enabled, reads the statement mail through Proton
                    Bridge and regexes the amount out of the body.
  template mode  -- Bridge not configured yet. Carries the last known figures
                    and asks you to supply the real ones.

`--set` overrides either, for the months you would rather just read the two
numbers off your phone and be done.

The one rule this script is built around: EVERY run ends in a push. It is
either a message you can forward as-is, or a specific account of what stopped
it. A monthly job that fails quietly is invisible for thirty days, by which
point you have stopped expecting it -- which is the problem this was written
to solve in the first place.
"""
import argparse
import csv
import datetime as dt
import email
import email.policy
import os
import re
import sys
import tomllib

# The Bridge/IMAP/ntfy layer is shared with inbox_tidy. sys.path rather than
# PYTHONPATH so this still works when run by hand, not only from the unit.
sys.path.insert(0, os.path.expanduser("~/.local/lib/pymail"))
from protonimap import (                                    # noqa: E402
    Fatal, imap_connect, imap_date, mailboxes, notify_failure, push,
    _body_text, _quote, _sent_at)

CONFIG = os.path.expanduser("~/docs/leases/bills.toml")
CSV_PATH = os.path.expanduser("~/docs/leases/utilities.csv")

OK, MISSING, AMBIGUOUS, ERROR = "ok", "missing", "ambiguous", "error"

# --discover sweeps back from today rather than around a month: it exists to
# find senders, not to resolve a particular month's bill.
DISCOVER_DAYS = 120




# --------------------------------------------------------------------------
# small helpers
# --------------------------------------------------------------------------



def prev_month(today=None):
    today = today or dt.date.today()
    first = today.replace(day=1)
    last_prev = first - dt.timedelta(days=1)
    return "%04d-%02d" % (last_prev.year, last_prev.month)


def money(x):
    return "%.2f" % x






def month_window(month, slack_days):
    """(since, before) bracketing the month a bill BELONGS to.

    Anchored on the requested month, never on today. A plain "last N days"
    lookback makes --month a label only: ask for July in September and you get
    August's bills filed under July, with nothing looking wrong. The observed
    arrival days are AT&T 6th-7th, LG&E 10th-11th, water 14th-15th, so a few
    days of slack around the calendar month captures the right bills without
    reaching into the neighbouring month's.
    """
    y, m = int(month[:4]), int(month[5:7])
    first = dt.date(y, m, 1)
    nxt = dt.date(y + (m == 12), (m % 12) + 1, 1)
    return (first - dt.timedelta(days=slack_days),
            nxt + dt.timedelta(days=slack_days))


# --------------------------------------------------------------------------
# ntfy
# --------------------------------------------------------------------------





# --------------------------------------------------------------------------
# CSV
# --------------------------------------------------------------------------

def load_csv():
    """-> {month: {utility: amount}}. Long format, one row per reading."""
    out = {}
    try:
        with open(CSV_PATH, newline="") as fh:
            for row in csv.DictReader(fh):
                out.setdefault(row["month"], {})[row["utility"]] = float(row["amount"])
    except FileNotFoundError:
        raise Fatal(
            "utilities.csv is missing.\n\n"
            "Expected it at %s. If this is a fresh checkout, restore it from "
            "utilities.csv.pre-automation in the same directory." % CSV_PATH)
    except (KeyError, ValueError) as exc:
        raise Fatal(
            "utilities.csv could not be parsed: %s\n\n"
            "It should be long format with the header `month,utility,amount` "
            "and one row per reading. The pre-conversion copy is beside it as "
            "utilities.csv.pre-automation." % exc)
    return out


def write_csv(table):
    rows = sorted(
        ((m, u, money(a)) for m, us in table.items() for u, a in us.items()),
        key=lambda t: (t[0], t[1]))
    tmp = CSV_PATH + ".tmp"
    with open(tmp, "w", newline="") as fh:
        # lineterminator is explicit: csv.writer defaults to CRLF, which would
        # quietly convert a file that has always been LF.
        w = csv.writer(fh, lineterminator="\n")
        w.writerow(["month", "utility", "amount"])
        w.writerows(rows)
    os.replace(tmp, CSV_PATH)


# --------------------------------------------------------------------------
# IMAP / Proton Bridge
# --------------------------------------------------------------------------











def message_bodies(conn, imap_cfg, sender, since, before):
    """Most-recent-first list of (subject, body-text) from `sender`.

    Ordered by the Date header across every mailbox searched, so "most
    recent" still means most recent when a bill was filed out of INBOX.
    """
    found = []
    for box in mailboxes(imap_cfg):
        typ, _ = conn.select(_quote(box), readonly=True)
        if typ != "OK":
            continue          # folder renamed or gone; the others still count
        typ, data = conn.search(
            None, '(SINCE "%s" BEFORE "%s" FROM "%s")'
            % (imap_date(since), imap_date(before), sender))
        if typ != "OK":
            continue
        for num in data[0].split():
            typ, raw = conn.fetch(num, "(RFC822)")
            if typ != "OK" or not raw or not raw[0]:
                continue
            msg = email.message_from_bytes(raw[0][1], policy=email.policy.default)
            found.append((_sent_at(msg),
                          str(msg.get("Subject", "(no subject)")),
                          _body_text(msg)))
    found.sort(key=lambda t: t[0], reverse=True)
    return [(subj, body) for _, subj, body in found]


def mailbox_total(conn, imap_cfg, since, before):
    """How many messages exist in the window at all, from anyone.

    Distinguishes "the bill did not arrive" from "Bridge has not finished
    syncing yet". A first sync takes ~90 minutes, and a partially-synced
    mailbox answers searches successfully with incomplete results -- so a
    missing bill and an unsynced mailbox look identical unless this is
    checked. Reporting the wrong one would send you hunting for a sender that
    was right all along.
    """
    total = 0
    try:
        for box in mailboxes(imap_cfg):
            typ, _ = conn.select(_quote(box), readonly=True)
            if typ != "OK":
                continue
            typ, data = conn.search(None, '(SINCE "%s" BEFORE "%s")'
                                    % (imap_date(since), imap_date(before)))
            if typ == "OK":
                total += len(data[0].split())
        return total
    except Exception:
        return None


def extract_amount(bodies, pattern):
    """-> (status, value, detail)."""
    rx = re.compile(pattern)
    # Every matching message in the window is considered, not just the newest.
    # Taking the newest would quietly resolve a window that had slipped and
    # caught two months of bills -- the failure this must not have. Within a
    # correctly scoped month there is exactly one bill, so two DIFFERENT
    # figures means something is wrong and is worth refusing. The same figure
    # repeated is fine: LG&E sends the identical mail twice, and AT&T's
    # "Bill total" appears once per bill.
    found, where = {}, {}
    for subject, body in bodies:
        for m in rx.finditer(body):
            v = m.group(1).replace(",", "")
            found[v] = found.get(v, 0) + 1
            where.setdefault(v, subject)
    if len(found) == 1:
        v = next(iter(found))
        return OK, float(v), where[v]
    if len(found) > 1:
        vals = ", ".join("$%s (%r)" % (v, where[v][:44]) for v in sorted(found))
        return AMBIGUOUS, None, (
            "matched %d DIFFERENT amounts in the window -- %s. Refusing to "
            "guess which is this month's." % (len(found), vals))
    if bodies:
        return MISSING, None, (
            "%d message(s) from that sender, but the pattern matched no dollar "
            "amount in any of them. Most recent was %r"
            % (len(bodies), bodies[0][0]))
    return MISSING, None, "no message from that sender in the search window"


# --------------------------------------------------------------------------
# rendering
# --------------------------------------------------------------------------

def render_message(cfg, amounts):
    """The text to forward to them. Layout matches the original utilities.py.

    One deliberate difference: the original's Total line printed `$$` because
    total_amount already carried a '$' and the f-string added another.
    """
    split = cfg["split"]
    lines = []
    full = their = 0.0
    for bill in cfg["bills"]:
        name = bill["name"]
        amt = amounts[name]
        share = amt * split
        full += amt
        their += share
        lines.append("%-15s $%-8s Your Share = $%s"
                     % (name + ":", money(amt) + ",", money(share)))
    lines.append("%-15s $%-8s Your Total Share = $%s"
                 % ("Total:", money(full) + ",", money(their)))
    body = "\n".join(lines)
    greeting = cfg.get("greeting", "").strip()
    return (greeting + "\n\n" + body) if greeting else body


def render_problems(month, problems, note=None):
    out = ["**Could not build the %s message.**" % month, ""]
    if note:
        out += [note, ""]
    for name, detail in problems:
        out.append("**%s**" % name)
        out.append(detail)
        out.append("")
    out.append("---")
    out.append("Nothing was written to utilities.csv. Once you have the real "
               "figures, run:")
    out.append("")
    out.append("```")
    # Join the continuations explicitly rather than appending a trailing "\\"
    # to every line -- the last one must NOT have it, or pasting the command
    # leaves the shell waiting on a continuation that never comes.
    cmd = ["rent_utilities --month %s" % month]
    cmd += ["    --set '%s=0.00'" % name for name, _ in problems]
    out.append(" \\\n".join(cmd))
    out.append("```")
    return "\n".join(out)


# --------------------------------------------------------------------------
# modes
# --------------------------------------------------------------------------

def discover(cfg):
    imap_cfg = cfg["imap"]
    conn = imap_connect(imap_cfg)
    try:
        conn.select(_quote(mailboxes(imap_cfg)[0]), readonly=True)
        since = dt.date.today() - dt.timedelta(days=DISCOVER_DAYS)
        typ, data = conn.search(None, '(SINCE "%s")' % imap_date(since))
        if typ != "OK":
            raise Fatal("IMAP search failed: %s" % typ)
        ids = data[0].split()
        print("%d messages since %s\n" % (len(ids), since))
        rx = re.compile(r"\$\s*([\d,]+\.\d{2})")
        for num in reversed(ids):
            typ, raw = conn.fetch(num, "(RFC822)")
            if typ != "OK" or not raw or not raw[0]:
                continue
            msg = email.message_from_bytes(raw[0][1], policy=email.policy.default)
            body = ""
            try:
                part = msg.get_body(preferencelist=("plain", "html"))
                body = part.get_content() if part is not None else ""
            except Exception:
                pass
            body = re.sub(r"<[^>]+>", " ", body)
            amounts = sorted({m.group(1) for m in rx.finditer(body)})
            print("from:    %s" % msg.get("From", "?"))
            print("subject: %s" % msg.get("Subject", "?"))
            print("date:    %s" % msg.get("Date", "?"))
            print("amounts: %s" % (", ".join("$" + a for a in amounts[:8]) or "none"))
            print()
    finally:
        try:
            conn.logout()
        except Exception:
            pass
    return 0


def collect(cfg, month, overrides):
    """-> (amounts, problems, used_fetch)."""
    amounts, problems = {}, []
    imap_cfg = cfg.get("imap", {})
    use_imap = imap_cfg.get("enabled", False)
    conn = None
    since, before = month_window(month, imap_cfg.get("window_slack_days", 5))

    try:
        for bill in cfg["bills"]:
            name = bill["name"]

            if name in overrides:
                amounts[name] = overrides[name]
                continue

            if "fixed" in bill:
                amounts[name] = float(bill["fixed"])
                continue

            if not use_imap:
                problems.append((name, "no amount supplied, and [imap].enabled "
                                       "is false so nothing was fetched."))
                continue

            sender, pattern = bill.get("sender", ""), bill.get("pattern", "")
            if not sender or not pattern:
                problems.append((name, (
                    "not configured for fetching: `sender` and/or `pattern` "
                    "are empty in bills.toml. Run `rent_utilities --discover` "
                    "to see the senders and candidate amounts in your recent "
                    "mail, then fill them in.")))
                continue

            if conn is None:
                conn = imap_connect(imap_cfg)
            try:
                bodies = message_bodies(conn, imap_cfg, sender, since, before)
            except Fatal:
                raise
            except Exception as exc:
                problems.append((name, "IMAP read failed: %r" % (exc,)))
                continue

            status, value, detail = extract_amount(bodies, pattern)
            if status == OK:
                amounts[name] = value
                continue

            note = ""
            if not bodies:
                total = mailbox_total(conn, imap_cfg, since, before)
                if total == 0:
                    note = ("\n\n**The mailbox returned NO messages at all in "
                            "this window, from anyone.** That is almost "
                            "certainly Bridge still syncing rather than a "
                            "missing bill -- a first sync takes ~90 minutes. "
                            "Check with:\n"
                            "  sudo podman logs --tail 20 "
                            "connor_protonmail-bridge_1\n"
                            "Do not go changing the sender in bills.toml yet.")
                elif total is not None:
                    note = ("\n\n(%d message(s) from other senders were found "
                            "in the same window, so the mailbox is reachable "
                            "and this sender genuinely has nothing.)" % total)
            problems.append((name, "%s\n\nSearched for `%s` between %s and %s in %s.%s"
                             % (detail, sender, since, before,
                                ", ".join(mailboxes(imap_cfg)), note)))
    finally:
        if conn is not None:
            try:
                conn.logout()
            except Exception:
                pass

    return amounts, problems, use_imap


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--month", help="target month as YYYY-MM (default: last month)")
    ap.add_argument("--set", action="append", default=[], metavar="NAME=AMOUNT",
                    help="supply an amount by hand; repeatable")
    ap.add_argument("--discover", action="store_true",
                    help="print recent senders and candidate amounts, then exit")
    ap.add_argument("--dry-run", action="store_true",
                    help="print what would be pushed instead of pushing")
    args = ap.parse_args()

    # Config load is outside the push-on-failure net: without it there is no
    # ntfy url or topic to push to. It surfaces on stderr, and systemd will
    # mark the unit failed.
    try:
        with open(CONFIG, "rb") as fh:
            cfg = tomllib.load(fh)
    except (OSError, tomllib.TOMLDecodeError) as exc:
        print("rent_utilities: cannot read config %s: %s" % (CONFIG, exc),
              file=sys.stderr)
        return 2

    if args.discover:
        try:
            return discover(cfg)
        except Fatal as exc:
            print("rent_utilities: %s" % exc, file=sys.stderr)
            return 1

    month = args.month or prev_month()
    if not re.fullmatch(r"\d{4}-\d{2}", month):
        print("rent_utilities: --month must be YYYY-MM, got %r" % month,
              file=sys.stderr)
        return 2

    overrides = {}
    for item in args.set:
        if "=" not in item:
            print("rent_utilities: --set wants NAME=AMOUNT, got %r" % item,
                  file=sys.stderr)
            return 2
        k, v = item.split("=", 1)
        try:
            overrides[k.strip()] = float(v)
        except ValueError:
            print("rent_utilities: %r is not a number in %r" % (v, item),
                  file=sys.stderr)
            return 2

    known = {b["name"] for b in cfg["bills"]}
    unknown = set(overrides) - known
    if unknown:
        print("rent_utilities: --set names not in bills.toml: %s (known: %s)"
              % (", ".join(sorted(unknown)), ", ".join(sorted(known))),
              file=sys.stderr)
        return 2

    try:
        table = load_csv()
        amounts, problems, used_fetch = collect(cfg, month, overrides)
    except Fatal as exc:
        body = "**%s could not run.**\n\n%s" % (month, exc)
        if args.dry_run:
            print(body)
        else:
            notify_failure(cfg, "Rent utilities: FAILED", body)
        return 1

    if problems:
        note = None
        if not used_fetch:
            last = max((m for m in table if m < month), default=None)
            if last:
                prev = ", ".join(
                    "%s $%s" % (u, money(a)) for u, a in sorted(table[last].items()))
                note = ("Template mode -- Proton Bridge is not wired up yet, so "
                        "nothing was fetched.\n\nLast figures on file (%s): %s"
                        % (last, prev))
        body = render_problems(month, problems, note)
        title = ("Rent utilities: %s needs numbers" % month if not used_fetch
                 else "Rent utilities: %s incomplete" % month)
        if args.dry_run:
            print(body)
        else:
            notify_failure(cfg, title, body)
        return 1

    msg = render_message(cfg, amounts)
    total = sum(amounts.values()) * cfg["split"]
    body = ("Their share for %s: **$%s**\n\nForward this:\n\n```\n%s\n```"
            % (month, money(total), msg))

    # Before the write, not after: a dry run must not touch utilities.csv.
    if args.dry_run:
        print(body)
        return 0

    table.setdefault(month, {}).update(amounts)
    try:
        write_csv(table)
    except OSError as exc:
        body = ("**Built the %s figures but could not save them.**\n\n"
                "Writing %s failed: %s\n\nThe message below is still correct, "
                "but it was NOT recorded -- next month will not have this "
                "month's history.\n\n```\n%s\n```"
                % (month, CSV_PATH, exc, render_message(cfg, amounts)))
        # No dry-run branch here: a dry run has already returned above, so
        # reaching this point means the write was real and really failed.
        notify_failure(cfg, "Rent utilities: %s not saved" % month, body)
        return 1

    try:
        push(cfg, "Rent utilities: %s" % month, body, "house,moneybag")
    except (urllib.error.URLError, OSError) as exc:
        print("rent_utilities: built the message but ntfy push failed: %s"
              % exc, file=sys.stderr)
        print(msg)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())
