#!/usr/bin/env python3
"""job_subscriber.py — the single entry point for observing Job events.

Subscribes to one job's ``<topic_prefix>/events`` (or, with ``--wait-any``, the
events of every running/pending job in the registry), prints one line to stdout
per accepted event, and exits on a terminal event or a timeout.

Design points (all flagged in the PLAN review):
  - terminal state machine: ``completed``/``error`` is acted on exactly once per
    job, so QoS-1 duplicates or an ``error``-after-``completed`` reorder are safe.
  - dual timeouts: a wall-clock ``--timeout`` (total budget, started at
    subscribe time so a cold start can't hang forever) AND an idle
    ``--idle-timeout`` (no new event for N seconds).
  - defensive parsing: undecodable payloads, ``schema_version`` mismatches, and
    ``job_id`` values we did not subscribe for are logged and dropped.

stdout = event lines only. Diagnostics go to stderr via logging.

Exit codes:
  0  all watched jobs reached ``completed``
  1  any watched job reached ``error``
  2  timed out (wall-clock or idle) before all jobs finished
"""
from __future__ import annotations

import argparse
import json
import logging
import queue
import sys
import time
from typing import Any, Dict, List, Optional, Set, Tuple

import mqtt_common
import registry
from mqtt_common import (
    DEFAULT_REGISTRY_DIR,
    SCHEMA_VERSION,
    broker_config_from_job,
    load_job,
    make_client,
)

logger = logging.getLogger("delegate_job.job_subscriber")

TERMINAL_EVENTS = ("completed", "error")


def _format_line(topic: str, payload: Dict[str, Any]) -> str:
    source_tag = f" [{payload['source']}]" if payload.get("source") else ""
    return (
        f"{payload.get('timestamp','-')}  "
        f"job={payload.get('job_id','?')}  "
        f"seq={payload.get('seq','?')}  "
        f"{payload.get('event','?') + source_tag:<20}  "
        f"{payload.get('detail','')}"
    )


def _check_disk_fallback(
    pending: Set[str],
    registry_dir: str,
    terminal: Dict[str, str],
) -> None:
    """Check local on-disk registry and audit logs for terminal status (B-15)."""
    for jid in list(pending):
        disk_status = None
        try:
            job_rec = load_job(jid, registry_dir)
            disk_status = job_rec.get("status")
        except Exception:
            pass
        if not disk_status or disk_status not in ("completed", "error", "cancelled"):
            try:
                disk_status = mqtt_common.read_logged_status(jid, mqtt_common.get_logs_dir())
            except Exception:
                pass

        if disk_status in ("completed", "error", "cancelled"):
            synth_event = "completed" if disk_status == "completed" else "error"
            synth_payload = {
                "job_id": jid,
                "event": synth_event,
                "seq": -1,
                "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                "detail": f"terminal status '{disk_status}' resolved via disk-fallback",
                "source": "disk-fallback",
            }
            synth_topic = f"mam/{jid}/events"
            print(_format_line(synth_topic, synth_payload), flush=True)
            terminal[jid] = synth_event
            pending.discard(jid)


