"""Connect the durable personal queue to the existing released-model runtime."""
import atexit
from html import escape
import json
from pathlib import Path
import time
import threading
import hashlib
import logging

from PIL import Image

from hf_space import release_backend, native_build, generation_progress, extraction_progress, ovoxel_results, example_assets, geometry_results
from hf_space.history_store import Store, NotFound
from hf_space.history_jobs import Runner
from hf_space.history_keepalive import WorkKeepAlive
from hf_space.history_routes import PrivateFileCache
from hf_space.history_preview import saved_preview
from hf_space.thumbnail_assets import thumbnail_path

ROOT = Path(__file__).resolve().parents[1]
LOG = logging.getLogger(__name__)


def deployment_status():
    from hf_space.gpu_mode import enabled
    native = native_build.status()
    models = release_backend.model_status()
    native_steps = 7  # CUDA setup, five native wheels, dependency validation.
    completed = (native_steps if native.get('stage') == 'compiled' else
                 1 + len(native.get('completed', [])) if native.get('stage') == 'compiling' else 0)
    native_progress = {'label': native.get('current') or 'Preparing native dependencies',
                       'completed': min(native_steps, completed), 'total': native_steps, 'unit': 'steps'}
    foreground = release_backend.foreground_status() if enabled() else None
    # A public readiness check must not expose traceback text, runtime file
    # paths, private configuration or job identifiers.
    return {'runtime_id': models.get('epoch'),
            **({'gpu_mode': 'zerogpu', 'foreground': {
                key: foreground[key] for key in ('stage', 'startup_progress') if key in foreground}} if enabled() else {}),
            'native_build': {**{key: native[key] for key in ('stage', 'current') if key in native},
                             'startup_progress': native_progress},
            'models': {key: models[key] for key in ('stage', 'startup_progress') if key in models}}


