"""Local colour-marker preview and deliberately armed camera servo test."""
from collections import deque
from io import BytesIO
from pathlib import Path
import base64
import json
import math
import os
import re
import secrets
import tempfile
import threading
import time

from PIL import Image, ImageChops

FOLLOW_GAIN = 600  # µs/s per unit of image offset (= old behaviour at 300 µs/s)
DEFAULT = dict(hue=85, tolerance=18, follow_speed=400, search_enabled=False, pan_min=1400, pan_max=1600,
               tilt_min=1400, tilt_max=1600, pan_reverse=False, tilt_reverse=False,
               pan_setup_min=1100, pan_setup_max=1900,
               tilt_setup_min=1100, tilt_setup_max=1900)


def validate(value):
    if not isinstance(value, dict):
        raise ValueError('Ungültige Nachführungseinstellungen')
    # Preserve the previous speed when loading existing configurations.
    value = dict(value)
    value.setdefault('follow_speed', 80)
    value.setdefault('search_enabled', False)
    # Old saved configurations retain their tested setup envelope.
    setup_keys = {a+"_setup_"+b for a in ("pan", "tilt") for b in ("min", "max")}
    if isinstance(value, dict) and set(value) == set(DEFAULT)-setup_keys:
        value = {**{k: DEFAULT[k] for k in setup_keys}, **value}
    if set(value) != set(DEFAULT):
        raise ValueError('Unvollständige Nachführungseinstellungen')
    result = {}
    for key, default in DEFAULT.items():
        v = value[key]
        if isinstance(default, bool):
            if type(v) is not bool:
                raise ValueError(f'{key}: Ja/Nein erwartet')
        else:
            if isinstance(v, bool) or not isinstance(v, (float, int)) or not math.isfinite(v):
                raise ValueError(f'{key}: gültige Zahl erwartet')
            v = int(v)
        result[key] = v
    if not 0 <= result['hue'] <= 255 or not 5 <= result['tolerance'] <= 35:
        raise ValueError('Ungültige Farbe oder Farbtoleranz')
    if not 40 <= result['follow_speed'] <= 800:
        raise ValueError('Nachführtempo muss zwischen 40 und 800 µs/s liegen')
    for axis in ('pan', 'tilt'):
        if not 500 <= result[axis+'_setup_min'] <= result[axis+'_min'] <= result[axis+'_max'] <= result[axis+'_setup_max'] <= 2500:
            raise ValueError('Einrichtbereich: 500–2500 µs; Minimum muss kleiner oder gleich Maximum sein; Bewegungsgrenzen müssen im Einrichtbereich liegen')
    return result


def midpoints(config):
    """Stored follow endpoints determine each start, rounded to whole µs."""
    return {axis: (config[axis+'_min'] + config[axis+'_max'] + 1)//2
            for axis in ('pan', 'tilt')}


def locate(jpeg, config):
    """Largest unambiguous saturated blob, circular hue distance (also red).

    Strict hue/saturation masks run first so a strongly coloured marker wins
    over similar but paler surroundings (e.g. bluish paper); weaker masks only
    apply when the strict ones find no unambiguous blob.
    """
    image = Image.open(BytesIO(jpeg))
    image.draft('RGB', (160, 120))  # JPEG decodes directly at reduced size
    image = image.convert('RGB')
    if image.size != (160, 120):
        image = image.resize((160, 120))
    planes = image.convert('HSV').split()
    tolerance = config['tolerance']
    for factor, min_saturation in ((.5, 140), (.75, 120), (1, 100)):
        result = _find_blob(planes, config['hue'], max(5, tolerance*factor), min_saturation)
        if result:
            return result
    return None