class _Watcher:
    """Holds the shared queue + the set of job_ids we accept events for."""

    def __init__(self, expected_job_ids: Set[str], expected_tokens: Dict[str, Optional[str]], expected_seqs: Dict[str, int]):
        self.events: "queue.Queue[Tuple[str, Dict[str, Any]]]" = queue.Queue()
        self.expected = set(expected_job_ids)
        self.tokens = expected_tokens  # job_id -> expected auth_token (or None)
        self.last_seq = dict(expected_seqs)

    def on_message(self, _client, _userdata, msg) -> None:
        # --- defensive parsing -------------------------------------------
        try:
            payload = json.loads(msg.payload.decode("utf-8"))
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            logger.warning("drop unparseable payload on %s: %s", msg.topic, exc)
            return
        if not isinstance(payload, dict):
            logger.warning("drop non-object payload on %s", msg.topic)
            return
        if payload.get("schema_version") != SCHEMA_VERSION:
            logger.warning("drop event with schema_version=%r (expected %d)",
                           payload.get("schema_version"), SCHEMA_VERSION)
            return
        jid = payload.get("job_id")
        if jid not in self.expected:
            logger.warning("drop event for unexpected job_id=%r on %s", jid, msg.topic)
            return
        # --- production auth check: data.auth_token must match if expected ---
        expected_token = self.tokens.get(jid)
        if not mqtt_common.verify_hmac(payload, expected_token):
            logger.warning("drop event for job %s: HMAC verify failed", jid)
            return
        # --- replay attack defense: check monotonic sequence ---
        seq = payload.get("seq")
        if seq is None or not isinstance(seq, int):
            logger.warning("drop event for job %s: missing or invalid seq", jid)
            return
        if seq <= self.last_seq.get(jid, 0):
            logger.warning("drop event for job %s: seq %d is not monotonically increasing (last %d)",
                           jid, seq, self.last_seq.get(jid, 0))
            return
        self.last_seq[jid] = seq
        # Persistent audit log from the *subscriber's* vantage point: every event
        # that survives defensive parsing is recorded here, including ones a
        # different host published. This is the external-observer record that
        # backstops the publisher's own "published" line if it never wrote one.
        mqtt_common.append_event(jid, {
            "event": "received",
            "source_event": payload.get("event"),
            "seq": payload.get("seq"),
            "topic": msg.topic,
            "timestamp": payload.get("timestamp"),
            "detail": payload.get("detail", ""),
        })
        self.events.put((msg.topic, payload))


def _collect_jobs(args) -> List[Dict[str, Any]]:
    """Resolve the list of job records this invocation should watch."""
    if args.wait_any:
        jobs = [r for r in registry.list_jobs(args.registry_dir)
                if r.get("status") in ("pending", "running")]
        if not jobs:
            logger.error("no pending/running jobs to wait for")
        return jobs
    job = load_job(args.job, args.registry_dir)  # raises FileNotFoundError
    return [job]


