#!/usr/bin/env python3
"""Live baseline V4L2/CAN capture + native TensorRT + CMS v9.10 risk."""
from __future__ import annotations
import argparse
from collections import Counter, deque
from contextlib import ExitStack
import copy
import json
import signal
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
import cv2
import numpy as np
ROOT=Path(__file__).resolve().parent
sys.path.insert(0,str(ROOT/'pipeline'))
from native_backend import NativeBackend
from pose_engine import NativePoseEngine
from latest_result import LatestResult, result_to_dict
from door_mask_risk import DoorMaskSequence, DoorRiskEvaluator
from door_zone_risk import DoorZoneConfig, DoorZoneAssigner, build_door_risks
from visualization import render_frame, visualization_skeletons, timing_stats


def select_active_door(mode, snapshot, side='left', timeout_ms=1000):
    if mode!='can': return mode, True
    p=snapshot['positions'];offset=0 if side=='left' else 1
    if snapshot['age_ms']<0 or snapshot['age_ms']>timeout_ms or min(p[offset],p[offset+2])<0:
        return 'none',False
    front,rear=p[offset]==0,p[offset+2]==0
    return ('both' if front and rear else 'front' if front else 'rear' if rear else 'none'),True


def gate_risks(raw, assignments, active, can_valid):
    doors,risk=build_door_risks(raw,assignments,active)
    if not can_valid:
        for door,item in doors.items():
            item.update(risk_type=door.upper()+'_UNKNOWN',risk_level='unknown',risk_score=None,decision='CAN_UNAVAILABLE_OR_STALE')
        risk=dict(risk,risk_type='CAN_UNKNOWN',risk_level='unknown',risk_score=None,decision='CAN_UNAVAILABLE_OR_STALE')
    return doors,risk


class JsonlWriter:
    """Bound the disk usage of continuous camera runs (current + two backups)."""
    def __init__(self,path,max_bytes):
        self.path=path;self.max_bytes=max_bytes;self.size=0
        self.handle=path.open('w',encoding='utf-8',buffering=1)
    def __enter__(self):return self
    def __exit__(self,*args):self.close()
    def close(self):self.handle.close()
    def write(self,text):
        size=len(text.encode('utf-8'))
        if self.max_bytes and self.size and self.size+size>self.max_bytes:
            self.handle.close()
            backup=self.path.with_suffix(self.path.suffix+'.1')
            if backup.exists():backup.replace(self.path.with_suffix(self.path.suffix+'.2'))
            self.path.replace(backup)
            self.handle=self.path.open('w',encoding='utf-8',buffering=1);self.size=0
        self.handle.write(text);self.size+=size


def render_output(frame, record, risk, mask, assignments, status):
    canvas = render_frame(frame, record, risk, mask, assignments)
    cv2.putText(canvas,status,(24,frame.shape[0]-30),cv2.FONT_HERSHEY_SIMPLEX,.8,(0,180,255),2)
    return canvas


