# -*- coding: utf-8 -*-
"""Feedback watcher for a dsh (DeepSeek Harness) background sub-session.

The contract: run me in the background; my EXIT is the wake signal.
------------------------------------------------------------------
A coordinator spawns a dsh session to do a piece of work, then launches THIS
script as a background job and goes back to its own business. dsh does not push
a "your sub is done / your sub is stuck" event to the coordinator, so this
watcher stands in for one: it polls the session and exits the instant the
session finishes, stalls (typically sitting on a permission/approval prompt in
the dsh web UI), or errors. The coordinator's harness delivers the watcher's
stdout when it exits, and that exit is what wakes the coordinator back up.

It prints exactly ONE line and exits with a code that says which of the five
things happened, so the coordinator can branch on the code alone:

    DONE     exit 0   turn/end reached, activity counter frozen
    CUT_OFF  exit 5   turn/end reached but the reply was truncated at the
                      model's output cap (finish reason "max-tokens")
    ERROR    exit 2   RPC error / result.ok false / session not found
    STALLED  exit 3   activity counter frozen past --stall, turn not ended
    TIMEOUT  exit 4   overall --max-minutes cap hit, session still active

How liveness is judged
----------------------
dsh serves a local JSON-RPC at POST <base>/api/<method>. We call
`session.history` and read four small things out of the (possibly huge)
response:

  * `value.projections.asOfSeq` -- a monotonically increasing ACTIVITY counter.
    It climbs while the session does work and FREEZES while the session sits
    idle or blocked on a prompt. This is the liveness signal.
  * `value.events[-1]["event"]["type"]` -- the tail event type. A finished/idle
    turn ends with `turn/end`.
  * the last `assistant/chunk` event whose `data.chunk.type == "finish"` --
    its `data.chunk.reason.kind` says HOW the model's turn ended. dsh emits a
    `finish` chunk on a normal stop (`reason.kind == "stop"`) AND when the
    reply is truncated by the output cap (`reason.kind == "max-tokens"`), so
    this is what tells a genuinely finished turn from a cut-off one.
  * `value.projections.values.permissions.currentValue` -- the current sandbox
    level, surfaced in the STALLED line so the coordinator can see whether the
    sub is blocked wanting a higher level.

The events array can carry tens of thousands of entries; we never print it.
The tail event is enough for liveness; distinguishing DONE from CUT_OFF means
scanning back through the events for the last `finish` chunk, which is a cheap
walk over the same array we already hold. `session.history` + `asOfSeq` is a
correct, sufficient baseline. A lighter RPC (e.g. a per-session running/idle
flag from `session.list`/`session.get`) could be swapped into `_poll` if one is
confirmed available; nothing else would have to change.

CPython 3, stdlib only. Lives beside rvt.py / qna_watch.py, not under the
IronPython extension.
"""

import argparse
import collections
import json
import os
import sys
import time
import uuid

try:                                    # py3
    import urllib.request as _urlreq
    import urllib.error as _urlerr
except ImportError:                     # pragma: no cover - py2 never runs this
    import urllib2 as _urlreq
    _urlerr = _urlreq

DEFAULT_BASE_URL = 'http://127.0.0.1:3080'
DASH = u'—'                        # em dash used in every output line

# One decision per cycle. kind is for programs/tests, message is the one line
# printed to stdout, code is the process exit status.
Verdict = collections.namedtuple('Verdict', 'kind code message')


def _base_url(base_url=None):
    if base_url:
        return base_url
    return os.environ.get('DSH_BASE_URL', DEFAULT_BASE_URL)


def _post_json(url, body, timeout):
    """The one and only network call, isolated so a test can monkeypatch it.

    Returns the decoded JSON response (a dict). Raises on any transport or
    decode failure -- the caller (`_poll`) turns that into an ok:False dict.
    """
    data = json.dumps(body).encode('utf-8')
    req = _urlreq.Request(url, data=data,
                          headers={'Content-Type': 'application/json'})
    resp = _urlreq.urlopen(req, timeout=timeout)
    try:
        raw = resp.read()
    finally:
        try:
            resp.close()
        except Exception:
            pass
    return json.loads(raw.decode('utf-8'))


