"""Read-only ComfyUI telemetry, embedded in the existing private normalizers.

Only the requesting Huobao job's client connection is observed. No queue mutation
or image previews are used. Sampling ETA is never presented as whole-video ETA.
"""
import asyncio as _hp_asyncio
import json as _hp_json
import time as _hp_time
import math as _hp_math
import statistics as _hp_stats
from urllib.parse import quote as _hp_quote
import httpx as _hp_httpx
import websockets as _hp_ws

_hp_watchers = {}
_hp_histories = {}


def _hp_phase(kind):
    kind = str(kind).lower()
    if 'sampler' in kind: return 'sampling'
    if 'decode' in kind: return 'decoding'
    if 'savevideo' in kind or 'combine' in kind or 'createvideo' in kind: return 'encoding'
    if 'loader' in kind or kind == 'loadimage': return 'loading'
    if 'caption' in kind: return 'postprocessing'
    if 'encode' in kind or 'tovideo' in kind: return 'conditioning'
    return 'running'


def _hp_signature(graph):
    # Keep graph topology, model weights, image count, dimensions, FPS, duration,
    # steps and sampler. Exclude content, random seeds and output filenames.
    omitted = {'text', 'prompt', 'positive', 'negative', 'seed', 'noise_seed',
               'filename_prefix', 'image', 'audio', 'video', 'file', 'path'}
    result = {}
    for node, value in graph.items():
        inputs = {}
        for key, item in value.get('inputs', {}).items():
            if key in omitted and not isinstance(item, list): continue
            if any(word in key.lower() for word in ['prompt', 'text', 'caption']) and not isinstance(item, list): continue
            inputs[key] = item
        result[node] = {'class_type': value.get('class_type'), 'inputs': inputs}
    return _hp_json.dumps(result, sort_keys=True, separators=(',', ':'))


def _hp_history_estimate(graph, history):
    signature = _hp_signature(graph)
    durations = []
    for entry in history.values():
        status = entry.get('status', {})
        prompt = entry.get('prompt', [])
        if status.get('status_str') != 'success' or not status.get('completed') or len(prompt) < 3: continue
        previous = prompt[2]
        if not isinstance(previous, dict) or _hp_signature(previous) != signature: continue
        times = {}; cached = []
        for event in status.get('messages', []):
            if len(event) < 2 or not isinstance(event[1], dict): continue
            times[event[0]] = event[1].get('timestamp')
            if event[0] == 'execution_cached': cached = event[1].get('nodes', [])
        if any('sampler' in str(previous.get(str(n), {}).get('class_type', '')).lower() for n in cached): continue
        start, end = times.get('execution_start'), times.get('execution_success')
        if not isinstance(start, (int, float)) or not isinstance(end, (int, float)): continue
        duration = (end - start) / 1000
        if 2 <= duration <= 86400: durations.append(duration)
    result = {'historySamples': len(durations)}
    if len(durations) >= 2:
        # A range, not a guarantee; model loading and shared GPU contention vary.
        result.update(historyMinSeconds=round(min(durations) * .8), historyMaxSeconds=round(max(durations) * 1.2))
    return result


def _hp_event(state, message, timestamp):
    data = message.get('data', {})
    if not isinstance(data, dict): return
    event = message.get('type')
    if data.get('prompt_id') not in (None, state['prompt']): return
    if event == 'executing':
        node = data.get('node')
        if node is None: return
        node = str(node)
        if node != state.get('node'):
            state.update(node=node, samples=[], value=None, total=None)
        state['phase'] = _hp_phase(state['classes'].get(node))
        state['updated'] = timestamp
    elif event == 'progress' and data.get('prompt_id') == state['prompt']:
        value, total = data.get('value'), data.get('max')
        if not all(type(x) in (int, float) and _hp_math.isfinite(x) for x in (value, total)): return
        if not 0 <= value <= total <= 1000000 or total <= 0: return
        node = str(data.get('node', state.get('node', '')))
        if node != state.get('node') or total != state.get('total') or value < (state.get('value') or 0):
            state.update(node=node, samples=[])
        if state.get('value') != value or not state['samples']:
            state['samples'] = (state['samples'] + [(timestamp, value)])[-8:]
        state.update(value=value, total=total, updated=timestamp, phase=_hp_phase(state['classes'].get(node)))


