#!/usr/bin/env python3
"""WI-252 Edge workflow client. Default is read-only; --run enables calibration.

Live: python3 scripts/scale_edge_sim.py --host BROKER --prefix mes/gateway/UID --run --capacity 50
Host-only: python3 scripts/scale_edge_sim.py --simulate --capacity 50 --increments 5,20,20
The simulation uses a real local MQTT broker and the production C++ gateway adapter.
"""
import argparse
from decimal import Decimal, InvalidOperation
import json
import os
from pathlib import Path
import queue
import socket
import subprocess
import tempfile
import threading
import time

import paho.mqtt.client as mqtt

ROOT = Path(__file__).resolve().parents[1]


class Rejected(RuntimeError):
    def __init__(self, reason): self.reason=reason;super().__init__('Rejected: '+reason)


class Link:
    def __init__(self, host, port, prefix, tls=False, log=None, weight_topic=None):
        self.prefix, self.log = prefix.rstrip('/'), log
        self.weight_topic = weight_topic or self.prefix + '/scale/weight'
        self.inbox = queue.Queue()
        self.connected = threading.Event()
        self.client = mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
        if tls:
            self.client.tls_set()
        if os.environ.get('MQTT_USER'):
            self.client.username_pw_set(os.environ['MQTT_USER'], os.environ.get('MQTT_PASSWORD'))
        def connect(client, userdata, flags, code, properties):
            if code == 0:
                client.subscribe(self.weight_topic, qos=1)
        def subscribed(client, userdata, mid, codes, properties):
            if all(code.value < 128 for code in codes): self.connected.set()
        def received(client, userdata, message):
            if message.retain: return  # retained historic status is not evidence
            try: payload = json.loads(message.payload)
            except (ValueError, UnicodeError): return
            if message.topic != self.weight_topic or not isinstance(payload, dict): return
            calibration = payload.get('calibration')
            if not isinstance(calibration, dict): return
            received_at = time.monotonic()
            data = dict(calibration, timestamp=payload.get('timestamp'))
            self.inbox.put((received_at, 'status', data))
            if data.get('sampleEvent'):
                self.inbox.put((received_at, 'sample', data))
            item = (received_at, 'weight', payload)
            if self.log:
                self.log.write(json.dumps({'receivedAt': item[0], 'topic': message.topic, 'payload': payload}) + '\n')
                self.log.flush()
        self.client.on_connect = connect
        self.client.on_subscribe = subscribed
        self.client.on_message = received
        self.client.on_disconnect = lambda *args: self.connected.clear()
        self.client.connect(host, port, 30)
        self.client.loop_start()
        if not self.connected.wait(10): raise TimeoutError('MQTT subscription not confirmed')

    def send(self, command):
        if not self.connected.is_set(): raise ConnectionError('MQTT disconnected; stop and query after reconnect')
        result = self.client.publish(self.prefix + '/cmd/calibration', json.dumps(command), qos=1, retain=False)
        if result.rc != mqtt.MQTT_ERR_SUCCESS: raise ConnectionError('publish rejected')

    def close(self):
        self.client.disconnect()
        self.client.loop_stop()