class Preview:
    def __init__(self,host,port,results):
        self.jpeg=None;self.results=results;self.lock=threading.Lock()
        self.pending=threading.Condition();self.packet=None;self.stopped=False
        outer=self
        class Handler(BaseHTTPRequestHandler):
            def do_GET(self):
                with outer.lock:
                    if self.path.split('?')[0]=='/frame.jpg': payload=outer.jpeg or b'';mime='image/jpeg'
                    elif self.path=='/result.json':payload=json.dumps(result_to_dict(outer.results.load()) or {},ensure_ascii=False,allow_nan=False).encode();mime='application/json'
                    else:
                        payload=b'''<!doctype html><meta charset="utf-8"><title>Vision anti-pinch 0923</title><body style="background:#101820;color:white;font:18px sans-serif"><h2>Vision anti-pinch 0923 / Live camera</h2><p id="state"></p><img id="cam" style="max-width:100%;max-height:83vh"><script>let img=document.querySelector('#cam'); function next(){setTimeout(()=>img.src='/frame.jpg?t='+Date.now(),150)}img.onload=next;img.onerror=next;next();setInterval(async()=>{let r=await(await fetch('/result.json')).json();document.querySelector('#state').textContent=JSON.stringify({frame:r.frame_index,active_door:r.active_door,risk:r.risk,can:r.can,calibration_verified:r.calibration_verified})},1000)</script>''';mime='text/html; charset=utf-8'
                self.send_response(200 if payload else 503);self.send_header('Content-Type',mime);self.send_header('Cache-Control','no-store');self.send_header('Content-Length',str(len(payload)));self.end_headers()
                try:self.wfile.write(payload)
                except (BrokenPipeError,ConnectionResetError):pass
            def log_message(self,*args):pass
        self.server=ThreadingHTTPServer((host,port),Handler)
        self.thread=threading.Thread(target=self.server.serve_forever,daemon=True);self.thread.start()
        self.renderer=threading.Thread(target=self._render_loop,name='risk-preview',daemon=True)
        self.renderer.start()
    def submit(self, frame, record, risk, mask, assignments, status):
        # All inputs belong to a finished frame; inference never modifies them
        # after submission. Retain only one pending frame, replacing older work.
        with self.pending:
            if not self.stopped:
                self.packet=(frame,record,risk,mask,assignments,status)
                self.pending.notify()

    def _render_loop(self):
        try:
            while True:
                with self.pending:
                    self.pending.wait_for(lambda:self.packet is not None or self.stopped)
                    if self.stopped:return
                    packet=self.packet;self.packet=None
                self.update(render_output(*packet))
        except Exception as exc:
            with self.pending:self.stopped=True;self.packet=None
            with self.lock:self.jpeg=None
            print('Visualization disabled; inference continues: '+str(exc),file=sys.stderr,flush=True)

    def update(self,image):
        w=960;h=round(image.shape[0]*w/image.shape[1]);image=cv2.resize(image,(w,h))
        ok,jpeg=cv2.imencode('.jpg',image,[cv2.IMWRITE_JPEG_QUALITY,80])
        if ok:
            with self.lock:self.jpeg=jpeg.tobytes()
    def close(self):
        with self.pending:self.stopped=True;self.packet=None;self.pending.notify_all()
        self.renderer.join()
        self.server.shutdown();self.server.server_close();self.thread.join()


def parse_args(argv=None):
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--video',type=Path,help='Diagnostic video replay through the same TensorRT/risk chain');p.add_argument('--camera',default='/dev/video5');p.add_argument('--width',type=int,default=1920);p.add_argument('--height',type=int,default=1536);p.add_argument('--fps',type=int,default=30)
    p.add_argument('--engines',type=Path,default=ROOT/'orin/models/engines' if (ROOT/'orin').is_dir() else ROOT/'models/engines')
    p.add_argument('--constant-door-mask',type=Path,default=ROOT/'annotations/constant_door_mask.json')
    p.add_argument('--door-zone-config',type=Path,default=ROOT/'annotations/door_zone_multifeature.json')
    p.add_argument('--active-door',choices=['can','front','rear','both','none'],default='can')
    p.add_argument('--can-side',choices=['left','right'],default='left');p.add_argument('--can-timeout-ms',type=float,default=1000);p.add_argument('--no-can',action='store_true')
    p.add_argument('--max-people',type=int,default=3);p.add_argument('--det-conf',type=float,default=.30)
    p.add_argument('--max-frames',type=int,default=0);p.add_argument('--duration',type=float,default=0)
    p.add_argument('--jsonl-max-mb',type=float,default=32,help='Rotate each JSONL at this size; keep two backups; 0 disables rotation');p.add_argument('--output',type=Path,default=ROOT/'outputs'/time.strftime('%Y%m%d-%H%M%S'))
    p.add_argument('--visualize',action='store_true',help='Enable optional HTTP visualization; disabled by default')
    p.add_argument('--save-results',action='store_true',help='Save per-frame JSONL and periodic latest.json')
    p.add_argument('--save-snapshots',action='store_true',help='Render and save latest.jpg every 10 seconds')
    p.add_argument('--preview-port',type=int,default=8765);p.add_argument('--preview-host',default='0.0.0.0');p.add_argument('--preview-every',type=int,default=3)
    p.add_argument('--write-video',action='store_true');p.add_argument('--save-pose-records',action='store_true');p.add_argument('--calibration-verified',action='store_true')
    a=p.parse_args(argv)
    if not 1<=a.max_people<=16 or not 0<a.det_conf<1 or min(a.width,a.height,a.fps,a.preview_every)<=0 or a.can_timeout_ms<=0 or a.jsonl_max_mb<0:p.error('Invalid camera/inference configuration')
    if a.no_can and a.active_door=='can':p.error('--no-can requires explicit --active-door front/rear/both/none')
    if not 0<=a.preview_port<=65535 or (a.visualize and a.preview_port==0):p.error('--visualize requires a port between 1 and 65535')
    if a.max_frames<0 or a.duration<0:p.error('Frame count and duration cannot be negative')
    a.det_roi_y_max_ratio=1.;return a


