#!/usr/bin/env python3
"""An example consumer of the feed's AMQP delivery, with pika.

It follows the client guide step by step: it connects with the pinned certificate
authority, binds its price queue, joins (heartbeat, join answer, held messages),
keeps every row and event by `version`, checks every heartbeat, stores every result
before acknowledging it, and rejoins whenever a rule says so. While it is not
joined, every price is suspended.

Configuration comes from environment variables (see README.md). What an application
would receive is printed on stdout, one JSON object per line; logs go to stderr and
never contain the password or the API key.
"""

import gzip
import json
import logging
import os
import random
import re
import signal
import ssl
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

import pika
import pika.exceptions

# The feed's client parameters, in seconds.
CLIENT_SILENCE_LIMIT = 3.0  # rule 1: no heartbeat for this long
CLIENT_LAG_LIMIT = 3.0  # rule 3: a heartbeat received this long after its sent_at
JOIN_TIMEOUT = 10.0  # step 3: waiting for a usable heartbeat
JOIN_REQUEST_TIMEOUT = 30.0  # step 4: one join request
JOIN_GIVE_UP = 60.0  # step 4: no answer within this, back to step 3
AMQP_CONNECTION_HEARTBEAT = 10  # never 0
RESULTS_PREFETCH = 1000  # at most 1,000

log = logging.getLogger("feed-consumer")


class Config:
    """Settings from the environment."""

    def __init__(self, env):
        url = urllib.parse.urlsplit(env["FEED_AMQP_URL"])
        if url.scheme != "amqps":
            raise SystemExit("FEED_AMQP_URL must be amqps://<broker address>:5671/feed")
        if url.username or url.password:
            raise SystemExit("put the user and password in FEED_AMQP_USER and FEED_AMQP_PASSWORD")
        self.host = url.hostname
        self.port = url.port or 5671
        self.vhost = urllib.parse.unquote(url.path[1:]) if len(url.path) > 1 else "feed"
        self.user = env["FEED_AMQP_USER"]
        self.password = env["FEED_AMQP_PASSWORD"]
        self.ca_file = env["FEED_AMQP_CA_FILE"]
        client = env["FEED_CLIENT"]
        queue_set = int(env.get("FEED_QUEUE_SET", "1"))
        self.prices_queue = f"c_{client}.prices.{queue_set}"
        self.results_queue = f"c_{client}.results.{queue_set}"
        self.api_url = env["FEED_API_URL"].rstrip("/")
        self.api_key = env["FEED_API_KEY"]
        self.api_ca_file = env.get("FEED_API_CA_FILE") or None
        self.bindings = split(env.get("FEED_BINDINGS", "#"))
        self.unbind = [k for k in split(env.get("FEED_UNBIND", "")) if k not in self.bindings]
        self.events = env.get("FEED_EVENTS", "0") == "1"
        self.results_file = env.get("FEED_RESULTS_FILE", "results.jsonl")
        self.filters = filters_of(self.bindings)


def split(text):
    return [k.strip() for k in text.split(",") if k.strip()]


def filters_of(bindings):
    """The join filters for a set of bindings: at most four, one per key shape."""
    if "#" in bindings:
        return [{}]
    sports, competitions, events = [], [], []
    for key in bindings:
        words = key.split(".")
        if len(words) == 2 and words[1] == "#" and words[0] not in ("*", "#"):
            sports.append(words[0])
        elif len(words) == 3 and words[0] == "*" and words[2] == "*" and words[1] not in ("*", "#"):
            competitions.append(words[1])
        elif len(words) == 3 and words[:2] == ["*", "*"] and words[2] not in ("*", "#"):
            events.append(words[2])
        else:
            raise SystemExit(f"binding {key!r} is not #, <sport>.#, *.<competition_id>.* or *.*.<event_id>")
    pairs = (("sports", sports), ("competitions", competitions), ("events", events))
    return [{name: ids} for name, ids in pairs if ids]


