"""Sparse training admission from the single compressed-cache owner."""
from collections import Counter
from dataclasses import replace
import json
import os
import struct
import time

from app.training_store import TrainingStore, root_path, MARKER


class TrainingCapture:
    def __init__(self, spec):
        from app.provider_delivery import runtime_plan
        selected = runtime_plan(spec)
        self.provider = next((p for p in selected.providers if p.person_owner), None)
        self.store = TrainingStore(root_path()) if self.provider and root_path() else None
        self.pending = None
        self.refresh_at = self.retry_at = 0
        self.due_at = self.admission_refresh_at = 0
        self.cfg = None
        self.metrics = Counter()
        from app.intrusion import IntrusionConfig
        from app.loitering import LoiteringConfig
        from app.fall_runtime import FallSettings
        from app.fight_runtime import FightSettings
        self.thresholds = []
        for name, parser in (('intrusion',IntrusionConfig),('loitering',LoiteringConfig),
                             ('fall',FallSettings),('fight',FightSettings)):
            cfg = parser.parse(spec.get(name) or {})
            if cfg.enabled:
                self.thresholds.append((name.upper(),cfg.confidence,name=='intrusion'))
        from app.pose_config import PoseConfig
        self.pose = PoseConfig.parse(spec.get('pose') or {})

    def candidates(self, result):
        commands = tuple(c for c in result.commands if c.detection.class_id == 0
                         and c.detection.model_key == self.provider.model)
        if self.provider.key == 'pose' and self.pose.nms_enabled:
            from app.pose_nms import pose_nms_indices
            kept = set(pose_nms_indices([c.detection for c in commands], self.pose.nms_iou_threshold))
            commands = tuple(c for i,c in enumerate(commands) if i in kept)
        return tuple(c for c in commands if any(c.detection.score > threshold if strict else
                     c.detection.score >= threshold for _,threshold,strict in self.thresholds))

    def consider(self, result):
        if self.store is None or self.pending:
            return
        try:
            now = time.monotonic()
            if now >= self.refresh_at:
                current = self.store.settings()
                if self.cfg is None or current['revision'] != self.cfg['revision']:
                    self.due_at = self.admission_refresh_at = self.retry_at = 0
                self.cfg = current
                self.refresh_at = now+1
                self.metrics['settings_reads'] += 1
            cfg = self.cfg['configured']
            if not cfg['enabled']:
                return
            if now < self.retry_at or now < self.due_at:
                self.metrics['early_interval_skip'] += 1
                return
            if now >= self.admission_refresh_at:
                delay = self.store.admission_wait(result.frame.source_id, cfg)
                self.metrics['admission_reads'] += 1
                self.due_at = now+delay
                self.admission_refresh_at = now+1
                if delay > 0:
                    self.metrics['early_interval_skip'] += 1
                    return
            self.metrics['candidate_checks'] += 1
            commands = self.candidates(result)
            if len(commands) < cfg['minimum_persons'] or result.frame.pts_ns is None:
                return
            job, reason = self.store.reserve(result.frame.source_id, cfg)
            if reason:
                self.metrics['skipped_'+reason] += 1
                self.retry_at = now+1 if reason == 'interval' else now+min(30, cfg['interval_seconds'])
                self.admission_refresh_at = 0
                return
            # Freeze prediction values, not SDK/GPU buffers or mutable scenario state.
            snapshot = replace(result, commands=tuple(replace(c, visible=True) for c in commands),
                               events=(), nms_kept=None).snapshot(visible_only=False)
            self.pending = dict(job=job, snapshot=snapshot, cfg=self.cfg, created=now,
                                count=len(commands))
            self.metrics['reserved'] += 1
        except Exception as error:
            self.metrics['error_'+type(error).__name__] += 1
            self.retry_at = time.monotonic()+30

    def collect(self, packets, *, stopped=False):
        if not self.pending:
            return
        from app.event_capture import PACKET
        p = self.pending
        pts = p['snapshot']['ptsNs']
        rows = list(packets)
        now = time.monotonic()
        if not stopped and now-p['created'] < 2 and (not rows or max(t for t,_ in rows) < pts+500_000_000):
            return
        try:
            keys = [i for i,(t,b) in enumerate(rows) if t <= pts and PACKET.unpack_from(b)[5]]
            if not keys or not any(t == pts for t,_ in rows):
                raise ValueError('NO_MATCHING_FRAME_OR_KEYFRAME')
            stage = self.store.safe('.staging/'+p['job'])
            with (stage/'capture.bin').open('xb') as f:
                # Decode order retained, including trailing reference frames for B-frames.
                for _, raw in rows[keys[-1]:]:
                    f.write(struct.pack('<I',len(raw))); f.write(raw)
            data = dict(owner=MARKER, jobId=p['job'], snapshot=p['snapshot'], settings=p['cfg'],
                        providerKey=self.provider.key, providerModel=self.provider.model,
                        modelVersion=os.environ.get('CONTEXTRA_TRAINING_MODEL_VERSION', 'unknown'),
                        personCount=p['count'], predictionKind='MODEL_PREDICTION_NOT_GROUND_TRUTH',
                        selection=dict(scope='FULL_FRAME_PERSON_SCORE_UNION', thresholds=self.thresholds,
                                       roi=False, tracking=False, temporal=False))
            with (stage/'job.json').open('x') as f:
                json.dump(data,f,allow_nan=False)
            with self.store.lock(), self.store.db() as db:
                db.execute("UPDATE jobs SET state='QUEUED' WHERE id=? AND state='CAPTURING'", (p['job'],))
            self.metrics['queued'] += 1
        except Exception as error:
            self.store.finish(p['job'], type(error).__name__)
            self.metrics['capture_failed'] += 1
        finally:
            self.pending = None
            self.retry_at = now+1
            self.due_at = self.admission_refresh_at = 0

    def status(self):
        return dict(self.metrics, pending=bool(self.pending), enabled=self.store is not None)