def _poll(sid, base_url=None, timeout=10.0):
    """One liveness sample. Returns a compact dict the loop can classify:

        {"ok": True, "asOfSeq": <int>, "tail": <str|None>, "perm": <str|None>,
         "finish": <str|None>}
        {"ok": False, "reason": <str>}

    `finish` is the `reason.kind` of the LAST assistant/finish chunk in the
    events ("stop" on a normal end, "max-tokens" when the reply was cut off by
    the output cap, None when no finish chunk is present or it cannot be read).

    Never raises: a transport hiccup, a malformed body, result.ok false, or a
    missing session all come back as ok:False with a reason. The loop retries a
    couple of times before it believes an ok:False, so a single hiccup is not a
    dead session.
    """
    url = _base_url(base_url).rstrip('/') + '/api/session.history'
    body = {'type': 'client-request',
            'rpcId': str(uuid.uuid4()),
            'method': 'session.history',
            'payload': {'sessionId': sid}}
    try:
        resp = _post_json(url, body, timeout)
    except Exception as exc:                       # transport / decode failure
        return {'ok': False, 'reason': 'poll failed: {0}'.format(exc)}

    if not isinstance(resp, dict):
        return {'ok': False, 'reason': 'malformed response (not an object)'}

    result = resp.get('result')
    if not isinstance(result, dict) or not result.get('ok'):
        reason = None
        if isinstance(result, dict):
            reason = result.get('error') or result.get('reason')
        return {'ok': False,
                'reason': reason or 'rpc result not ok (session not found?)'}

    value = result.get('value') or {}
    proj = value.get('projections') or {}
    seq = proj.get('asOfSeq')
    if seq is None:
        return {'ok': False, 'reason': 'no asOfSeq in projections'}

    perms = (proj.get('values') or {}).get('permissions') or {}
    perm = perms.get('currentValue')

    tail = None
    finish = None
    events = value.get('events') or []
    if events:
        last = events[-1]
        ev = last.get('event') if isinstance(last, dict) else None
        if isinstance(ev, dict):
            tail = ev.get('type')
        # The LAST assistant/finish chunk carries how the model's turn ended:
        # its reason.kind is "stop" on a normal end and "max-tokens" when the
        # reply was truncated by the output cap. Walk back so a finish chunk
        # after unrelated events is still found; read defensively throughout
        # (any dict in the chain may be missing a key).
        for entry in reversed(events):
            if not isinstance(entry, dict):
                continue
            ev = entry.get('event')
            if not isinstance(ev, dict):
                continue
            if ev.get('type') != 'assistant/chunk':
                continue
            data = ev.get('data')
            if not isinstance(data, dict):
                continue
            chunk = data.get('chunk')
            if not isinstance(chunk, dict):
                continue
            if chunk.get('type') != 'finish':
                continue
            reason = chunk.get('reason')
            if isinstance(reason, dict):
                finish = reason.get('kind')
            break

    try:
        seq = int(seq)
    except (TypeError, ValueError):
        return {'ok': False, 'reason': 'asOfSeq not an int: {0!r}'.format(seq)}

    return {'ok': True, 'asOfSeq': seq, 'tail': tail, 'perm': perm,
            'finish': finish}


def _poll_with_retry(sid, base_url, sleep, attempts=3, retry_delay=1.0):
    """Call `_poll` up to `attempts` times; a transient failure is not a dead
    session. Looks `_poll` up on the module at call time so tests can swap it.
    """
    reason = 'poll failed'
    for i in range(attempts):
        try:
            res = _poll(sid, base_url)
        except Exception as exc:                   # a monkeypatched/broken poll
            res = {'ok': False, 'reason': 'poll raised: {0}'.format(exc)}
        if res.get('ok'):
            return res
        reason = res.get('reason') or reason
        if i < attempts - 1:
            sleep(retry_delay)
    return {'ok': False, 'reason': reason}