def _find_blob(planes, centre, tolerance, min_saturation):
    hue, saturation, value = planes
    hue_table = [255 if min(abs(h-centre), 256-abs(h-centre)) <= tolerance else 0 for h in range(256)]
    mask_image = ImageChops.darker(
        ImageChops.darker(hue.point(hue_table), saturation.point([255 if i >= min_saturation else 0 for i in range(256)])),
        value.point([255 if i >= 60 else 0 for i in range(256)]))
    mask = {m.start() for m in re.finditer(b'\xff', mask_image.tobytes())}
    blobs = []
    while mask:
        seed = mask.pop()
        queue, blob = deque([seed]), [seed]
        while queue:
            i = queue.popleft()
            x, y = i % 160, i // 160
            for n in ((i-1 if x else -1), (i+1 if x < 159 else -1),
                      (i-160 if y else -1), (i+160 if y < 119 else -1)):
                if n in mask:
                    mask.remove(n); queue.append(n); blob.append(n)
        if len(blob) >= 12:
            blobs.append(blob)
    blobs.sort(key=len, reverse=True)
    if not blobs or len(blobs[0]) > 160*120*.35:
        return None
    if len(blobs) > 1 and len(blobs[1]) >= len(blobs[0])*.6:
        return None
    blob = blobs[0]
    return {'x': sum(i % 160 for i in blob)/len(blob)/160,
            'y': sum(i // 160 for i in blob)/len(blob)/120,
            'radius': math.sqrt(len(blob)/math.pi)/160}


class PCA9685:
    """Dedicated camera board, bus 1/address 0x40, channels 0 and 1 only.

    Register protocol: https://www.nxp.com/docs/en/data-sheet/PCA9685.pdf
    No I2C access until an explicit hardware start. No bus scanning/reset broadcast.
    """
    def __init__(self, bus_factory=None):
        if bus_factory is None:
            from smbus2 import SMBus
            bus_factory = SMBus
        self.bus = bus_factory(1)
        try:
            self.off()
            # Sleep, internal oscillator, auto-increment. 25 MHz / (4096*122).
            self.bus.write_byte_data(0x40, 0x00, 0x30)
            self.bus.write_byte_data(0x40, 0xFE, 121)
            self.bus.write_byte_data(0x40, 0x01, 0x04)  # non-inverted totem-pole
            self.bus.write_byte_data(0x40, 0x00, 0x20)
            time.sleep(.005)
        except Exception:
            self.close()
            raise

    def move(self, pan, tilt):
        for channel, pulse in enumerate((pan, tilt)):
            if not math.isfinite(pulse) or not 500 <= pulse <= 2500:
                raise ValueError('Puls außerhalb der absoluten Testgrenze')
            count = round(pulse * 25 / 122)  # microseconds -> oscillator ticks / prescale
            self.bus.write_i2c_block_data(0x40, 0x06+4*channel,
                                         [0, 0, count & 255, count >> 8])

    def off(self):
        failure = None
        # Do not use a block write before auto-increment has been initialized.
        for channel in (0, 1):
            try:
                self.bus.write_byte_data(0x40, 0x09+channel*4, 0x10)
            except Exception as exc:
                failure = exc
        if failure:
            raise failure

    def close(self):
        try:
            self.off()
        finally:
            self.bus.close()


class FollowService:
    def __init__(self, camera, directory, driver_factory=PCA9685, clock=time.monotonic):
        self.camera, self.clock, self.factory = camera, clock, driver_factory
        self.lock = threading.RLock()
        self.path = Path(directory)/'follow.json'
        self.config = dict(DEFAULT)
        self.error = None
        self.config_fault = False
        try:
            if self.path.exists():
                self.config = validate(json.loads(self.path.read_text()))
        except (OSError, ValueError, TypeError):
            self.error = 'Nachführungseinstellungen beschädigt; neu speichern'
            self.config_fault = True
        self.active = False
        self.operation = 'follow'
        self.jog_direction = None
        self.jog_until = 0.
        self.jog_sequence = -1
        self.center_target = None
        self.driver = None
        self.setup_origin = midpoints(self.config)
        self.pan, self.tilt = self.setup_origin['pan'], self.setup_origin['tilt']
        self.last_ping = self.clock()
        self.last_result = self.last_seen = None
        self.result_times = deque(maxlen=45)
        self.analysis_ms = None
        self._reset_search()
        self.number = 0
        self.target = None
        self.jpeg = None
        self.sample = None
        self.sample_token = None
        self.session = None
        self.revision = 0
        self.thread = None
        self.quit = threading.Event()

    def _reset_search(self):
        self.search_started = None
        self.search_path = deque()
        self.last_target = None
        self.reacquire_count = 0
        self.waiting_for_target = False

    def _wait_for_target(self):
        # Keep the selected colour and the armed session, but issue no motion.
        # Existing pulses hold position; halt/watchdogs still release the driver.
        self.search_started = None
        self.search_path.clear()
        self.reacquire_count = 0
        self.waiting_for_target = True

    def _search_step(self, now, dt):
        """Finite sweep of stored follow limits; never use the setup envelope."""
        if self.search_started is None:
            self.search_started = now
            # First look towards the last image position, respecting reversal.
            first = {}
            other = {}
            for axis, coordinate in (('pan', 'x'), ('tilt', 'y')):
                sign = 1 if self.last_target[coordinate] >= .5 else -1
                if self.config[axis+'_reverse']: sign *= -1
                first[axis] = self.config[axis+('_max' if sign > 0 else '_min')]
                other[axis] = self.config[axis+('_min' if sign > 0 else '_max')]
            self.search_path = deque([
                (first['pan'], self.tilt), (other['pan'], self.tilt),
                (other['pan'], first['tilt']), (first['pan'], first['tilt']),
                (first['pan'], other['tilt']), (other['pan'], other['tilt'])])
        if now-self.search_started >= 60:
            self._wait_for_target()
            return
        while self.search_path and (self.pan, self.tilt) == self.search_path[0]:
            self.search_path.popleft()
        if not self.search_path:
            self._wait_for_target()
            return
        step = min(200, self.config['follow_speed'])*dt
        for axis, goal in zip(('pan', 'tilt'), self.search_path[0]):
            position = getattr(self, axis)
            value = position+max(-step, min(step, goal-position))
            setattr(self, axis, max(self.config[axis+'_min'], min(self.config[axis+'_max'], value)))
        self.driver.move(self.pan, self.tilt)

    def _save(self, value):
        config = validate(value)
        with self.lock:
            self.path.parent.mkdir(parents=True, exist_ok=True)
            fd, name = tempfile.mkstemp(dir=self.path.parent, prefix='.follow-')
            try:
                with os.fdopen(fd, 'w') as f:
                    json.dump(config, f)
                os.replace(name, self.path)
            finally:
                if os.path.exists(name): os.unlink(name)
            self.config, self.error, self.config_fault = config, None, False
            self.revision += 1
            return dict(config)

    def configure(self, value):
        with self.lock:
            if self.active:
                raise ValueError('Vor Einstellungsänderungen Nachführung stoppen')
            return self._save(value)

    def _direction(self, direction):
        choices = {'left': ('pan', -1), 'right': ('pan', 1),
                   'up': ('tilt', -1), 'down': ('tilt', 1)}
        if direction not in choices:
            raise ValueError('Ungültige Richtung')
        axis, sign = choices[direction]
        return axis, -sign if self.config[axis+'_reverse'] else sign

    def _require_setup(self, session):
        self.watchdog()
        if not self.active or self.operation != 'setup' or session != self.session:
            raise ValueError('Einrichtsitzung beendet oder nicht freigegeben')

    def jog(self, session, direction, sequence):
        with self.lock:
            self._require_setup(session)
            if type(sequence) is not int or not 0 <= sequence <= 1_000_000_000:
                raise ValueError('Ungültige Tastensequenz')
            if direction is not None: self._direction(direction)
            if sequence <= self.jog_sequence:
                return {'accepted': False}
            self.jog_sequence = sequence
            self.center_target = None
            self.jog_direction = direction
            self.jog_until = self.clock()+.35 if direction else 0.
            return {'accepted': True}

    def center(self, session, sequence):
        """Explicit slow recentering within an already armed setup session."""
        with self.lock:
            accepted = self.jog(session, None, sequence)
            if accepted['accepted']:
                self.center_target = midpoints(self.config)
            return accepted

    def remember(self, session, direction):
        with self.lock:
            self._require_setup(session)
            if self.center_target is not None:
                raise ValueError('Mittelstellungsfahrt zuerst anhalten oder abwarten')
            if self.jog_direction and self.clock() < self.jog_until:
                raise ValueError('Richtungstaste vor dem Speichern loslassen')
            axis, sign = self._direction(direction)
            # Store a slightly inward commanded limit, never an inferred physical stop.
            value = round(getattr(self, axis)) + (10 if sign < 0 else -10)
            key = axis+('_min' if sign < 0 else '_max')
            origin = self.setup_origin[axis]
            if sign < 0 and value > origin or sign > 0 and value < origin:
                raise ValueError('Für diese Grenze zuerst in die zugehörige Richtung von der Startmitte dieser Sitzung weg fahren; mindestens 10 µs Abstand zur Mitte nötig')
            return self._save(dict(self.config, **{key:value}))

    def capture(self):
        with self.lock:
            if self.active:
                raise ValueError('Vor dem Farbfoto Nachführung stoppen')
            _, jpeg, captured = self.camera.latest_timed()
            if self.camera.error or jpeg is None or captured is None or self.clock()-captured > 1:
                raise ValueError('Kein frisches Kamerabild')
            self.sample, self.sample_token = jpeg, secrets.token_hex(16)
            return {'image_base64': base64.b64encode(jpeg).decode(), 'token': self.sample_token}

    def select(self, token, x, y):
        with self.lock:
            if not self.sample or token != self.sample_token:
                raise ValueError('Beispielfoto erneut aufnehmen')
            if not all(isinstance(v, (int, float)) and not isinstance(v, bool) and
                       math.isfinite(v) and 0 <= v <= 1 for v in (x, y)):
                raise ValueError('Punkt außerhalb des Bildes')
            image = Image.open(BytesIO(self.sample)).convert('HSV')
            px, py = int(x*(image.width-1)), int(y*(image.height-1))
            colors = list(image.crop((max(0,px-2), max(0,py-2),
                                      min(image.width,px+3), min(image.height,py+3))).getdata())
            colors = [(h,s,v) for h,s,v in colors if s >= 100 and v >= 60]
            if len(colors) < 5:
                raise ValueError('Bitte einen hellen, kräftig farbigen Punkt wählen')
            # Circular mean avoids averaging red at both ends into cyan.
            angle = math.atan2(sum(math.sin(h*2*math.pi/256) for h,_,_ in colors),
                               sum(math.cos(h*2*math.pi/256) for h,_,_ in colors))
            config = dict(self.config, hue=round(angle*256/(2*math.pi)) % 256)
            return self.configure(config)

    def begin(self, hardware=False, acknowledged=False, operation='follow'):
        with self.lock:
            if operation not in ('follow', 'setup'):
                raise ValueError('Unbekannter Kameramodus')
            if self.active:
                raise ValueError('Kameramodus läuft bereits')
            if self.config_fault:
                raise ValueError(self.error)
            self.error = None
            _, jpeg, captured = self.camera.latest_timed()
            if self.camera.error or jpeg is None or captured is None or self.clock()-captured > 1:
                raise ValueError('Kein frisches Kamerabild')
            if hardware and (self.camera.demo or not acknowledged):
                raise ValueError('Echte Kamera und bestätigter freier Servotest erforderlich')
            self.setup_origin = midpoints(self.config)
            self.pan, self.tilt = self.setup_origin['pan'], self.setup_origin['tilt']
            try:
                if hardware:
                    self.driver = self.factory()
                    self.driver.move(self.pan, self.tilt)
            except Exception as exc:
                self.halt(f'Servostart fehlgeschlagen: {exc}')
                raise ValueError(self.error) from exc
            self.active = True
            self.operation = operation
            self.jog_direction, self.jog_until, self.jog_sequence = None, 0., -1
            self.center_target = None
            self.session = secrets.token_hex(16)
            self.last_ping = self.clock()
            self.last_result = self.last_seen = self.clock()
            self._reset_search()
            self.result_times.clear()
            self.analysis_ms = None
            self.jpeg = None
            self.target = None
            self.revision += 1
            return self.session

    def heartbeat(self, session):
        with self.lock:
            if not self.active or session != self.session:
                raise ValueError('Nachführungssitzung beendet')
            self.last_ping = self.clock()

    def halt(self, error=None):
        with self.lock:
            self.active = False
            self._reset_search()
            self.session = None
            self.center_target = None
            self.jog_direction = None
            self.jog_until = 0.
            self.revision += 1
            self.target = None
            if self.driver:
                driver, self.driver = self.driver, None
                try:
                    driver.close()
                except Exception as exc:
                    error = f'PWM-Abschaltung nicht bestätigt: {exc}; Servoversorgung ausschalten'
            if error: self.error = error

    def watchdog(self):
        with self.lock:
            if not self.active: return
            now = self.clock()
            _, _, captured = self.camera.latest_timed()
            if self.camera.error or captured is None or now-captured > 1:
                self.halt('Kamerabild ausgefallen; erneut starten')
            elif now-self.last_ping > 1.5:
                self.halt('Bedienverbindung verloren; erneut starten')
            elif self.last_result is not None and now-self.last_result > 1:
                self.halt('Bildauswertung ausgefallen; erneut starten')

    def process(self, number, jpeg, captured, revision, target, analysis_ms=None):
        with self.lock:
            now = self.clock()
            if not self.active or revision != self.revision or captured is None or now-captured > 1:
                return
            self.watchdog()
            if not self.active: return
            # Do not catch up a delayed frame with a large movement jump.
            previous = self.last_result if self.last_result is not None else now
            max_dt = .10 if self.operation == 'follow' else .25
            dt = min(max_dt, max(0, now-previous))
            self.result_times.append(now)
            self.analysis_ms = analysis_ms
            self.last_result, self.number, self.jpeg, self.target = now, number, jpeg, target
            if self.operation == 'setup':
                self.target = None
                if self.center_target is not None:
                    for axis, goal in self.center_target.items():
                        position = getattr(self, axis)
                        setattr(self, axis, position + max(-40*dt, min(40*dt, goal-position)))
                    if self.driver: self.driver.move(self.pan, self.tilt)
                    if all(getattr(self, a) == goal for a, goal in self.center_target.items()):
                        self.center_target = None
                    return
                if self.jog_direction and now < self.jog_until:
                    axis, sign = self._direction(self.jog_direction)
                    value = max(self.config[axis+'_setup_min'], min(self.config[axis+'_setup_max'], getattr(self,axis)+sign*40*dt))
                    setattr(self,axis,value)
                    if self.driver: self.driver.move(self.pan,self.tilt)
                else:
                    self.jog_direction = None
                return
            if self.search_started is not None and now-self.search_started >= 60:
                self._wait_for_target()
                return
            if target:
                if self.search_started is not None or self.waiting_for_target:
                    # Stop searching immediately; require two fresh detections
                    # before returning control to ordinary tracking.
                    self.reacquire_count += 1
                    if self.reacquire_count < 2: return
                    self.search_started = None
                    self.search_path.clear()
                    self.waiting_for_target = False
                self.reacquire_count = 0
                self.last_target = dict(target)
                self.last_seen = now
                for axis, coordinate in (('pan','x'), ('tilt','y')):
                    error = target[coordinate]-.5
                    if abs(error) < .06: continue
                    sign = -1 if self.config[axis+'_reverse'] else 1
                    # The tempo is a speed limit; the loop gain stays capped because
                    # camera/servo delay makes a higher gain swing back and forth.
                    speed = self.config['follow_speed']
                    value = getattr(self, axis) + sign*max(-speed, min(speed, error*FOLLOW_GAIN))*dt
                    setattr(self, axis, max(self.config[axis+'_min'], min(self.config[axis+'_max'], value)))
                if self.driver: self.driver.move(self.pan, self.tilt)
            elif self.driver and self.last_seen is not None:
                self.reacquire_count = 0
                if self.waiting_for_target:
                    return
                if self.config['search_enabled'] and self.last_target is not None:
                    if now-self.last_seen >= 1.5:
                        self._search_step(now, dt)
                else:
                    self._wait_for_target()

    def status(self):
        with self.lock:
            times = self.result_times
            fps = ((len(times)-1)/(times[-1]-times[0])
                   if self.active and len(times)>1 and times[-1]>times[0] and self.clock()-times[-1]<1 else 0.)
            return {'fps': round(fps, 1), 'analysis_ms': self.analysis_ms,
                    'searching': self.search_started is not None,
                    'reacquiring': (self.search_started is not None or self.waiting_for_target) and self.reacquire_count > 0,
                    'waiting_for_target': self.waiting_for_target,
                    'search_seconds': round(self.clock()-self.search_started, 1) if self.search_started is not None else 0,
                    'active': self.active, 'hardware': self.driver is not None,
                    'operation': self.operation, 'demo': self.camera.demo,
                    'error': self.error, 'config': dict(self.config),
                    'pan_us': round(self.pan), 'tilt_us': round(self.tilt),
                    'setup_origin': dict(self.setup_origin),
                    'centering': self.center_target is not None,
                    'target': self.target, 'frame_number': self.number,
                    'image_base64': base64.b64encode(self.jpeg).decode() if self.active and self.jpeg else None}

    def start(self):
        self.thread = threading.Thread(target=self._run, daemon=True, name='colour-follow')
        self.thread.start()

    def _run(self):
        last = 0
        while not self.quit.wait(.01):
            self.watchdog()
            with self.lock:
                if not self.active: continue
                revision, config, operation = self.revision, dict(self.config), self.operation
            number, jpeg, captured = self.camera.latest_timed()
            if number == last or jpeg is None: continue
            last = number
            try:
                began = time.monotonic()
                target = locate(jpeg, config) if operation == 'follow' else None
                analysis_ms = round((time.monotonic()-began)*1000, 1)
                self.process(number, jpeg, captured, revision, target, analysis_ms)
            except Exception as exc:
                with self.lock:
                    if revision == self.revision: self.halt(f'Nachführungsfehler: {exc}')

    def stop(self):
        self.quit.set()
        self.halt()
        if self.thread: self.thread.join(timeout=2)