#!/usr/bin/env python3
"""ntfy -> Matrix bridge.

Watches a set of ntfy topics over the streaming JSON API and forwards every
notification to a Matrix room. Designed to run as a long-lived daemon:
  * one streaming connection per mapping, with exponential-backoff reconnect
  * per-mapping resume using ntfy's `since=<message_id>` so nothing is lost
    across reconnects or process restarts
  * a single Matrix sender with retry/backoff; the resume cursor only advances
    after a message is actually accepted by the homeserver (at-least-once)

This handles UNENCRYPTED Matrix rooms only. The target room must not be
end-to-end encrypted, or messages will be sent in the clear and ignored by
clients.
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import html
import json
import logging
import os
import re
import signal
import sys
from collections import deque
from typing import Any, Optional

import aiohttp
import yaml
from nio import AsyncClient, AsyncClientConfig, RoomSendResponse

log = logging.getLogger("ntfy-matrix")

# --------------------------------------------------------------------------- #
# Presentation helpers
# --------------------------------------------------------------------------- #

# A small subset of ntfy's emoji short-code -> emoji map. Unknown tags are
# shown as text labels instead. (ntfy itself does something similar.)
EMOJI_MAP = {
    "warning": "⚠️", "rotating_light": "🚨", "fire": "🔥", "skull": "💀",
    "bug": "🐛", "tada": "🎉", "partying_face": "🥳", "white_check_mark": "✅",
    "heavy_check_mark": "✔️", "x": "❌", "no_entry": "⛔", "bell": "🔔",
    "loudspeaker": "📢", "package": "📦", "computer": "💻", "floppy_disk": "💾",
    "zzz": "💤", "robot": "🤖", "lock": "🔒", "unlock": "🔓",
    "green_circle": "🟢", "yellow_circle": "🟡", "red_circle": "🔴",
    "+1": "👍", "-1": "👎",
}

# Prepended for elevated priorities (4 = high, 5 = max/urgent).
PRIORITY_EMOJI = {4: "🟠", 5: "🔴"}


def _split_tags(tags: list[str]) -> tuple[list[str], list[str]]:
    emojis, texts = [], []
    for t in tags:
        e = EMOJI_MAP.get(t)
        (emojis.append(e) if e else texts.append(t))
    return emojis, texts


def build_content(mapping: dict, msg: dict) -> dict:
    """Turn an ntfy message object into a Matrix m.room.message content dict."""
    title = msg.get("title")
    body = msg.get("message", "")
    priority = msg.get("priority", 3)
    tags = msg.get("tags") or []
    click = msg.get("click")
    attachment = msg.get("attachment") or {}

    emoji_tags, text_tags = _split_tags(tags)
    lead = "".join(emoji_tags)
    if priority in PRIORITY_EMOJI:
        lead = PRIORITY_EMOJI[priority] + lead

    # ---- plain text body ----
    lines: list[str] = []
    header = " ".join(p for p in (lead, title) if p).strip()
    if header:
        lines.append(header)
    if body:
        lines.append(body)
    if text_tags:
        lines.append("Tags: " + ", ".join(text_tags))
    if click:
        lines.append("Link: " + click)
    if attachment.get("url"):
        lines.append("Attachment: " + attachment["url"])
    plain = "\n".join(lines) if lines else "(empty notification)"

    # ---- HTML formatted body ----
    esc = html.escape
    h: list[str] = []
    head_html = " ".join(
        p for p in (lead, f"<strong>{esc(title)}</strong>" if title else "") if p
    ).strip()
    if head_html:
        h.append(head_html)
    if body:
        h.append(esc(body).replace("\n", "<br>"))
    if text_tags:
        h.append(f"<em>Tags: {esc(', '.join(text_tags))}</em>")
    if click:
        h.append(f'<a href="{esc(click)}">{esc(click)}</a>')
    if attachment.get("url"):
        a_url = attachment["url"]
        a_name = attachment.get("name", a_url)
        h.append(f'📎 <a href="{esc(a_url)}">{esc(a_name)}</a>')
    formatted = "<br>".join(h)

    return {
        "msgtype": mapping.get("msgtype", "m.text"),
        "body": plain,
        "format": "org.matrix.custom.html",
        "formatted_body": formatted,
    }


# --------------------------------------------------------------------------- #
# Config + state
# --------------------------------------------------------------------------- #

_ENV_RE = re.compile(r"\$\{([^}]+)\}")


def _expand_env(obj: Any) -> Any:
    """Recursively replace ${VAR} in string values with environment variables.

    Raises if a referenced variable is unset or empty, so a missing secret
    fails loudly at startup instead of silently degrading (e.g. an empty token
    turning into an anonymous request that 403s later).
    """
    if isinstance(obj, str):
        def repl(m: "re.Match") -> str:
            var = m.group(1)
            val = os.environ.get(var)
            if not val:
                raise ValueError(
                    f"config references ${{{var}}} but that environment "
                    f"variable is unset or empty"
                )
            return val
        return _ENV_RE.sub(repl, obj)
    if isinstance(obj, list):
        return [_expand_env(x) for x in obj]
    if isinstance(obj, dict):
        return {k: _expand_env(v) for k, v in obj.items()}
    return obj


def load_config(path: str) -> dict:
    with open(path) as f:
        cfg = _expand_env(yaml.safe_load(f))

    if "matrix" not in cfg or "mappings" not in cfg:
        raise ValueError("config must contain 'matrix' and 'mappings' sections")

    m = cfg["matrix"]
    if not m.get("access_token"):
        m["access_token"] = os.environ.get("MATRIX_ACCESS_TOKEN", "")
    if not m.get("access_token"):
        raise ValueError(
            "no Matrix access token (set matrix.access_token or "
            "the MATRIX_ACCESS_TOKEN env var)"
        )
    for key in ("homeserver", "user_id"):
        if not m.get(key):
            raise ValueError(f"matrix.{key} is required")

    names = set()
    for mp in cfg["mappings"]:
        for key in ("name", "ntfy_server", "topics", "room"):
            if not mp.get(key):
                raise ValueError(f"each mapping needs '{key}' (offending: {mp})")
        if mp["name"] in names:
            raise ValueError(f"duplicate mapping name: {mp['name']}")
        names.add(mp["name"])
    return cfg


class State:
    """Tiny JSON-backed store of the last forwarded message id per mapping."""

    def __init__(self, path: str):
        self.path = path
        self.data: dict[str, dict] = {}
        self._lock = asyncio.Lock()

    def load(self) -> None:
        if os.path.exists(self.path):
            with open(self.path) as f:
                self.data = json.load(f)

    def last_id(self, name: str) -> Optional[str]:
        return self.data.get(name, {}).get("last_id")

    def set_last_id(self, name: str, mid: str) -> None:
        self.data.setdefault(name, {})["last_id"] = mid

    async def save(self) -> None:
        async with self._lock:
            tmp = self.path + ".tmp"
            with open(tmp, "w") as f:
                json.dump(self.data, f, indent=2)
            os.replace(tmp, self.path)  # atomic replace


# --------------------------------------------------------------------------- #
# ntfy side: one streaming subscription per mapping
# --------------------------------------------------------------------------- #

def _auth_headers(mapping: dict) -> dict:
    if mapping.get("token"):
        return {"Authorization": f"Bearer {mapping['token']}"}
    if mapping.get("username") is not None:
        raw = f"{mapping['username']}:{mapping.get('password', '')}".encode()
        return {"Authorization": "Basic " + base64.b64encode(raw).decode()}
    return {}


async def subscribe_loop(
    mapping: dict,
    queue: "asyncio.Queue",
    state: State,
    stop: asyncio.Event,
) -> None:
    name = mapping["name"]
    base = mapping["ntfy_server"].rstrip("/")
    topics = ",".join(mapping["topics"])
    url = f"{base}/{topics}/json"
    headers = _auth_headers(mapping)
    prio_min = mapping.get("priority_min")
    auth_mode = (
        "bearer" if mapping.get("token")
        else "basic" if mapping.get("username") is not None
        else "none"
    )
    log.info("[%s] target %s (auth: %s)", name, url, auth_mode)

    # Bounded dedup of recently seen ids, seeded with the persisted boundary so
    # a replayed `since=<id>` message isn't forwarded twice after a restart.
    seen: deque[str] = deque(maxlen=2000)
    seen_set: set[str] = set()
    boundary = state.last_id(name)
    if boundary:
        seen.append(boundary)
        seen_set.add(boundary)

    def remember(mid: str) -> None:
        if len(seen) == seen.maxlen:
            seen_set.discard(seen[0])
        seen.append(mid)
        seen_set.add(mid)

    timeout = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=120)
    backoff = 1

    async with aiohttp.ClientSession(timeout=timeout) as session:
        while not stop.is_set():
            params = {}
            last = state.last_id(name)
            if last:
                params["since"] = last
            try:
                async with session.get(url, headers=headers, params=params) as resp:
                    if resp.status == 401 or resp.status == 403:
                        log.error("[%s] auth failed (HTTP %s); check token/credentials",
                                  name, resp.status)
                        await asyncio.sleep(min(backoff, 60))
                        backoff = min(backoff * 2, 60)
                        continue
                    resp.raise_for_status()
                    log.info("[%s] subscribed to %s (topics: %s)", name, base, topics)
                    backoff = 1  # connected; reset backoff for the next drop

                    async for raw in resp.content:
                        if stop.is_set():
                            break
                        line = raw.strip()
                        if not line:
                            continue
                        try:
                            obj = json.loads(line)
                        except json.JSONDecodeError:
                            log.warning("[%s] non-JSON line skipped: %r", name, line[:200])
                            continue

                        if obj.get("event") != "message":
                            continue  # open / keepalive / poll_request
                        mid = obj.get("id")
                        if not mid or mid in seen_set:
                            continue
                        if prio_min and obj.get("priority", 3) < prio_min:
                            remember(mid)
                            continue

                        remember(mid)
                        await queue.put((mapping, obj))
            except asyncio.CancelledError:
                raise
            except Exception as e:  # noqa: BLE001 - daemon must keep running
                log.warning("[%s] stream error: %s; reconnecting in %ds",
                            name, e, backoff)
                try:
                    await asyncio.wait_for(stop.wait(), timeout=backoff)
                except asyncio.TimeoutError:
                    pass
                backoff = min(backoff * 2, 60)


# --------------------------------------------------------------------------- #
# Matrix side: a single sender drains the queue with retry/backoff
# --------------------------------------------------------------------------- #

async def sender_loop(
    client: AsyncClient,
    queue: "asyncio.Queue",
    state: State,
    stop: asyncio.Event,
) -> None:
    while not stop.is_set():
        mapping, msg = await queue.get()
        try:
            delay = 1
            while not stop.is_set():
                try:
                    resp = await client.room_send(
                        room_id=mapping["room"],
                        message_type="m.room.message",
                        content=build_content(mapping, msg),
                    )
                    if isinstance(resp, RoomSendResponse):
                        state.set_last_id(mapping["name"], msg["id"])
                        await state.save()
                        log.info("[%s] forwarded %s -> %s",
                                 mapping["name"], msg["id"], mapping["room"])
                        break
                    log.warning("[%s] send rejected: %s; retrying in %ds",
                                mapping["name"], resp, delay)
                except asyncio.CancelledError:
                    raise
                except Exception as e:  # noqa: BLE001
                    log.warning("[%s] send error: %s; retrying in %ds",
                                mapping["name"], e, delay)
                try:
                    await asyncio.wait_for(stop.wait(), timeout=delay)
                except asyncio.TimeoutError:
                    pass
                delay = min(delay * 2, 60)
        finally:
            queue.task_done()


# --------------------------------------------------------------------------- #
# Wiring
# --------------------------------------------------------------------------- #

async def run(cfg: dict) -> None:
    state = State(cfg.get("state_file", "./state.json"))
    state.load()

    mc = cfg["matrix"]
    client = AsyncClient(
        mc["homeserver"],
        mc["user_id"],
        config=AsyncClientConfig(request_timeout=30),
    )
    client.access_token = mc["access_token"]
    client.user_id = mc["user_id"]
    if mc.get("device_id"):
        client.device_id = mc["device_id"]

    queue: asyncio.Queue = asyncio.Queue()
    stop = asyncio.Event()

    loop = asyncio.get_running_loop()
    for sig in (signal.SIGINT, signal.SIGTERM):
        try:
            loop.add_signal_handler(sig, stop.set)
        except NotImplementedError:
            pass  # e.g. on Windows

    tasks = [asyncio.create_task(sender_loop(client, queue, state, stop))]
    for mp in cfg["mappings"]:
        tasks.append(asyncio.create_task(subscribe_loop(mp, queue, state, stop)))

    log.info("bridge started with %d mapping(s)", len(cfg["mappings"]))
    await stop.wait()
    log.info("shutting down...")

    for t in tasks:
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)
    await state.save()
    await client.close()
    log.info("stopped cleanly")


def main() -> None:
    ap = argparse.ArgumentParser(description="Forward ntfy notifications to Matrix.")
    ap.add_argument("-c", "--config", default="config.yaml", help="path to config YAML")
    args = ap.parse_args()

    logging.basicConfig(
        level=os.environ.get("LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)-7s %(message)s",
    )

    try:
        cfg = load_config(args.config)
    except (OSError, ValueError) as e:
        log.error("config error: %s", e)
        sys.exit(1)

    try:
        asyncio.run(run(cfg))
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()