def watch(sid, interval=15.0, stall=60.0, max_minutes=120.0, base_url=None,
          sleep=None, now=None):
    """Poll `sid` until it finishes, stalls, errors, or the cap is hit.

    Returns a Verdict(kind, code, message). Pure of I/O: `main` prints and
    exits with it. `sleep` and `now` are injectable so tests run instantly with
    a fake clock; production uses time.sleep / time.monotonic.
    """
    sleep = sleep or time.sleep
    now = now or time.monotonic

    prev_seq = None
    start = now()
    last_change = start

    while True:
        res = _poll_with_retry(sid, base_url, sleep)
        if not res.get('ok'):
            return Verdict('ERROR', 2,
                           u'ERROR {0} {1} {2}'.format(sid, DASH, res.get('reason')))

        seq = res['asOfSeq']
        tail = res.get('tail')
        perm = res.get('perm')
        finish = res.get('finish')
        t = now()

        advanced = (prev_seq is None) or (seq != prev_seq)
        if advanced:
            last_change = t
        prev_seq = seq

        frozen = t - last_change
        elapsed = t - start

        # DONE / CUT_OFF: the turn has ended and the counter is no longer
        # advancing. We require a non-advancing cycle so we never call a turn
        # done off the very poll that emitted turn/end (that poll itself
        # advances the seq). The last finish chunk's reason.kind splits the two
        # cases: "max-tokens" means dsh ended the turn because the reply hit
        # the model's output cap -- a CUT_OFF the coordinator should resume
        # with "continue", not treat as a finished job.
        if tail == 'turn/end' and not advanced:
            if finish == 'max-tokens':
                return Verdict(
                    'CUT_OFF', 5,
                    u'CUT_OFF {0} {1} reply cut off at max-tokens (turn/end '
                    u'reached); send "continue" to resume'.format(sid, DASH))
            return Verdict('DONE', 0,
                           u'DONE {0} steps={1} {2} turn ended'.format(sid, seq, DASH))

        # STALLED: counter frozen past the threshold and the turn is NOT over --
        # the classic "waiting on a permission/approval prompt" shape.
        if frozen >= stall and tail != 'turn/end':
            return Verdict('STALLED', 3,
                           u'STALLED {0} {1} asOfSeq frozen {2}s at seq={3}, '
                           u'tail={4}, perm={5}; check the dsh web UI for a '
                           u'pending approval/prompt'.format(
                               sid, DASH, int(round(frozen)), seq, tail, perm))

        # TIMEOUT: overall cap; the session is still doing something.
        if elapsed > max_minutes * 60.0:
            return Verdict('TIMEOUT', 4,
                           u'TIMEOUT {0} {1} watched {2}m, still active'.format(
                               sid, DASH, int(round(elapsed / 60.0))))

        sleep(interval)


def _emit(verdict):
    try:
        sys.stdout.reconfigure(encoding='utf-8', errors='replace')
    except Exception:
        pass
    try:
        print(verdict.message)
    except UnicodeEncodeError:
        sys.stdout.write(verdict.message.encode('ascii', 'replace').decode('ascii'))
        sys.stdout.write('\n')
    return verdict.code


def main(argv=None):
    ap = argparse.ArgumentParser(
        description='Watch a dsh background sub-session and exit when it '
                    'finishes, stalls, or errors.')
    ap.add_argument('session_id', help='the dsh sessionId to watch')
    ap.add_argument('--interval', type=float, default=15.0,
                    help='seconds between polls (default 15)')
    ap.add_argument('--stall', type=float, default=60.0,
                    help='seconds of frozen asOfSeq before STALLED (default 60)')
    ap.add_argument('--max-minutes', type=float, default=120.0,
                    help='overall watch cap in minutes (default 120)')
    ap.add_argument('--base-url', default=_base_url(),
                    help='dsh RPC base URL (default $DSH_BASE_URL or '
                         + DEFAULT_BASE_URL + ')')
    args = ap.parse_args(argv)

    verdict = watch(args.session_id,
                    interval=args.interval,
                    stall=args.stall,
                    max_minutes=args.max_minutes,
                    base_url=args.base_url)
    return _emit(verdict)


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