def run_pipeline(a, results, stop, summary_sink=None):
    """Produce results independently of optional visualization and disk outputs."""
    save_any = a.save_results or a.save_snapshots or a.save_pose_records or a.write_video
    if save_any:
        a.output.mkdir(parents=True,exist_ok=True)
    cv2.setNumThreads(2)
    native=None;preview=None;writer=None;video=None;count=0;dropped=0;levels=Counter();latencies=deque(maxlen=3000);pose_ms=deque(maxlen=3000);risk_ms=deque(maxlen=3000);ages=deque(maxlen=3000)
    started=time.monotonic();run_started=None;error=None
    try:
        if a.video:
            video=cv2.VideoCapture(str(a.video))
            if not video.isOpened():raise RuntimeError('Cannot open video: '+str(a.video))
            a.width=int(video.get(cv2.CAP_PROP_FRAME_WIDTH));a.height=int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        evaluator=DoorRiskEvaluator(DoorMaskSequence(constant_mask=a.constant_door_mask,require_every_frame=True))
        assigner=DoorZoneAssigner(DoorZoneConfig.load(a.door_zone_config))
        # Fail before acquiring camera if calibration mask is invalid.
        initial=evaluator.sequence.load(0,(a.height,a.width))
        if initial.status!='loaded' or initial.mask is None or not np.any(initial.mask):raise ValueError('Door mask has no usable pixels')
        native=NativeBackend(a);pose=NativePoseEngine(a,native)
        if a.visualize:
            try:preview=Preview(a.preview_host,a.preview_port,results)
            except Exception as exc:
                print('Visualization unavailable; inference continues: '+str(exc),file=sys.stderr,flush=True)
        if not video:native.open_camera()
        run_started=time.monotonic();last_report=run_started;last_snapshot=0
        if a.write_video:
            writer=cv2.VideoWriter(str(a.output/'rtmw_risk_overlay.mp4'),cv2.VideoWriter_fourcc(*'mp4v'),a.fps,(a.width,a.height))
            if not writer.isOpened():raise RuntimeError('Video writer could not open')
        with ExitStack() as output_stack:
            out=output_stack.enter_context(JsonlWriter(a.output/'causal_results.jsonl',int(a.jsonl_max_mb*1024*1024))) if a.save_results else None
            pose_out=JsonlWriter(a.output/'pose_records.jsonl',int(a.jsonl_max_mb*1024*1024)) if a.save_pose_records else None
            try:
                while not stop.is_set():
                    if (a.max_frames and count>=a.max_frames) or (a.duration and time.monotonic()-run_started>=a.duration):break
                    if video:
                        ok,frame=video.read()
                        if not ok:break
                        capture={'sequence':count,'capture_monotonic_ms':time.monotonic()*1000,'dropped_before_inference':0,'latency_clock':'video_decode_complete'}
                    else:frame,capture=native.read()
                    if frame is None:break
                    tick=time.monotonic();record=pose.infer(frame,count);pose_ms.append((time.monotonic()-tick)*1000)
                    rt=time.monotonic();snapshot=native.can_snapshot();active,valid=select_active_door(a.active_door,snapshot,a.can_side,a.can_timeout_ms)
                    assignments=assigner.update(count,record,(a.height,a.width));raw,mask=evaluator.evaluate(count,record,(a.height,a.width))
                    doors,risk=gate_risks(raw,assignments,active,valid);risk_ms.append((time.monotonic()-rt)*1000)
                    result={'frame_index':count,'track_ids':record['track_ids'],'boxes':record['boxes'],'skeletons':visualization_skeletons(record),'active_door':active,'door_assignments':assignments,'door_risks':doors,'risk':risk,'ungated_risk':raw,'capture':capture,'can':dict(snapshot,valid=valid,source='hardware' if a.active_door=='can' else 'simulation',side=a.can_side),'calibration_verified':a.calibration_verified}
                    result['capture_to_result_ms']=round(time.monotonic()*1000-capture['capture_monotonic_ms'],3)
                    results.publish(result)
                    if out:out.write(json.dumps(result,ensure_ascii=False,allow_nan=False)+'\n')
                    if pose_out:pose_out.write(json.dumps(record,ensure_ascii=False,allow_nan=False)+'\n')
                    now=time.monotonic()
                    snapshot_due=now-last_snapshot>=10
                    if a.save_results and snapshot_due:
                        (a.output/'latest.json').write_text(json.dumps(result,ensure_ascii=False,indent=2))
                    status=('CAN LIVE' if a.active_door=='can' else 'CAN SIMULATION: '+active)+' | '+('CALIBRATED' if a.calibration_verified else 'CALIBRATION UNVERIFIED')
                    if preview and count%a.preview_every==0:
                        preview.submit(frame,record,risk,mask,assignments,status)
                    if writer or (a.save_snapshots and snapshot_due):
                        canvas=render_output(frame,record,risk,mask,assignments,status)
                        if writer:writer.write(canvas)
                        if a.save_snapshots and snapshot_due:
                            cv2.imwrite(str(a.output/'latest.jpg'),canvas)
                    if snapshot_due:last_snapshot=now
                    count+=1;dropped+=capture['dropped_before_inference'];levels[risk['risk_level']]+=1;ages.append(result['capture_to_result_ms']);latencies.append((time.monotonic()-tick)*1000)
                    if now-last_report>=5:
                        print(json.dumps({'frames':count,'fps':round(count/(now-run_started),2),'people':len(record['boxes']),'risk':risk['risk_type'],'latency_ms':round(float(np.mean(latencies)),1),'capture_age_ms':round(float(np.mean(ages)),1),'can_valid':valid}),flush=True);last_report=now
            finally:
                if pose_out:pose_out.close()
    except Exception as exc:
        error=str(exc);raise
    finally:
        ended=time.monotonic()
        if video:video.release()
        if native:native.close()
        if writer:writer.release()
        if preview:preview.close()
        elapsed=ended-(run_started or started)
        summary={'mode':'baseline_v4l2_can_native_trt_cms_v9_10','camera':a.camera if not a.video else None,'video':str(a.video) if a.video else None,'frames':count,'elapsed_seconds':round(elapsed,3),'throughput_fps':round(count/max(elapsed,.001),3),'dropped_latest_frames':dropped,'risk_levels':dict(levels),'active_door_mode':a.active_door,'can_side':a.can_side,'calibration_verified':a.calibration_verified,'visualize_requested':a.visualize,'save_results':a.save_results,'save_snapshots':a.save_snapshots,'latency_window_frames':3000,'latency_ms':{'total':timing_stats(list(latencies)),'pose':timing_stats(list(pose_ms)),'risk':timing_stats(list(risk_ms)),'capture_to_result':timing_stats(list(ages))},'error':error}
        if summary_sink is not None:summary_sink.update(summary)
        if save_any:(a.output/'summary.json').write_text(json.dumps(summary,ensure_ascii=False,indent=2))
        print(json.dumps(summary,ensure_ascii=False),flush=True)