def _hp_snapshot(state, timestamp):
    result = {'phase': state.get('phase', 'running')}
    if not state.get('connected') or timestamp - state.get('updated', 0) > 600:
        return {'phase': 'running'}
    if state.get('value') is None: return result
    result.update(step=state['value'], totalSteps=state['total'], stepUpdatedAt=round(state['updated'] * 1000))
    samples = state['samples']
    if state.get('phase') == 'sampling' and len(samples) >= 3:
        seconds = samples[-1][0] - samples[0][0]; steps = samples[-1][1] - samples[0][1]
        if seconds >= 5 and steps > 0 and state['value'] < state['total']:
            speed = seconds / steps
            if timestamp - state['updated'] <= max(30, speed * 2.5):
                remaining = max(1, speed * (state['total'] - state['value']) - (timestamp - state['updated']))
                result.update(samplingMinSeconds=round(remaining * .8), samplingMaxSeconds=round(remaining * 1.3))
    return result


async def _hp_watch(state, base, headers, client_id):
    uri = base.replace('https://', 'wss://').replace('http://', 'ws://') + '/ws?clientId=' + _hp_quote(client_id, safe='')
    while _hp_time.time() - state['requested'] < 120:
        try:
            async with _hp_ws.connect(uri, additional_headers=headers, max_size=32 * 1024 * 1024, open_timeout=5) as ws:
                state['connected'] = True
                # This observer needs text progress only, never preview images.
                await ws.send(_hp_json.dumps({'type':'feature_flags','data':{'supports_preview_metadata':True}}))
                while _hp_time.time() - state['requested'] < 120:
                    try: message = await _hp_asyncio.wait_for(ws.recv(), 15)
                    except _hp_asyncio.TimeoutError: continue
                    if isinstance(message, str):
                        try: _hp_event(state, _hp_json.loads(message), _hp_time.time())
                        except (ValueError, TypeError, AttributeError): pass
        except _hp_asyncio.CancelledError: raise
        except Exception:
            await _hp_asyncio.sleep(5)
        finally: state['connected'] = False


async def huobao_progress(provider, prompt_id, queue, headers):
    """Failure of progress reporting must never fail or resubmit a generation."""
    timestamp = _hp_time.time()
    running = queue.get('queue_running', []); pending = queue.get('queue_pending', [])
    row = next((x for x in running + pending if len(x) > 3 and x[1] == prompt_id), None)
    if row is None: return None
    result = {'observedAt': round(timestamp * 1000), 'phase': 'queued' if row in pending else 'running'}
    if row in pending:
        result.update(queuePosition=pending.index(row) + 1, queueAhead=len(running) + pending.index(row))
    try:
        previous = _hp_watchers.get(provider.id)
        if previous and (previous['prompt'] != prompt_id and row in running or previous['task'].done()):
            previous['task'].cancel(); _hp_watchers.pop(provider.id, None); previous = None
        if row in running:
            client_id = row[3].get('client_id')
            # Huobao's normalizer generates unique canvas-* client IDs. Never
            # replace a human's browser websocket or guess another job's state.
            if isinstance(client_id, str) and client_id.startswith('canvas-'):
                if previous is None:
                    previous = {'prompt':prompt_id, 'classes':{str(k):v.get('class_type') for k,v in row[2].items()},
                                'requested':timestamp, 'samples':[], 'connected':False}
                    previous['task'] = _hp_asyncio.create_task(_hp_watch(previous, provider.base_url, headers, client_id))
                    _hp_watchers[provider.id] = previous
                previous['requested'] = timestamp
                if previous['prompt'] == prompt_id: result.update(_hp_snapshot(previous, timestamp))
        cached = _hp_histories.get(provider.id)
        if cached is None or timestamp - cached[0] > 90:
            async with _hp_httpx.AsyncClient(timeout=5) as client:
                response = await client.get(provider.base_url + '/history', params={'max_items':80}, headers=headers)
                response.raise_for_status(); history = response.json()
                _hp_histories[provider.id] = (timestamp, history)
        else: history = cached[1]
        result.update(_hp_history_estimate(row[2], history))
    except Exception:
        pass
    return result