def topic_match(pattern, key):
    """AMQP topic matching: `*` is one word, `#` zero or more."""
    p, k = pattern.split("."), key.split(".")

    def match(i, j):
        if i == len(p):
            return j == len(k)
        if p[i] == "#":
            return any(match(i + 1, n) for n in range(j, len(k) + 1))
        return j < len(k) and p[i] in ("*", k[j]) and match(i + 1, j + 1)

    return match(0, 0)


def row_key(row):
    event = row["event"]
    return f"{event['sport']}.{event['competition_id']}.{event['id']}"


def event_key(event):
    return f"{event['sport']}.{event['competition']['id']}.{event['id']}"


def utc_seconds(text):
    """An RFC 3339 time with `Z` as seconds since the epoch."""
    return datetime.fromisoformat(text.replace("Z", "+00:00")).timestamp()


def backoff(attempt):
    """Exponential from 1 s to 30 s, with jitter: [1, 2], [2, 4], ... [15, 30]."""
    step = min(30.0, 2.0 ** min(attempt, 6))
    return max(1.0, random.uniform(step / 2, step))


class RetryAfter(Exception):
    def __init__(self, seconds):
        super().__init__(f"retry after {seconds} s")
        self.seconds = seconds


class ResultJournalFailure(RuntimeError):
    """Stop rather than append after a possibly partial journal write."""