class CameraPipeline:
    """One lifecycle per instance, with the 0920 start/latest/stop/join pattern.

    The worker owns inference. Downstream threads read the latest immutable
    snapshot directly, without JSON files, HTTP, callbacks or a result queue.
    """
    def __init__(self, args=None):
        self.args = copy.deepcopy(args if args is not None else parse_args([]))
        self._results = LatestResult()
        self._stop = threading.Event()
        self._lifecycle_lock = threading.Lock()
        self._thread = None
        self._error = None
        self.summary = {}

    @property
    def latest_result(self):
        """The latest complete result variable, or None before the first frame."""
        return self._results.load()

    def latest(self):
        return self.latest_result

    def wait_for_result(self, after_frame_index=-1, timeout=None):
        return self._results.wait(after_frame_index, timeout)

    def start(self):
        with self._lifecycle_lock:
            if self._thread is not None:
                raise RuntimeError('CameraPipeline can only be started once')
            self._thread = threading.Thread(target=self._run, name='vision-risk')
            self._thread.start()
        return self

    def _run(self):
        try:
            run_pipeline(self.args, self._results, self._stop, self.summary)
        except BaseException as exc:
            self._error = exc
        finally:
            self._results.close()

    def running(self):
        return self._thread is not None and self._thread.is_alive()

    def request_stop(self):
        self._stop.set()

    def join(self, timeout=None):
        if self._thread is not None:
            self._thread.join(timeout)
            if not self._thread.is_alive() and self._error is not None:
                raise self._error

    def __enter__(self):
        return self.start()

    def __exit__(self, *args):
        self.request_stop()
        self.join()


def main():
    pipeline=CameraPipeline(parse_args())
    for sig in (signal.SIGINT,signal.SIGTERM):
        signal.signal(sig,lambda *_:pipeline.request_stop())
    pipeline.start()
    # Periodic joins allow Python to deliver stop signals on the main thread.
    while pipeline.running():
        pipeline.join(.2)
    pipeline.join()


if __name__=='__main__':main()