"""Worker-local, bounded asynchronous inference. No operator-network dependencies."""

import hashlib
import json
import threading
import time
import traceback
from pathlib import Path


class LatestInference:
    def __init__(self, factory, policy, directory: Path, clock=time.monotonic):
        self.factory, self.policy, self.directory, self.clock = factory, policy, directory, clock
        self.condition = threading.Condition()
        self.enabled = self.closed = self.ready = False
        self.ever_ready = False
        self.epoch = self.count = self.dropped = 0
        self.pending = self.result = None
        self.error = None
        self.thread = threading.Thread(target=self._run, name="polygon-inference", daemon=True)

    def start(self):
        self.thread.start()

    def enable(self, enabled):
        with self.condition:
            if enabled != self.enabled:
                self.enabled = enabled
                self.epoch += 1
                self.pending = self.result = None
            self.condition.notify_all()

    def submit(self, rgb, frame_id, captured_at, simulation_time_ns, observation=None):
        with self.condition:
            if not self.enabled or self.closed:
                return
            if self.pending is not None:
                self.dropped += 1
            self.pending = (rgb, frame_id, captured_at, simulation_time_ns, self.epoch, observation)
            self.condition.notify_all()

    def command(self, now, frame_deadline=0.5, command_deadline=0.5):
        with self.condition:
            result = self.result
            if not self.enabled:
                return 0.0, 0.0, "paused", result
            if self.error:
                return 0.0, 0.0, "inference-error", result
            if result is None or now - result["captured_at"] > frame_deadline:
                return 0.0, 0.0, "stale-camera", result
            if now - result["completed_at"] > command_deadline:
                return 0.0, 0.0, "stale-command", result
            decision = result["decision"]
            return decision["speed_mps"], decision["yaw_rate_rps"], "none", result

    def _run(self):
        model = None
        policy_epoch = None
        self.directory.mkdir(parents=True, exist_ok=True)
        try:
            with (self.directory / "decisions.jsonl").open("a", encoding="utf-8") as journal:
                while True:
                    with self.condition:
                        self.condition.wait_for(
                            lambda: self.closed or (self.enabled and self.pending is not None)
                        )
                        if self.closed:
                            return
                        rgb, frame_id, captured_at, sim_ns, epoch, observation = self.pending
                        self.pending = None
                    try:
                        if model is None:
                            model = self.factory()
                            model.ready()
                        started = self.clock()
                        evidence = {}
                        if self.policy is None:
                            if policy_epoch is not None and policy_epoch != epoch:
                                model.reset()
                            policy_epoch = epoch
                            decision, boxes, evidence = model.infer_observation(
                                rgb, observation, frame_id
                            )
                        else:
                            road, boxes = model.infer(rgb)
                        with self.condition:
                            if epoch != self.epoch or not self.enabled:
                                continue
                            if self.policy is not None and policy_epoch != epoch:
                                self.policy.reset()
                                policy_epoch = epoch
                            if self.policy is not None:
                                decision = self.policy.decide(road, boxes).model_dump()
                        completed = self.clock()
                        record = dict(
                            frame_id=frame_id,
                            captured_at=captured_at,
                            simulation_time_ns=sim_ns,
                            completed_at=completed,
                            inference_ms=(completed - started) * 1000,
                            decision=decision,
                        )
                        with self.condition:
                            self.ready, self.ever_ready, self.error = True, True, None
                            if epoch != self.epoch or not self.enabled:
                                continue  # A pause/resume fences every earlier in-flight decision.
                            self.result = record
                            self.count += 1
                        # Archive source RGB on Worker, independently of the video stream.
                        from PIL import Image

                        image_path = self.directory / f"{frame_id:08d}.jpg"
                        Image.fromarray(rgb).save(image_path, "JPEG", quality=90)
                        record = {
                            **record,
                            "image": image_path.name,
                            "image_sha256": hashlib.sha256(image_path.read_bytes()).hexdigest(),
                            "obstacles": boxes,
                            **evidence,
                        }
                        journal.write(json.dumps(record, allow_nan=False) + "\n")
                        journal.flush()
                    except Exception as exc:
                        with (self.directory / "errors.jsonl").open(
                            "a", encoding="utf-8"
                        ) as errors:
                            errors.write(
                                json.dumps(
                                    {
                                        "frame_id": frame_id,
                                        "monotonic": self.clock(),
                                        "error": str(exc),
                                        "type": type(exc).__name__,
                                        "traceback": traceback.format_exc(),
                                    }
                                )
                                + "\n"
                            )
                        with self.condition:
                            self.error, self.ready, self.result = type(exc).__name__, False, None
                        if model is not None:
                            model.close()
                            model = None
                        with self.condition:
                            self.condition.wait(timeout=0.5)
        finally:
            if model is not None:
                model.close()

    def close(self):
        with self.condition:
            self.closed = True
            self.condition.notify_all()
        self.thread.join(timeout=12)