class ResultStore:
    """Every result received, by market, kept in a JSON lines file. A result is
    acknowledged only once it is written and synced."""

    def __init__(self, path):
        self.path = path
        self.versions = {}
        if os.path.exists(path):
            with open(path, encoding="utf-8") as f:
                for line in f:
                    if line.strip():
                        row = json.loads(line)
                        validate_result(row)
                        self.versions[row["market_id"]] = max(row["version"], self.versions.get(row["market_id"], 0))
        self.file = open(path, "a", encoding="utf-8")
        # Persist the directory entry too, before any delivery can be acknowledged.
        parent = os.open(os.path.dirname(os.path.abspath(path)), os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(parent)
        finally:
            os.close(parent)

    def store(self, row):
        """Stores a result if it is newer than the one held; True if it was."""
        validate_result(row)
        if row["version"] <= self.versions.get(row["market_id"], 0):
            return False
        encoded = json.dumps(row, separators=(",", ":")) + "\n"
        try:
            self.file.write(encoded)
            self.file.flush()
            os.fsync(self.file.fileno())
        except OSError as error:
            raise ResultJournalFailure("result journal write failed; reconcile before restart") from error
        self.versions[row["market_id"]] = row["version"]
        return True


def validate_result(row):
    """Known result messages must be valid before deduplication or acknowledgement."""
    result = row.get("result") if isinstance(row, dict) else None
    if (not isinstance(row, dict) or row.get("type") != "result"
            or not isinstance(row.get("market_id"), str) or not row["market_id"].strip()
            or type(row.get("version")) is not int or not 0 < row["version"] < 2**53
            or not isinstance(result, dict)
            or type(result.get("revision")) is not int or result["revision"] <= 0
            or type(result.get("void")) is not bool or "winner" not in result
            or (result["winner"] is not None and not isinstance(result["winner"], str))
            or not isinstance(result.get("resolved_at"), str)
            or not re.fullmatch(r"\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(?:\.\d+)?(?:Z|[+-]\d{2}:\d{2})", result["resolved_at"])):
        raise ValueError("invalid result message")
    utc_seconds(result["resolved_at"])


class Printer:
    """What the application receives, as JSON lines on stdout."""

    def out(self, **fields):
        fields["t"] = round(time.time(), 3)
        sys.stdout.write(json.dumps(fields, separators=(",", ":")) + "\n")
        sys.stdout.flush()

    def joined(self, consumer):
        for row in consumer.rows.values():
            self.row(row, via="join")
        for event in consumer.events.values():
            self.event(event, via="join")
        self.out(kind="status", joined=True, session=consumer.session, join_version=consumer.join_version,
                 rows=len(consumer.rows), events=len(consumer.events))

    def suspended(self, reason):
        self.out(kind="status", joined=False, reason=reason)

    def row(self, row, via="update"):
        self.out(kind="row", via=via, market_id=row["market_id"], event_id=row["event"]["id"],
                 version=row["version"], state=row["state"], status=row["market"]["status"])

    def event(self, event, via="update"):
        self.out(kind="event", via=via, event_id=event["id"], version=event["version"], status=event["status"])

    def dropped(self, event_id, market_ids):
        self.out(kind="drop", event_id=event_id, market_ids=sorted(market_ids))

    def result(self, row, queue):
        result = row.get("result") or {}
        self.out(kind="result", queue=queue, market_id=row["market_id"], version=row["version"],
                 revision=result.get("revision"), winner=result.get("winner"), void=result.get("void"))

    def heartbeat(self, hb):
        self.out(kind="heartbeat", session=hb["session"], seq=hb["seq"], as_of=hb["as_of"])


class FeedConsumer:
    """The join state machine: connect (step 1), wait (step 3), join (step 4), live
    (step 5). Every price is suspended outside live."""

    def __init__(self, cfg, app):
        self.cfg = cfg
        self.app = app
        self.results = ResultStore(cfg.results_file)
        self.api_ssl = ssl.create_default_context(cafile=cfg.api_ca_file) if cfg.api_ca_file else None
        self.state = "connect"
        self.joined = False
        self.rows = {}  # market_id -> the newest row
        self.events = {}  # event_id -> the newest event (FEED_EVENTS=1)
        self.removed_events = {}  # event_id -> last removal version, until the next join
        self.join_version = 0
        self.session = None
        self.seq = 0
        self.held = []  # messages received while joining, in order
        self.last_heartbeat = 0.0
        self.wait_since = 0.0
        self.join_future = None
        self.join_started = 0.0
        self.join_generation = 0
        self.joins = ThreadPoolExecutor(max_workers=4, thread_name_prefix="join")
        self.conn = None
        self.cancelled = False
        self.stopping = False

    # Prices for the application ------------------------------------------------

    def price(self, market_id):
        """The current row of a market, suspended while not joined (with the
        example's own reason `disconnected`, as the client libraries do)."""
        row = self.rows.get(market_id)
        if row is None or (self.joined and time.monotonic() - self.last_heartbeat <= CLIENT_SILENCE_LIMIT):
            return row
        return dict(row, state="suspended", reason="disconnected", uncertainty=None,
                    outcomes=[dict(o, prob=None, fair_odds=None) for o in row["outcomes"]])

    def bound(self, key):
        return any(topic_match(b, key) for b in self.cfg.bindings)

    # Step 1 ---------------------------------------------------------------------

    def run(self):
        attempt = 0
        while not self.stopping:
            self.suspend("connecting")
            if attempt:
                delay = backoff(attempt)
                log.info("reconnecting in %.1f s", delay)
                self.sleep(delay)
                if self.stopping:
                    break
            attempt += 1
            try:
                self.connect()
                self.consume()
            except pika.exceptions.ProbableAuthenticationError:
                log.error("login refused: check FEED_AMQP_USER and FEED_AMQP_PASSWORD")
            except pika.exceptions.ProbableAccessDeniedError:
                log.error("access to the virtual host refused: the account is disabled, contact support")
            except (pika.exceptions.AMQPError, OSError) as e:
                log.warning("connection or channel closed: %s", describe(e))
            finally:
                if self.state == "live" or self.joined:
                    attempt = 1  # the backoff starts again after a successful join
                self.to_state("connect")
                self.close()
        self.close()

    def connect(self):
        cfg = self.cfg
        tls = ssl.create_default_context(cafile=cfg.ca_file)  # trusts the pinned CA only
        tls.minimum_version = ssl.TLSVersion.TLSv1_2
        params = pika.ConnectionParameters(
            host=cfg.host, port=cfg.port, virtual_host=cfg.vhost,
            credentials=pika.PlainCredentials(cfg.user, cfg.password),
            ssl_options=pika.SSLOptions(tls, server_hostname=cfg.host),
            heartbeat=AMQP_CONNECTION_HEARTBEAT, connection_attempts=1, socket_timeout=10,
            client_properties={"connection_name": f"{cfg.user} example consumer"})
        log.info("connecting to %s:%s as %s", cfg.host, cfg.port, cfg.user)
        self.conn = pika.BlockingConnection(params)
        self.cancelled = False
        ch = self.conn.channel()
        rch = self.conn.channel()  # results on a channel of their own
        exchanges = ["feed.prices"] + (["feed.events"] if cfg.events else [])
        for key in cfg.unbind:  # bindings an earlier run left, which the broker keeps
            for exchange in ("feed.prices", "feed.events"):
                ch.queue_unbind(cfg.prices_queue, exchange, key)
        for key in cfg.bindings:
            for exchange in exchanges:
                ch.queue_bind(cfg.prices_queue, exchange, key)
        ch.add_on_cancel_callback(self.on_cancel)
        rch.add_on_cancel_callback(self.on_cancel)
        self.to_state("wait")
        # Step 2: prices with automatic acknowledgement, results with manual ones.
        ch.basic_consume(cfg.prices_queue, self.on_prices_queue, auto_ack=True)
        rch.basic_qos(prefetch_count=RESULTS_PREFETCH)
        rch.basic_consume(cfg.results_queue, self.on_results_queue, auto_ack=False)
        self.ch, self.rch = ch, rch
        log.info("consuming %s and %s; waiting for a heartbeat", cfg.prices_queue, cfg.results_queue)

    def consume(self):
        while not self.stopping:
            self.check_timers()
            self.conn.process_data_events(time_limit=0.1)
            if self.cancelled:
                log.warning("the broker cancelled the consumer: rejoining from step 1")
                return
            if self.ch.is_closed or self.rch.is_closed:
                log.warning("a channel closed: rejoining from step 1")
                return
            self.check_timers()

    def close(self):
        if self.conn is not None and self.conn.is_open:
            try:
                self.conn.close()
            except (pika.exceptions.AMQPError, OSError):
                pass
        self.conn = None

    # Steps 3 to 5 -----------------------------------------------------------------

    def to_state(self, state):
        if state != "live":
            self.held = []
            self.join_generation += 1  # a join in flight is abandoned
            self.suspend({"connect": "connection closed", "wait": "rejoining"}.get(state, state))
        if state == "wait":
            self.wait_since = time.monotonic()
        self.state = state

    def suspend(self, reason):
        if self.joined:
            self.joined = False
            log.warning("every price suspended: %s", reason)
            self.app.suspended(reason)

    def check_timers(self):
        now = time.monotonic()
        if self.state in ("join", "live") and now - self.last_heartbeat > CLIENT_SILENCE_LIMIT:
            self.rejoin("no heartbeat for 3 s")  # rule 1
        elif self.state == "wait" and now - self.wait_since > JOIN_TIMEOUT:
            log.warning("no usable heartbeat for 10 s (a standby consumer, or the feed is not sending); waiting")
            self.wait_since = now
        elif self.state == "join" and now - self.join_started >= JOIN_GIVE_UP:
            self.rejoin("no join answer within 60 s")
        elif self.state == "join" and self.join_future.done():
            self.finish_join()

    def rejoin(self, reason):
        log.warning("rejoining: %s", reason)
        self.suspend(reason)
        self.to_state("wait")

    def on_prices_queue(self, _ch, _method, _props, body):
        # A callback after a stalled loop must not erase an already missed deadline.
        if self.state in ("join", "live") and time.monotonic() - self.last_heartbeat > CLIENT_SILENCE_LIMIT:
            self.rejoin("no heartbeat for 3 s")
        received = time.time()  # when this code takes the message from pika
        try:
            m = json.loads(body)
        except ValueError:
            log.warning("ignored a message that is not JSON")
            return
        kind = m.get("type")
        if kind == "result":
            self.on_result(m, "prices")
            return
        if self.state == "wait":
            # Step 3: everything up to and including the join heartbeat is dropped.
            if kind == "heartbeat" and not m.get("final") and received - utc_seconds(m["sent_at"]) <= CLIENT_LAG_LIMIT:
                self.start_join(m)
            return
        if kind == "heartbeat":
            problem = self.heartbeat_problem(m, received)
            if problem:
                self.rejoin(problem)
                return
            self.seq = m["seq"]
            self.last_heartbeat = time.monotonic()
            if self.state == "live":
                self.app.heartbeat(m)
        elif self.state == "join":
            if len(self.held) >= 200_000:
                self.rejoin("join buffer full")
            else:
                self.held.append(m)  # applied once the join answer is in
        else:
            self.apply(m)

    def heartbeat_problem(self, hb, received):
        """Rules 2 and 3; None when the heartbeat is in order."""
        if hb.get("session") != self.session:
            return "the feed started a new session"
        if hb.get("seq") != self.seq + 1:
            return "a heartbeat is missing"
        if hb.get("final"):
            return "the feed is stopping"
        if received - utc_seconds(hb["sent_at"]) > CLIENT_LAG_LIMIT:
            return "more than 3 s behind the queue"
        return None

    def start_join(self, hb):
        self.session, self.seq = hb["session"], hb["seq"]
        self.last_heartbeat = time.monotonic()
        self.join_started = self.last_heartbeat
        self.held = []
        self.state = "join"
        self.join_generation += 1
        self.join_future = self.joins.submit(self.request_join, self.join_generation)
        log.info("heartbeat %s/%s: requesting the join answer", self.session, self.seq)

    def request_join(self, generation):
        """Step 4, on a thread of its own so that messages keep arriving: the join
        answer and, with events, every page of GET /v1/events."""
        deadline = time.monotonic() + JOIN_GIVE_UP
        attempt = 0
        while time.monotonic() < deadline and generation == self.join_generation and not self.stopping:
            try:
                answer = self.http("POST", "/v1/prices/join", {"filters": self.cfg.filters})
                events = self.read_events() if self.cfg.events else {}
                return answer, events
            except RetryAfter as e:
                delay = e.seconds
                log.info("join: %s", e)
            except urllib.error.HTTPError as e:
                attempt += 1
                delay = backoff(attempt)
                if e.code >= 500:
                    log.warning("join: HTTP %s, retrying in %.1f s", e.code, delay)
                else:  # 400, 401, 403: a configuration problem; prices stay suspended
                    log.error("join refused with HTTP %s (check FEED_API_KEY and its AMQP access); retrying in %.1f s",
                              e.code, delay)
            except (urllib.error.URLError, OSError, ValueError) as e:
                attempt += 1
                delay = backoff(attempt)
                log.warning("join: %s, retrying in %.1f s", describe(e), delay)
            end = min(deadline, time.monotonic() + delay)
            while time.monotonic() < end and generation == self.join_generation and not self.stopping:
                time.sleep(0.1)
        return None

    def http(self, method, path, body=None):
        request = urllib.request.Request(
            self.cfg.api_url + path, method=method,
            data=None if body is None else json.dumps(body).encode(),
            headers={"Authorization": f"Bearer {self.cfg.api_key}", "Accept-Encoding": "gzip",
                     "Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(request, timeout=JOIN_REQUEST_TIMEOUT, context=self.api_ssl) as r:
                data = r.read()
                if r.headers.get("Content-Encoding") == "gzip":
                    data = gzip.decompress(data)
                return json.loads(data)
        except urllib.error.HTTPError as e:
            if e.code in (429, 503):
                retry = e.headers.get("Retry-After", "1")
                raise RetryAfter(int(retry) if retry.isdigit() else 1) from None
            raise

    def read_events(self):
        events, cursor = {}, None
        while True:
            query = "limit=500" + (f"&cursor={urllib.parse.quote(cursor)}" if cursor else "")
            page = self.http("GET", f"/v1/events?{query}")
            for event in page["data"]:
                if self.bound(event_key(event)):
                    events[event["id"]] = event
            cursor = page.get("next_cursor")
            if not cursor:
                return events

    def finish_join(self):
        try:
            result = self.join_future.result()
        except Exception as e:  # noqa: BLE001 - any failure of the join thread
            log.warning("join failed: %s", describe(e))
            result = None
        if result is None:
            self.rejoin("no join answer within 60 s")
            return
        answer, events = result
        self.rows = {row["market_id"]: row for row in answer["data"]}
        self.events = events
        self.removed_events = {}
        self.join_version = answer["version"]
        held, self.held = self.held, []
        for m in held:  # step 5: what was held since the join heartbeat, in order
            self.apply(m)
        self.state = "live"
        self.joined = True
        log.info("joined at version %s: %d rows, %d events", self.join_version, len(self.rows), len(self.events))
        self.app.joined(self)

    def apply(self, m):
        kind = m.get("type")
        if kind == "result":
            self.on_result(m, "prices")
        elif kind == "price":
            event = m["event"]["id"]
            if m["version"] <= max(self.version_held(self.rows.get(m["market_id"])),
                                    self.removed_events.get(event, 0)):
                return
            if not self.bound(row_key(m)):
                self.drop_event(event, m["version"])  # it left the bindings
            elif m["version"] > self.version_held(self.rows.get(m["market_id"])):
                self.rows[m["market_id"]] = m
                if self.joined:
                    self.app.row(m)
        elif kind == "event" and self.cfg.events:
            if m["version"] <= max(self.version_held(self.events.get(m["id"])),
                                    self.removed_events.get(m["id"], 0)):
                return
            if not self.bound(event_key(m)):
                self.drop_event(m["id"], m["version"])
            elif m["version"] > self.version_held(self.events.get(m["id"])):
                self.events[m["id"]] = m
                if self.joined:
                    self.app.event(m)
        # anything else is ignored (API-015)

    def version_held(self, held):
        """A row or event not held counts as held at the join version."""
        return held["version"] if held is not None else self.join_version

    def drop_event(self, event_id, version):
        markets = [mid for mid, row in self.rows.items() if row["event"]["id"] == event_id]
        newest = max([self.version_held(self.events.get(event_id)), self.removed_events.get(event_id, 0)]
                     + [self.rows[mid]["version"] for mid in markets])
        if version <= newest:
            return
        self.removed_events[event_id] = version
        for mid in markets:
            del self.rows[mid]
        had_event = self.events.pop(event_id, None) is not None
        if self.joined and (markets or had_event):
            self.app.dropped(event_id, markets)

    # Results ----------------------------------------------------------------------

    def on_results_queue(self, ch, method, _props, body):
        self.on_result(json.loads(body), "results")
        ch.basic_ack(method.delivery_tag)  # only after it is stored

    def on_result(self, row, queue):
        """Results are kept by market whatever key they came under, never filtered."""
        if row.get("type") == "result" and self.results.store(row):
            self.app.result(row, queue)

    def on_cancel(self, _frame):
        self.cancelled = True

    def sleep(self, seconds):
        end = time.monotonic() + seconds
        while time.monotonic() < end and not self.stopping:
            time.sleep(0.1)

    def stop(self, *_):
        self.stopping = True


def describe(error):
    """An error in a few words, without anything secret."""
    text = str(error) or type(error).__name__
    return f"{type(error).__name__}: {text}"[:300]


def main():
    logging.basicConfig(stream=sys.stderr, level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    logging.getLogger("pika").setLevel(logging.ERROR)
    consumer = FeedConsumer(Config(os.environ), Printer())
    signal.signal(signal.SIGTERM, consumer.stop)
    signal.signal(signal.SIGINT, consumer.stop)
    consumer.app.suspended("starting")
    consumer.run()
    consumer.joins.shutdown(wait=False, cancel_futures=True)
    log.info("stopped")


if __name__ == "__main__":
    main()