class Edge:
    def __init__(self, link, slave=3, stable_timeout=300):
        self.link, self.slave = link, slave
        self.stable_timeout = stable_timeout
        self.boot, self.sid, self.seq, self.steps = None, 0, 0, 3
        self.revision, self.last_sample = 0, 0
        self.last_ping = 0
        self.pending = None
        self.stop = threading.Event()
        self.heartbeat = None

    def receive(self, deadline):
        while time.monotonic() < deadline:
            if self.sid and time.monotonic() - self.last_ping >= 10:
                self.link.send({'op': 'ping', 'boot': self.boot, 'sid': self.sid})
                self.last_ping = time.monotonic()
            try: received, kind, data = self.link.inbox.get(timeout=min(.2, max(.001, deadline-time.monotonic())))
            except queue.Empty: continue
            if self.boot and data.get('boot') != self.boot:
                raise RuntimeError('Gateway rebooted: previous result unknown; do not resume automatically')
            return received, kind, data
        raise TimeoutError('No matching gateway response; result may be unknown')

    def query(self):
        self.link.send({'op': 'query'})
        deadline = time.monotonic() + 10
        while True:
            _, kind, data = self.receive(deadline)
            if kind == 'status':
                self.boot = data['boot']
                if not self.boot: raise RuntimeError('Gateway boot identity unavailable')
                return data

    def request(self, op, **fields):
        if self.pending: raise RuntimeError('Unresolved request: query before further commands')
        command = dict(op=op, boot=self.boot, sid=self.sid, seq=self.seq+1, slave=self.slave, steps=self.steps, **fields)
        self.pending = command
        # Retry only the identical request, once. Gateway owns idempotency.
        for attempt in range(2):
            self.link.send(command)
            deadline = time.monotonic()+12
            while time.monotonic() < deadline:
                try: _, kind, data = self.receive(deadline)
                except TimeoutError: break
                if kind != 'status' or data.get('requestSid') != self.sid or data.get('requestSeq') != command['seq']: continue
                if data.get('result') == 'accepted': continue
                if data.get('result') == 'applied':
                    expected={'begin':'await_zero','zero':'collecting','point':'collecting','finish':'complete','cancel':'partial'}[op]
                    if data.get('state')!=expected:
                        raise RuntimeError('Historical ACK received but session state changed; stop and query')
                    self.seq = command['seq']; self.revision = data['revision']; self.pending = None
                    return data
                if data.get('result') == 'rejected':
                    self.pending = None
                    raise Rejected(data.get('reason','unknown'))
                if data.get('result') == 'unknown':
                    raise RuntimeError('Controller outcome unknown; stop. Query and explicitly cancel/restart; no automatic replay')
        raise TimeoutError('ACK absent after identical retry. Stop; do not advance to next point')

    def wait_stable(self, timeout=None):
        if timeout is None: timeout = self.stable_timeout
        deadline = time.monotonic()+timeout
        seen = set(); last_received = None; proof = None
        # Purge queued pre-confirmation samples. Gateway also validates sample/revision.
        while True:
            try: self.link.inbox.get_nowait()
            except queue.Empty: break
        while True:
            try: received, kind, data = self.receive(deadline)
            except TimeoutError:
                raise TimeoutError(f'No fresh stable samples within {timeout}s; this step was not sent. Stop and query') from None
            if data.get('sid')==self.sid and data.get('state') in ('unknown','partial','recovery_required'):
                raise RuntimeError('Calibration interrupted: '+data['state']+'; stop and query')
            if kind != 'sample' or data.get('sid') != self.sid or data.get('revision') != self.revision: continue
            if time.monotonic()-received > 2.5: continue
            sample = data.get('sample',0)
            reading = data.get('reading',0) if data.get('v',1)>=2 else sample
            if reading <= self.last_sample: continue
            self.last_sample = reading
            if data.get('v',1)>=2 and sample!=proof:
                seen.clear(); last_received=None; proof=sample
            if not data.get('valid') or not data.get('stable'):
                seen.clear(); last_received=None; continue
            if last_received is not None and received-last_received > 2.5: seen.clear()
            seen.add(reading); last_received=received
            print(f"sample={sample} weight={Decimal(data['weightRaw'])/(10**data['decimals'])} kg state={data['state']}", flush=True)
            if len(seen)>=3 and data.get('ready'):
                return {'sample': sample, 'revision': self.revision}

    def run(self, capacity, increments, fixture=None, tolerance="0.1"):
        snapshot=self.query()
        if snapshot['state'] in ('await_zero','collecting','unknown'):
            raise RuntimeError('Existing session found. Query/cancel explicitly; do not take it over')
        self.sid=snapshot['sid']+1; self.steps=len(increments)
        print(self.request('begin', unit='kg', model='tda08b', capacityKg=str(capacity),toleranceKg=str(tolerance)))
        # Keep the session alive while a person is placing weights at the console.
        # A closed client/network loss stops renewals; gateway enforces the lease.
        def renew():
            while not self.stop.wait(10):
                if not self.sid: return
                try: self.link.send({'op':'ping','boot':self.boot,'sid':self.sid})
                except ConnectionError: return
        self.heartbeat=threading.Thread(target=renew,daemon=True);self.heartbeat.start()
        if fixture: fixture(0)
        else: input('清空秤盤，確認面板單位 kg，再按 Enter：')
        print(self.stable_request('zero'))
        total=Decimal(0)
        for increment in increments:
            total+=increment
            if fixture: fixture(total)
            else: input(f'再放 {increment} kg；秤盤總量應為 {total} kg。放穩後按 Enter：')
            result=self.stable_request('point',totalKg=str(total))
            print(f"第 {result['step']} 段 ACK：總重量 {total} kg", flush=True)
        print(self.stable_request('finish'))
        self.sid=0

    def stable_request(self, op, **fields):
        for attempt in range(3):
            try: return self.request(op,**fields,**self.wait_stable())
            except Rejected as error:
                if error.reason!='not_fresh_stable' or attempt==2: raise
                # This rejection guarantees no writes. Keep the SAME total and sequence.

    def close(self):
        self.stop.set()
        if self.heartbeat: self.heartbeat.join(1)

    def cancel(self, session):
        snapshot=self.query()
        if snapshot['sid']!=session: raise RuntimeError('Session changed; refusing to cancel another session')
        self.sid=session;self.seq=snapshot['seq'];self.steps=snapshot['steps']
        print(self.request('cancel'))
        self.sid=0


def build_host(output):
    subprocess.run(['c++','-std=c++17','-I','test/test_measurement_polling','-I','src',
                    '-I','.pio/libdeps/opta/ArduinoJson/src','scripts/tests/scale_gateway_host.cpp','-o',str(output)], cwd=ROOT, check=True)


