"""One committed, restart-safe example refresh through the ordinary FIFO queue.

This controller admits work only through Store's normal single-writer methods.
It neither starts another GPU worker nor exposes an administrative HTTP route.
Old public examples remain authoritative until an operator publishes all eight
validated candidates together.
"""
from __future__ import annotations

from copy import deepcopy
import hashlib
import json
from pathlib import Path
import re
import struct
import tempfile
import threading
import time
import uuid

import requests
from urllib3.exceptions import HTTPError as TransportError

from hf_space import artifact_delivery, example_assets, release_backend

PLAN = Path(__file__).with_name('curated-refresh-plan.json')
STATE_ROOT = Path('/data/generated/curated-refresh-v1')
BATCH_ID = '20261021-latest-code-v1'
MAX_RECEIPT_BYTES = 2 * 1024 * 1024
STORAGE_BACKOFF_SECONDS = (5, 10)
_ID = re.compile(r'^[0-9a-f]{32}$')
_SHA = re.compile(r'^[0-9a-f]{64}$')
_REVISION = re.compile(r'^[0-9a-f]{40}$')
_start_lock = threading.Lock()
_coordinator = None


class RefreshError(RuntimeError):
    pass


def _digest(path):
    value = hashlib.sha256()
    with Path(path).open('rb') as stream:
        while chunk := stream.read(4 * 1024 * 1024):
            value.update(chunk)
    return value.hexdigest()


def validate_plan(value):
    """A checked-in data file cannot select owners, paths, code or new inputs."""
    allowed = {'schema', 'enabled', 'batch_id', 'seed', 'model_revision',
               'code_revision', 'checkpoint', 'requested_from_revision', 'examples'}
    if not isinstance(value, dict) or set(value) != allowed:
        raise RefreshError('Invalid curated refresh plan.')
    if (value['schema'] != 1 or type(value['enabled']) is not bool
            or value['batch_id'] != BATCH_ID or type(value['seed']) is not int
            or value['seed'] != 42):
        raise RefreshError('Unsupported curated refresh plan.')
    if (value['model_revision'] != release_backend.MODEL_REVISION
            or value['code_revision'] != release_backend.CODE_REVISION
            or value['checkpoint'] != release_backend.CHECKPOINT
            or not isinstance(value['requested_from_revision'], str)
            or not _REVISION.fullmatch(value['requested_from_revision'])):
        raise RefreshError('The refresh plan does not match the released pipeline.')
    entries = value['examples']
    if (not isinstance(entries, list) or len(entries) != len(example_assets.EXAMPLE_NAMES)
            or tuple(item.get('name') for item in entries if isinstance(item, dict))
            != example_assets.EXAMPLE_NAMES):
        raise RefreshError('The refresh plan must contain exactly the eight approved examples.')
    for entry in entries:
        if (set(entry) != {'name', 'source_id', 'input_sha256', 'input_rgba_sha256'}
                or not isinstance(entry['source_id'], str) or not 1 <= len(entry['source_id']) <= 96
                or any(ord(c) < 32 for c in entry['source_id'])
                or not all(isinstance(entry[key], str) and _SHA.fullmatch(entry[key])
                           for key in ('input_sha256', 'input_rgba_sha256'))):
            raise RefreshError('Invalid approved example provenance.')
    return deepcopy(value)


def _validate_glb(path, expected_size):
    path = Path(path)
    size = path.stat().st_size
    if size != expected_size or size < 20:
        raise RefreshError('The exported GLB size is invalid.')
    with path.open('rb') as stream:
        if struct.unpack('<4sII', stream.read(12)) != (b'glTF', 2, size):
            raise RefreshError('The exported GLB header is invalid.')
        length, kind = struct.unpack('<II', stream.read(8))
        if kind != 0x4e4f534a or not 0 < length <= min(size - 20, 16 * 1024 * 1024):
            raise RefreshError('The exported GLB JSON is invalid.')
        document = json.loads(stream.read(length))
    if (document.get('asset', {}).get('version') != '2.0'
            or not isinstance(document.get('meshes'), list) or not document['meshes']):
        raise RefreshError('The exported GLB has no mesh.')
    for entry in document.get('images', []) + document.get('buffers', []):
        if entry.get('uri') and not entry['uri'].startswith('data:'):
            raise RefreshError('The exported GLB has external dependencies.')
    return {'sha256': _digest(path), 'glb_bytes': size}