class Service:
    def __init__(self, *, store=None):
        self._boot_started = time.monotonic()
        self.store = store if store is not None else Store()
        self.private_files = PrivateFileCache()
        self.runner = Runner(self.store, {'generate': self.generate, 'texture': self.texture, 'export': self.export,
                                         'recover': self.recover, 'foreground': self.foreground}, ready=self.ready,
                             progress=self.progress)
        self.keepalive = WorkKeepAlive(self.active_work)
        self._start_lock = threading.Lock()
        self._started = False

    def active_work(self):
        state = deployment_status()
        # Failed startup must not keep a paid machine awake indefinitely.
        failed = any(str(state[key].get('stage', '')).startswith(('failed', 'blocked'))
                     for key in ('native_build', 'models'))
        if failed:
            return False
        if not self.ready('startup'):
            # A visitor may close the page while a genuine cold boot is still
            # running. Keep only that finite controller awake, never idle GPU.
            return time.monotonic()-self._boot_started < getattr(release_backend,'STARTUP_TOTAL_SECONDS',7200)
        return self.store.has_pending()

    def start(self):
        with self._start_lock:
            if self._started:
                return
            self._started = True
        from hf_space.gpu_mode import enabled
        if enabled():
            # History is available independently of models. Opening the page,
            # polling status, and restoring saved records never prepare weights.
            def initialize_history():
                try:
                    self.store.initialize()
                except Exception:
                    LOG.exception('Could not initialize personal history')
            threading.Thread(target=initialize_history, name='history-startup', daemon=True).start()
            atexit.register(self.runner.stop)
            return
        # One small read-only HTTP bootstrap lets existing history cards appear
        # during model loading. It never traverses the Bucket mount, writes task
        # state, triggers model preparation or waits on the weight controller.
        def history_reads():
            for attempt in range(3):
                if self.runner._stop.is_set() or self.store.ready:
                    return
                try:
                    if self.store.initialize_reads():
                        return
                except Exception:
                    LOG.warning('History startup snapshot is temporarily unavailable')
                if self.runner._stop.wait(10 * (attempt + 1)):
                    return
        threading.Thread(target=history_reads, name='history-read-startup', daemon=True).start()
        # Full durable recovery and the FIFO worker still wait for the unique
        # model controller to finish its resource-heavy startup.
        def after_models():
            while not self.runner._stop.wait(1):
                state = deployment_status()
                stages = [str(state[key].get('stage','')) for key in ('native_build','models')]
                if any(stage.startswith(('failed','blocked')) for stage in stages):
                    return
                if self.ready('initialize'):
                    self.runner.start()
                    return
        threading.Thread(target=after_models, name='history-after-model-startup', daemon=True).start()
        # Operator-approved curated refreshes enqueue work through this same
        # durable FIFO. They never create a second inference/export worker.
        from hf_space.curated_refresh import start as start_curated_refresh
        start_curated_refresh(self)
        self.keepalive.start()
        atexit.register(self.runner.stop)
        atexit.register(self.keepalive.stop)

    @staticmethod
    def prepare_models(kind='generate'):
        """Start CPU preparation only after an explicit owner-checked task."""
        from hf_space.gpu_mode import enabled
        if enabled():
            if kind == 'foreground':
                release_backend.start_foreground_preparation()
            else:
                release_backend.start_preparation()

    @staticmethod
    def ready(kind):
        state = deployment_status()
        if state.get('gpu_mode') == 'zerogpu' and kind == 'foreground':
            return state['foreground'].get('stage') == 'ready_released'
        native = state['native_build'].get('stage') == 'compiled'
        # GPU execution requires the CPU pool and native extensions. Page,
        # history and task admission have an independent workspace gate.
        return native and state['models'].get('stage') == 'ready_released'

    def generate(self, job, update):
        token = job['id']
        generation_progress.update(token, started_unix=time.time(), status='queued', label='Waiting for GPU')
        prepared = bool(job.get('foreground_prepared'))
        input_path = self.store.resolve_file(job['owner'], job['id'], 'foreground.png' if prepared else 'input.png')
        with Image.open(input_path) as source:
            image = source.convert('RGBA')
        split = job.get('workflow') == 'geometry_texture'
        update(label='Generating geometry' if split else 'Generating geometry and reference materials')
        generate = release_backend.generate_geometry if split else release_backend.generate
        mesh, _asset, report = generate(image, int(job['seed']), generation_token=token,
                                                       **({'foreground_prepared': True} if prepared else {}))
        return self._publish_generated(mesh, report, job, token, update, geometry=split)

    def texture(self, job, update):
        """Resume from the saved white model; never regenerate its geometry."""
        import tempfile
        if job.get('workflow') != 'geometry_texture' or not job.get('geometry_asset_id'):
            raise ValueError('A saved geometry result is required before generating textures.')
        token = job['operation_id']
        update(label='Loading saved geometry')
        with tempfile.TemporaryDirectory(prefix='supratex-texture-input-') as temporary:
            folder = geometry_results.restore(job['geometry_asset_id'], Path(temporary))
            mesh, _asset, report = release_backend.generate_texture(
                folder, int(job['seed']), generation_token=token)
        return self._publish_generated(mesh, report, job, token, update)

    def _publish_generated(self, mesh, report, job, token, update, *, geometry=False):
        label = 'Saving geometry' if geometry else 'Saving the 3D asset'
        generation_progress.update(token, status='saving', label=label)
        update(label=label)
        identifier = (geometry_results if geometry else ovoxel_results).persist(mesh, report)
        local = Path(mesh).parent
        files, metadata = ovoxel_results.local_manifest(local)
        # Prime the first actual view from local generation output. This is
        # optional acceleration; durable publication remains authoritative.
        try:
            order = ovoxel_results.preview_view_order(metadata)
            first = f'{metadata.get("default_mode",metadata["modes"][0])}-{order[0]:02d}.webp'
            info = files['frames/'+first]
            durable = ovoxel_results.directory(identifier)/'views'/first
            self.private_files.seed(job['owner'], durable, local/'views'/first,
                                    descriptor=dict(info,path=str(durable)))
        except Exception:
            pass
        generation_progress.update(token, status='complete', label='Complete', finished_unix=time.time())
        return {'asset_id': identifier, 'file_manifest':files, 'preview_metadata':metadata}

    def foreground(self, job, update):
        update(label='Removing image background')
        path = self.store.resolve_upload_file(job['owner'], job['upload_id'], 'input.png')
        with Image.open(path) as source:
            image = source.convert('RGBA')
        cutout, metadata = release_backend.remove_background(image, timeout_seconds=120)
        update(label='Saving the background-removed image')
        result = self.store.write_foreground_output(job, cutout, metadata)
        # This optional local mirror avoids reading a just-published Bucket
        # object back through HTTP merely to display it in this browser.
        import tempfile
        try:
            with tempfile.TemporaryDirectory(prefix='supratex-foreground-') as folder:
                local = Path(folder) / 'foreground.png'
                cutout.save(local, format='PNG')
                durable = self.store.root / 'jobs' / job['id'] / 'foreground.png'
                info = result['foreground_file_manifest']['foreground.png']
                self.private_files.seed(job['owner'], durable, local, descriptor=dict(info, path=str(durable)))
        except Exception:
            pass
        return result

    def export(self, job, update):
        token = job['operation_id']
        extraction_progress.start(token)
        update(extraction_progress.status(token))
        file_manifest = {}
        def prepare_download(local, durable):
            from hf_space import artifact_delivery
            remote = artifact_delivery.bucket_url(local)
            if remote:
                info = artifact_delivery.describe(local)
            else:
                info = {'path':str(durable), 'size':Path(local).stat().st_size,
                        'version':hashlib.sha256((str(durable)+str(Path(local).stat().st_mtime_ns)).encode()).hexdigest()}
            file_manifest['asset.glb'] = {key:info[key] for key in ('size','version')}
            progress = lambda completed,total: extraction_progress.transfer(token,completed,total)
            if remote:
                _cached,release = self.private_files.acquire(job['owner'],durable,descriptor=info,progress=progress)
                release()
            else:
                self.private_files.seed(job['owner'], durable, local, progress=progress, descriptor=info)
        output = ovoxel_results.extract(job['asset_id'], token, prepare_download=prepare_download)
        if not Path(output).is_file():
            raise RuntimeError('The exported file is not available.')
        value = extraction_progress.status(token)
        update(value)
        result = {'elapsed_seconds': value.get('elapsed_seconds', 0), 'file_manifest':file_manifest}
        try:
            details = json.loads((Path(output).parent/'extraction.json').read_text(encoding='utf-8'))
            if not isinstance(details, dict):
                details = {}
            pair = (details.get('reprojection'), details.get('gutter'))
            if pair in (('applied', 'observation_guided'), ('skipped_missing_inputs', 'chart_isolated')):
                result['export_postprocess'] = {'reprojection': pair[0], 'gutter': pair[1]}
        except (OSError, ValueError, TypeError):
            # Completed exports from earlier releases may not have this receipt.
            pass
        return result

    def recover(self, job, update):
        from hf_space.preview_recovery import recover
        recover(job['asset_id'])
        return {}

    @staticmethod
    def progress(job, operation):
        kind = operation['kind']
        token = job['id'] if kind == 'generate' else operation['id']
        if kind == 'export':
            return extraction_progress.status(token)
        if kind == 'foreground':
            return None
        return generation_progress.status(token)

    def preview(self, owner, job, view='auto'):
        return saved_preview(self.store, owner, job, view=view)

    def prepare_example(self, owner, name):
        path = self.example_file(name, 'input')
        return self.store.save_upload(owner, path.read_bytes())

    @staticmethod
    def example_file(name, kind):
        if name not in example_assets.EXAMPLE_NAMES:
            raise NotFound('Example not found.')
        source = ROOT/'assets'/f'{name}.png'
        if kind == 'input':
            return source
        if kind == 'thumbnail':
            return Path(thumbnail_path(source))
        if kind == 'model':
            value = example_assets.lookup_name(name)
            if value:
                return Path(value)
        raise NotFound('Example file not found.')

    @staticmethod
    def examples():
        result = []
        for name in example_assets.EXAMPLE_NAMES:
            entry = example_assets.manifest()['examples'][name]
            fallback = f'/api/history/examples/{name}/model'
            # The same website renderer and pinned compressed public model files.
            html = (f'<div class="st-gallery-preview" data-case="result" data-model-public="true" '
                    f'data-model="{escape(entry["source_url"], quote=True)}" data-model-encoding="gzip" '
                    f'data-model-bytes="{entry["compressed_bytes"]}" '
                    f'data-model-decoded-bytes="{entry["glb_bytes"]}" data-model-fallback="{fallback}">'
                    '<div class="st-gallery-canvas"></div><div class="st-gallery-controls">'
                    '<button type="button" data-mode="pbr" aria-pressed="true">PBR</button>'
                    '<button type="button" data-mode="albedo" aria-pressed="false">Albedo</button>'
                    '<button type="button" data-reset>Reset view</button>'
                    '<span class="st-gallery-status" role="status">Loading 3D model…</span></div></div>')
            result.append({'name': name, 'title': name.replace('_', ' ').title(),
                           'image_url': f'/api/history/examples/{name}/input',
                           'thumbnail_url': f'/api/history/examples/{name}/thumbnail',
                           'preview_html': html})
        return result