def _run_subscriber(args) -> int:
    try:
        jobs = _collect_jobs(args)
    except FileNotFoundError as exc:
        logger.error("%s", exc)
        return 2
    if not jobs:
        return 2

    expected_ids: Set[str] = {j["job_id"] for j in jobs}
    tokens = {j["job_id"]: j.get("auth_token") for j in jobs}
    seqs = {}
    for j in jobs:
        jid = j["job_id"]
        last_seq = int(j.get("last_seq", 0))
        seqs[jid] = max(0, last_seq - 1)
    watcher = _Watcher(expected_ids, tokens, seqs)

    # Resolve timeouts from CLI, falling back to the (first) job's settings.
    base_job = jobs[0]
    wall_timeout = args.timeout if args.timeout is not None else float(base_job.get("timeout_sec", 3600))
    idle_timeout = args.idle_timeout if args.idle_timeout is not None else float(base_job.get("idle_timeout_sec", 120))

    # All watched jobs share a broker in practice; connect using the first
    # job's broker and subscribe to each job's events topic.
    config = broker_config_from_job(base_job)
    client = make_client("subscriber", config)
    client.on_message = watcher.on_message

    subscribed_topics = []
    for job in jobs:
        prefix = job.get("topic_prefix") or mqtt_common.topic_prefix_for(job["job_id"])
        subscribed_topics.append(f"{prefix}/events")

    def on_connect(_c, _u, _flags, reason_code, _props):
        if mqtt_common.reason_code_value(reason_code) != 0:
            logger.error("broker connection failed: rc=%s", reason_code)
            return
        for topic in subscribed_topics:
            _c.subscribe(topic, qos=1)
            logger.info("subscribed to %s", topic)

    def on_disconnect(_c, _u, _flags, reason_code, _props):
        rc = mqtt_common.reason_code_value(reason_code)
        if rc != 0:
            logger.warning("broker disconnected (rc=%s); will retry reconnect", reason_code)

    def on_subscribe(_c, _u, mid, granted_qos, _props=None):
        for topic in subscribed_topics:
            print(f"SUBSCRIBED {topic}", flush=True)

    client.on_connect = on_connect
    client.on_disconnect = on_disconnect
    client.on_subscribe = on_subscribe
    client.reconnect_delay_set(min_delay=1, max_delay=16)

    terminal: Dict[str, str] = {}        # job_id -> "completed"/"error"
    pending: Set[str] = set(expected_ids)
    start = time.monotonic()
    wall_deadline = start + wall_timeout
    last_event = start
    last_disk_check = 0.0
    retention_checked = not args.expect_retention
    connected = False

    # Check if all jobs are already in terminal state on disk before connecting
    _check_disk_fallback(pending, args.registry_dir, terminal)
    if not pending:
        if any(state == "error" for state in terminal.values()):
            return 1
        return 0

    try:
        mqtt_common.with_retry(
            lambda: client.connect(config.host, config.port, config.keepalive),
            attempts=3, base_delay=0.2, max_delay=1.0
        )()
        client.loop_start()
        connected = True
    except Exception as exc:
        logger.warning("broker connection failed (%s); checking disk fallback", exc)
        _check_disk_fallback(pending, args.registry_dir, terminal)
        if not pending:
            if any(state == "error" for state in terminal.values()):
                return 1
            return 0
        logger.error("broker connection failed: %s", exc)
        return 3

    try:
        while pending:
            now = time.monotonic()
            if now >= wall_deadline:
                _check_disk_fallback(pending, args.registry_dir, terminal)
                if not pending:
                    break
                logger.error("wall-clock timeout (%.0fs); still pending: %s",
                             wall_timeout, ", ".join(sorted(pending)))
                return 2
            idle_left = idle_timeout - (now - last_event)
            if idle_left <= 0:
                _check_disk_fallback(pending, args.registry_dir, terminal)
                if not pending:
                    break
                logger.error("idle timeout (%.0fs, no events); still pending: %s",
                             idle_timeout, ", ".join(sorted(pending)))
                return 2
            wait = min(wall_deadline - now, idle_left, 1.0)
            try:
                topic, payload = watcher.events.get(timeout=wait)
            except queue.Empty:
                if not retention_checked and (now - start) > 3.0:
                    logger.warning("--expect-retention set but no retained "
                                   "terminal event observed yet")
                    retention_checked = True

                # Step 2 (B-15): Local disk fallback throttled to every 3.0s
                if (now - last_disk_check) >= 3.0:
                    last_disk_check = now
                    _check_disk_fallback(pending, args.registry_dir, terminal)
                continue

            last_event = time.monotonic()
            retention_checked = True
            print(_format_line(topic, payload), flush=True)

            jid = payload["job_id"]
            event = payload.get("event")
            if event in TERMINAL_EVENTS:
                if jid in terminal:
                    # Already finalised: ignore duplicates / late reorders.
                    logger.info("ignoring duplicate terminal %s for %s", event, jid)
                    continue
                terminal[jid] = event
                pending.discard(jid)
    finally:
        if connected:
            client.loop_stop()
            try:
                client.disconnect()
            except Exception:  # pragma: no cover
                pass

    # All jobs reached a terminal state. error wins over completed.
    if any(state == "error" for state in terminal.values()):
        return 1
    return 0


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description="Subscribe to Job events on MQTT")
    target = parser.add_mutually_exclusive_group(required=True)
    target.add_argument("--job", help="job id to watch")
    target.add_argument("--wait-any", action="store_true",
                        help="watch every pending/running job in the registry")
    parser.add_argument("--timeout", type=float, default=None,
                        help="wall-clock budget in seconds (default: job.timeout_sec or 3600)")
    parser.add_argument("--idle-timeout", type=float, default=None,
                        help="max seconds with no new event (default: job.idle_timeout_sec or 120)")
    parser.add_argument("--expect-retention", action="store_true",
                        help="warn if no retained terminal event arrives promptly")
    parser.add_argument("--registry-dir", default=DEFAULT_REGISTRY_DIR)
    parser.add_argument("-v", "--verbose", action="store_true")
    args = parser.parse_args(argv)

    mqtt_common.setup_logging(logging.DEBUG if args.verbose else logging.WARNING)

    try:
        return _run_subscriber(args)
    except Exception as exc:
        logger.error("subscriber fatal error: %s", exc)
        return 3


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