class Coordinator:
    def __init__(self, service, *, plan=None, state_root=STATE_ROOT, poll_interval=2):
        self.service = service
        self.store = service.store
        self.plan = validate_plan(plan if plan is not None else json.loads(PLAN.read_text(encoding='utf-8')))
        self.fingerprint = hashlib.sha256(json.dumps(self.plan, sort_keys=True,
                                                    separators=(',', ':')).encode()).hexdigest()
        self.owner = 'system:curated-refresh:' + self.plan['batch_id']
        self.directory = Path(state_root) / self.plan['batch_id']
        self.state_path = self.directory / 'batch.json'
        self.poll_interval = poll_interval
        self.state = None
        self._dirty = False
        self._thread = None
        self._guard = threading.Lock()

    def _save(self):
        self._dirty = True
        self.state['updated_at'] = time.time()
        self.directory.mkdir(parents=True, exist_ok=True)
        temporary = self.state_path.with_name('.batch.' + uuid.uuid4().hex + '.part')
        try:
            temporary.write_text(json.dumps(self.state, indent=2) + '\n', encoding='utf-8')
            temporary.replace(self.state_path)
            self._dirty = False
        finally:
            temporary.unlink(missing_ok=True)

    def _load(self):
        if self.state_path.exists():
            state = json.loads(self.state_path.read_text(encoding='utf-8'))
            if (state.get('schema') != 1 or state.get('batch_id') != self.plan['batch_id']
                    or state.get('plan_sha256') != self.fingerprint
                    or state.get('stage') not in ('running', 'complete', 'failed')
                    or not isinstance(state.get('items'), dict)
                    or set(state['items']) - set(example_assets.EXAMPLE_NAMES)):
                raise RefreshError('The saved refresh state does not match this plan.')
            for name, item in state['items'].items():
                if (not isinstance(item, dict) or item.get('name') != name
                        or not isinstance(item.get('job_id'), str) or not _ID.fullmatch(item['job_id'])
                        or item.get('stage') not in ('generation_queued', 'export_queued', 'verified')):
                    raise RefreshError('The saved example state is invalid.')
            self.state = state
        else:
            self.state = {'schema': 1, 'batch_id': self.plan['batch_id'],
                          'plan_sha256': self.fingerprint, 'stage': 'running',
                          'created_at': time.time(), 'seed': self.plan['seed'],
                          'requested_from_revision': self.plan['requested_from_revision'],
                          'pipeline': {key: self.plan[key] for key in
                                       ('model_revision', 'code_revision', 'checkpoint')},
                          'items': {}, 'publication_candidates': [], 'error': None}
            self._save()

    def _key(self, name, kind):
        return self.plan['batch_id'] + ':' + name + ':' + kind

    def _admit(self, entry):
        """Repeated admission keys survive a crash before batch.json is saved."""
        from PIL import Image
        source = Path(self.service.example_file(entry['name'], 'input'))
        if _digest(source) != entry['input_sha256']:
            raise RefreshError('The approved input image changed.')
        with Image.open(source) as image:
            rgba = image.convert('RGBA')
            pixel_hash = hashlib.sha256(f'RGBA:{rgba.width}:{rgba.height}:'.encode()
                                        + rgba.tobytes()).hexdigest()
        if pixel_hash != entry['input_rgba_sha256']:
            raise RefreshError('The approved input pixels changed.')
        upload = self.store.save_upload(self.owner, source.read_bytes())
        job = self.store.submit(self.owner, upload['id'], self.plan['seed'],
                                self._key(entry['name'], 'generate'), workflow=None)
        self.state['items'][entry['name']] = {
            'name': entry['name'], 'source': deepcopy(entry), 'job_id': job['id'],
            'upload_id': job['upload_id'], 'input_sha256': job['input_sha256'],
            'stage': 'generation_queued', 'generation_key': self._key(entry['name'], 'generate'),
            'export_key': self._key(entry['name'], 'export')}
        self._save()
        self._log(entry['name'], 'generation queued')

    def _receipt(self, asset_id, name):
        """Prefer local output; remote receipts use finite authenticated HTTP."""
        from hf_space import ovoxel_results
        if name not in ('report.json', 'extraction.json') or not _ID.fullmatch(asset_id):
            raise RefreshError('Invalid export receipt.')
        local = ovoxel_results.JOBS / asset_id / name
        if not local.is_symlink() and local.is_file():
            if local.stat().st_size > MAX_RECEIPT_BYTES:
                raise RefreshError('The export receipt is too large.')
            return json.loads(local.read_text(encoding='utf-8'))
        source = self.store.asset_root / asset_id / name
        if artifact_delivery.bucket_url(source):
            info = artifact_delivery.describe(source)
            if not info or info['size'] > MAX_RECEIPT_BYTES:
                raise RefreshError('The export receipt is unavailable.')
            with tempfile.TemporaryDirectory(prefix='supratex-curated-receipt-') as folder:
                destination = Path(folder) / name
                artifact_delivery.download(source, destination, info['size'])
                return json.loads(destination.read_text(encoding='utf-8'))
        if source.is_symlink() or source.stat().st_size > MAX_RECEIPT_BYTES:
            raise RefreshError('The export receipt is unavailable.')
        return json.loads(source.read_text(encoding='utf-8'))

    def _verify(self, entry, item, job):
        asset_id = job.get('asset_id')
        if not isinstance(asset_id, str) or not _ID.fullmatch(asset_id):
            raise RefreshError('The completed example has no generated asset.')
        report, receipt = self._receipt(asset_id, 'report.json'), self._receipt(asset_id, 'extraction.json')
        if (report.get('status') != 'complete' or report.get('asset_id') != asset_id
                or report.get('project_revision') != self.plan['model_revision']
                or report.get('code_revision') != self.plan['code_revision']
                or report.get('checkpoint') != self.plan['checkpoint']
                or report.get('weight_origin') != 'released'
                or report.get('output_kind') != 'final_ovoxel'
                or report.get('automatic_glb_export') is not False
                or report.get('backview') is not False or report.get('refine') is not False):
            raise RefreshError('The example was not generated by the planned pipeline.')
        if (receipt.get('status') != 'complete' or receipt.get('reprojection') != 'applied'
                or receipt.get('gutter') != 'observation_guided'
                or receipt.get('postprocess_warning') is not None
                or job.get('export_postprocess') != {'reprojection': 'applied', 'gutter': 'observation_guided'}):
            raise RefreshError('The example is missing required GLB postprocessing.')
        descriptor = self.store.get_file_descriptor(self.owner, job['id'], 'asset.glb')
        expected = self.store.asset_root / asset_id / 'asset.glb'
        if not descriptor or descriptor.get('path') != str(expected):
            raise RefreshError('The completed GLB descriptor is unavailable.')
        local, release = self.service.private_files.acquire(self.owner, expected, descriptor=descriptor)
        try:
            details = _validate_glb(local, descriptor['size'])
        finally:
            release()
        candidate = {'name': entry['name'], 'source': deepcopy(entry), 'seed': self.plan['seed'],
                     'job_id': job['id'], 'asset_id': asset_id, 'path': str(expected),
                     'input_sha256': item['input_sha256'], **details,
                     'pipeline': deepcopy(self.state['pipeline']),
                     'reprojection': 'applied', 'gutter': 'observation_guided',
                     'verified_at': time.time()}
        item.update(stage='verified', asset_id=asset_id, result=candidate)
        # Candidates are private records, never a partial public manifest.
        self.state['publication_candidates'] = [self.state['items'][name]['result']
            for name in example_assets.EXAMPLE_NAMES
            if self.state['items'].get(name, {}).get('stage') == 'verified']
        self._save()
        self._log(entry['name'], 'GLB verified')

    @staticmethod
    def _log(name, phase):
        # Only the public batch/name and a fixed phase enter container logs.
        print('[curated-refresh] ' + BATCH_ID + ' ' + name + ' ' + phase, flush=True)

    def run_once(self):
        """Advance one durable boundary, leaving GPU execution to Runner."""
        if (not self.plan['enabled'] or self.service.runner._stop.is_set()
                or not self.store.ready or not self.service.ready('curated_refresh')):
            return False
        if self.state is None:
            self._load()
        if self.state['stage'] in ('complete', 'failed'):
            return False
        for entry in self.plan['examples']:
            item = self.state['items'].get(entry['name'])
            if item is None:
                self._admit(entry)
                return True
            if item['stage'] == 'verified':
                continue
            job = self.store.get_job(self.owner, item['job_id'])
            if (job.get('seed') != self.plan['seed'] or job.get('input_sha256') != item['input_sha256']
                    or job.get('upload_id') != item['upload_id']):
                raise RefreshError('The saved refresh job changed.')
            if job['status'] in ('failed', 'interrupted'):
                raise RefreshError('The example generation failed; no automatic resubmission was made.')
            if job['status'] != 'complete':
                return False
            if job.get('export_status') in ('failed', 'interrupted'):
                raise RefreshError('The example GLB extraction failed; no automatic resubmission was made.')
            if job.get('export_status') == 'complete':
                self._verify(entry, item, job)
                return True
            if item['stage'] != 'export_queued':
                self.store.submit_export(self.owner, job['id'], self._key(entry['name'], 'export'))
                item['stage'] = 'export_queued'
                self._save()
                self._log(entry['name'], 'GLB extraction queued')
                return True
            return False
        if len(self.state['publication_candidates']) != len(example_assets.EXAMPLE_NAMES):
            raise RefreshError('The completed refresh is missing publication candidates.')
        self.state.update(stage='complete', completed_at=time.time())
        self._save()
        self._log('all', 'eight GLBs ready for publication')
        return True

    def _run(self):
        storage_failures = 0
        while not self.service.runner._stop.is_set():
            try:
                # Retry an unacknowledged state write before crossing another
                # boundary; admission/export keys remain unchanged throughout.
                if self._dirty:
                    self._save()
                self.run_once()
                storage_failures = 0
            except Exception as error:
                retryable = (not isinstance(error, RefreshError) and (
                    isinstance(error, (OSError, requests.RequestException, TransportError))
                    or getattr(error, 'status_code', None) in (408, 429, 500, 502, 503, 504)
                    or type(error) is RuntimeError and str(error) ==
                    'Artifact storage is temporarily unavailable. Please retry.'))
                if retryable and storage_failures < len(STORAGE_BACKOFF_SECONDS):
                    delay = STORAGE_BACKOFF_SECONDS[storage_failures]
                    storage_failures += 1
                    self._log('all', 'waiting for storage; retry ' + str(storage_failures) + '/2')
                    self.service.runner._stop.wait(delay)
                    continue
                if self.state is not None:
                    self.state.update(stage='failed', error={'type': type(error).__name__,
                                                            'message': str(error)[:500]})
                    try:
                        self._save()
                    except Exception:
                        pass
                self._log('all', 'stopped: ' + type(error).__name__)
                return
            if self.state and self.state['stage'] in ('complete', 'failed'):
                return
            self.service.runner._stop.wait(self.poll_interval)

    def start(self):
        with self._guard:
            if self._thread is None:
                self._thread = threading.Thread(target=self._run, name='curated-example-refresh', daemon=True)
                self._thread.start()
        return self


def start(service):
    """Exactly one service-owned refresh controller in the application process."""
    global _coordinator
    with _start_lock:
        if _coordinator is None:
            try:
                candidate = Coordinator(service)
                candidate.start()
            except Exception as error:
                # A bad one-time plan must not prevent the public app or its
                # existing user queue from starting normally.
                print('[curated-refresh] ' + BATCH_ID + ' all not started: '
                      + type(error).__name__, flush=True)
                return None
            _coordinator = candidate
        return _coordinator