"""Bounded, cancellable persistent JSONL model process; never executes a shell."""
import json, os, selectors, signal, subprocess, threading, time
from robot_bt_coordinator.plan import canonical, strict_json

class InferenceError(RuntimeError):
    def __init__(self,code,message=''):self.code=code;super().__init__(message or code)

class FixtureBackend:
    model_version='fixture-only'
    def __init__(self,raw):self.raw=raw
    def infer(self,request,timeout,cancel=None):
        if cancel is not None and cancel.is_set():raise InferenceError('CANCELED')
        return self.raw(request) if callable(self.raw) else self.raw
    def close(self):pass

class ProcessBackend:
    """One in-flight call; timeout/cancel kills the entire inference process group.
    Restart reloads weights. The local GPU worker has no robot-control authority.
    """
    def __init__(self,argv,model_version):
        if not isinstance(argv,list) or not argv or any(not isinstance(x,str) or not x for x in argv):raise ValueError('explicit argv required')
        if not model_version:raise ValueError('pinned model version required')
        self.argv=argv;self.model_version=model_version;self.lock=threading.Lock();self.process=None
    def close(self):
        p=self.process;self.process=None
        if p:
            if p.poll() is None:
                os.killpg(p.pid,signal.SIGKILL)
            p.wait(timeout=5)
            p.stdin.close();p.stdout.close()
    def infer(self,request,timeout,cancel=None):
        if not 0<timeout<=3600:raise InferenceError('INVALID_TIMEOUT')
        if not self.lock.acquire(blocking=False):raise InferenceError('BUSY')
        deadline=time.monotonic()+timeout
        try:
            if cancel is not None and cancel.is_set():raise InferenceError('CANCELED')
            if self.process is None:
                self.process=subprocess.Popen(self.argv,stdin=subprocess.PIPE,stdout=subprocess.PIPE,stderr=None,start_new_session=True,bufsize=0)
            p=self.process
            payload=(canonical(request)+'\n').encode()
            if len(payload)>262144:raise InferenceError('INPUT_TOO_LARGE')
            # Nonblocking write/read includes process startup in the same deadline.
            os.set_blocking(p.stdin.fileno(),False);os.set_blocking(p.stdout.fileno(),False)
            sent=0;data=b''
            with selectors.DefaultSelector() as sel:
                sel.register(p.stdout,selectors.EVENT_READ);sel.register(p.stdin,selectors.EVENT_WRITE)
                while True:
                    if cancel is not None and cancel.is_set():raise InferenceError('CANCELED')
                    if time.monotonic()>=deadline:raise InferenceError('TIMEOUT')
                    for key,_ in sel.select(min(.05,max(0,deadline-time.monotonic()))):
                        if key.fileobj is p.stdin:
                            sent+=os.write(p.stdin.fileno(),payload[sent:])
                            if sent==len(payload):sel.unregister(p.stdin)
                        else:
                            chunk=os.read(p.stdout.fileno(),65536)
                            if not chunk:raise InferenceError('WORKER_EXITED')
                            data+=chunk
                            if len(data)>262144:raise InferenceError('OUTPUT_TOO_LARGE')
                            if b'\n' in data:
                                raw,extra=data.split(b'\n',1)
                                if extra.strip():raise InferenceError('WORKER_PROTOCOL')
                                msg=strict_json(raw.decode())
                                if not isinstance(msg,dict) or set(msg)!={'raw'} or not isinstance(msg['raw'],str):raise InferenceError('WORKER_PROTOCOL')
                                return msg['raw']
        except InferenceError:
            self.close();raise
        except Exception as ex:
            self.close();raise InferenceError('INFERENCE_FAILED',str(ex)) from ex
        finally:self.lock.release()