class HostBridge:
    """Test-only: real MQTT -> production C++ adapter -> RTU transport double."""
    def __init__(self, binary, port):
        self.process=subprocess.Popen([str(binary)], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, bufsize=1)
        self.inputs=queue.Queue(); self.stop=threading.Event(); self.ready=threading.Event()
        self.client=mqtt.Client(mqtt.CallbackAPIVersion.VERSION2)
        self.client.on_connect=lambda client,*args: client.subscribe('test/gateway/cmd/calibration',qos=1)
        self.client.on_subscribe=lambda *args: self.ready.set()
        self.client.on_message=lambda client,userdata,msg: self.inputs.put({'command':json.loads(msg.payload)})
        self.client.connect('127.0.0.1',port,30);self.client.loop_start()
        if not self.ready.wait(5): raise TimeoutError('bridge subscribe failed')
        self.thread=threading.Thread(target=self.work,daemon=True);self.thread.start()

    def work(self):
        while not self.stop.is_set():
            try: item=self.inputs.get(timeout=.1)
            except queue.Empty: item={}
            item['advanceMs']=100  # accelerated controlled host clock, not hardware timing evidence
            self.process.stdin.write(json.dumps(item)+'\n');self.process.stdin.flush()
            while True:
                line=self.process.stdout.readline()
                if not line: self.stop.set();return
                event=json.loads(line)
                if event.get('done'): break
                if 'topic' in event: self.client.publish(event['topic'],json.dumps(event['payload']),qos=1)

    def mass(self, kg): self.inputs.put({'weightRaw':int(kg*100)})
    def close(self):
        self.stop.set(); self.thread.join(3);self.client.disconnect();self.client.loop_stop()
        self.process.terminate();self.process.wait(timeout=3)


def main():
    parser=argparse.ArgumentParser(description=__doc__,formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--host',default='127.0.0.1');parser.add_argument('--port',type=int,default=1883)
    parser.add_argument('--prefix');parser.add_argument('--tls',action='store_true')
    parser.add_argument('--weight-topic',help='Existing weight topic; defaults to PREFIX/scale/weight; sole subscription')
    parser.add_argument('--run',action='store_true');parser.add_argument('--simulate',action='store_true')
    parser.add_argument('--cancel-session',type=int,help='Explicitly cancel this session; applied points are NOT rolled back')
    parser.add_argument('--capacity',type=Decimal);parser.add_argument('--increments',default='5,20,20')
    parser.add_argument('--tolerance',default='0.1',help='v2 stable peak-to-peak band in kg; default 0.1 (not +/-0.1)')
    parser.add_argument('--slave',type=int,default=3);parser.add_argument('--log',type=Path)
    parser.add_argument('--stable-timeout',type=int,default=300,help='Seconds to wait for stable samples after operator confirmation (default: 300)')
    args=parser.parse_args()
    if args.stable_timeout<=0: parser.error('--stable-timeout must be positive')
    try: increments=[Decimal(v) for v in args.increments.split(',')]
    except InvalidOperation: parser.error('increments must be decimal kg')
    if args.run or args.simulate:
        if len(increments) not in (3,5) or any(not v.is_finite() or v<=0 for v in increments): parser.error('require 3 or 5 positive increments')
        if args.capacity is None or not args.capacity.is_finite() or args.capacity<=0 or sum(increments)>args.capacity: parser.error('capacity must cover cumulative known mass')
    if not args.simulate and not args.prefix: parser.error('--prefix is required')
    if args.cancel_session is not None and (args.run or args.simulate): parser.error('cancel must be a separate operation')
    bridge=broker=link=log=edge=None
    try:
        with tempfile.TemporaryDirectory(prefix='wi252-') as tmp:
            if args.simulate:
                binary=Path(tmp)/'gateway';build_host(binary)
                with socket.socket() as sock: sock.bind(('127.0.0.1',0));args.port=sock.getsockname()[1]
                config=Path(tmp)/'mosquitto.conf';config.write_text(f'listener {args.port} 127.0.0.1\nallow_anonymous true\npersistence false\n')
                broker=subprocess.Popen(['mosquitto','-c',str(config)],stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL)
                for _ in range(50):
                    try:
                        with socket.create_connection(('127.0.0.1',args.port),timeout=.1): break
                    except OSError: time.sleep(.1)
                bridge=HostBridge(binary,args.port);args.prefix='test/gateway';args.host='127.0.0.1';args.tls=False
            if args.log: log=args.log.open('w')
            link=Link(args.host,args.port,args.prefix,args.tls,log,args.weight_topic)
            edge=Edge(link,args.slave,args.stable_timeout)
            if args.cancel_session is not None: edge.cancel(args.cancel_session)
            elif args.run or args.simulate: edge.run(args.capacity,increments,bridge.mass if bridge else None,args.tolerance)
            else: print(json.dumps(edge.query(),ensure_ascii=False,indent=2))
    finally:
        if edge: edge.close()
        if link: link.close()
        if bridge: bridge.close()
        if broker: broker.terminate();broker.wait(timeout=3)
        if log: log.close()


if __name__=='__main__': main()