"""任务服务：CPU 负责文件，GPU 租约仅覆盖真实 Blender 渲染。"""
from __future__ import annotations

import spaces  # 必须先于可能加载 CUDA 的库导入。
import contextlib
import importlib.metadata
import json
import os
from pathlib import Path
import re
import secrets
import shutil
import subprocess
import threading
import time
import zipfile

from blender_runtime import ensure_blender, run_task
from render_config import RenderConfig, parse_config, json_schema
from uploads import stage_uploads

BASE = Path(__file__).resolve().parent
DATA_ROOT = Path(os.environ.get('RENDER_DATA_DIR', BASE / 'runtime_data')).resolve()
JOB_ROOT = DATA_ROOT / 'jobs'
PUBLIC_ROOT = DATA_ROOT / 'public'
for folder in (JOB_ROOT, PUBLIC_ROOT):
    folder.mkdir(parents=True, exist_ok=True)
_LOCKS: dict[str, threading.Lock] = {}
_LOCKS_GUARD = threading.Lock()


def _atomic_json(path: Path, data: dict) -> None:
    tmp = path.with_suffix('.tmp-' + secrets.token_hex(4))
    tmp.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding='utf-8')
    os.replace(tmp, path)


def _job(job_id: str) -> Path:
    # ID 同时是访问凭证，不允许任意路径访问。
    if not re.fullmatch(r'[0-9a-f]{48}', job_id or ''):
        raise ValueError('任务 ID 无效。请先准备任务。')
    path = JOB_ROOT / job_id
    if not (path / 'job.json').is_file():
        raise ValueError('任务不存在或已过期；Space 临时磁盘重置后请重新上传。')
    return path


def _load(path: Path) -> dict:
    return json.loads((path / 'job.json').read_text(encoding='utf-8'))


def _save(path: Path, meta: dict) -> None:
    meta['updated_at'] = time.time()
    _atomic_json(path / 'job.json', meta)


@contextlib.contextmanager
def _lock(path: Path):
    with _LOCKS_GUARD:
        lock = _LOCKS.setdefault(path.name, threading.Lock())
    if not lock.acquire(blocking=False):
        raise ValueError('该任务正在执行，请等待或取消当前批次。')
    handle = None
    acquired = False
    try:
        handle = (path / '.lock').open('a+')
        # Linux 文件锁使重入和多进程请求也不能同时更改同一任务。
        import fcntl
        try:
            fcntl.flock(handle, fcntl.LOCK_EX | fcntl.LOCK_NB)
            acquired = True
        except BlockingIOError:
            raise ValueError('该任务正在另一个进程执行。') from None
        yield
    finally:
        if handle is not None:
            # ZeroGPU forks can inherit this descriptor. Closing the parent copy
            # alone leaves the lock held until all forked copies disappear.
            try:
                if acquired:
                    fcntl.flock(handle, fcntl.LOCK_UN)
            finally:
                handle.close()
        lock.release()


def _backend(cfg: RenderConfig | None = None) -> str:
    # Each job explicitly chooses CPU or GPU. Deployment hardware stays available.
    if cfg is None or cfg.device == 'CPU':
        return 'cpu'
    choice = os.environ.get('RENDER_BACKEND', 'auto').lower()
    from spaces.config import Config
    zero = bool(Config.zero_gpu)
    if choice in {'auto', 'cpu'}:
        return 'zerogpu' if zero else 'gpu'
    if choice == 'zerogpu' and not zero:
        raise ValueError('当前不是 ZeroGPU 运行环境；请选择 CPU 或配置可用 GPU。')
    if choice not in {'zerogpu', 'gpu'}:
        raise ValueError('RENDER_BACKEND 必须为 auto、zerogpu、gpu 或 cpu。')
    return choice


def _config(meta: dict) -> RenderConfig:
    saved = dict(meta['config'])
    # Pre-1.1 jobs were always Cycles. Never mix Eevee into a resumed sequence.
    saved.setdefault('engine', 'CYCLES')
    return parse_config(json.dumps(saved))


def _frame_path(path: Path, frame: int, config: RenderConfig) -> Path:
    return path / 'frames' / f"frame_{frame:06d}.{'png' if config.format == 'PNG' else 'exr'}"


def _completed(path: Path, config: RenderConfig, verify: bool = False) -> list[int]:
    result = []
    # 只信任 worker 原子改名完成的最终文件，不接纳 .part 或用户上传文件。
    for frame in config.frame_numbers():
        p = _frame_path(path, frame, config)
        if not p.is_file() or p.stat().st_size < 32:
            continue
        if verify and config.format == 'PNG':
            try:
                from PIL import Image
                with Image.open(p) as img:
                    if img.size != (config.output_width, config.output_height):
                        continue
                    img.verify()
            except Exception:
                continue
        result.append(frame)
    return result


def _public_report(value):
    """JSON 中的 path 会被 Gradio Client 当作可下载文件；内部路径只输出文件名。"""
    if isinstance(value, dict):
        return {('filename' if key == 'path' else key):
                (Path(item).name if key == 'path' and isinstance(item, str) else _public_report(item))
                for key, item in value.items()}
    if isinstance(value, list):
        return [_public_report(item) for item in value]
    return value


def _status(path: Path, verified_done: list[int] | None = None) -> dict:
    meta = _load(path)
    cfg = _config(meta)
    done = _completed(path, cfg) if verified_done is None else verified_done
    expected = cfg.frame_numbers()
    status = meta['status']
    if (path / 'cancel.flag').exists():
        status = 'cancelling' if status == 'rendering' else 'cancelled'
    return {
        'job_id': path.name, 'status': status, 'mode': cfg.mode,
        'backend': meta.get('backend'), 'completed_count': len(done),
        'total_frames': len(expected), 'completed_frames': done,
        'remaining_frames': [f for f in expected if f not in set(done)],
        'progress': round(len(done) / len(expected), 4),
        'config': meta['config'], 'inspection': _public_report(meta.get('inspection')),
        'last_render': _public_report(meta.get('last_render')), 'error': meta.get('error'),
        'created_at': meta['created_at'], 'updated_at': meta['updated_at'],
        'note': '任务 ID 是访问凭证，请勿公开；断点依赖当前 Space 磁盘。',
    }


def _cleanup() -> None:
    """准备新任务时清理过期的空闲任务，不删除仍持有执行锁的任务。"""
    ttl = int(os.environ.get('JOB_TTL_HOURS', '24')) * 3600
    now = time.time()
    for public in PUBLIC_ROOT.glob('probe-*'):
        if public.is_dir() and now - public.stat().st_mtime > ttl:
            shutil.rmtree(public, ignore_errors=True)
    for path in JOB_ROOT.iterdir():
        if not path.is_dir() or not (path / 'job.json').is_file():
            continue
        try:
            if now - _load(path)['updated_at'] > ttl:
                with _lock(path):
                    shutil.rmtree(PUBLIC_ROOT / path.name, ignore_errors=True)
                    shutil.rmtree(path)
                with _LOCKS_GUARD:
                    _LOCKS.pop(path.name, None)
        except (ValueError, OSError, KeyError, json.JSONDecodeError):
            continue


def environment_info() -> dict:
    """Describe deployment capabilities without allocating a GPU."""
    from spaces.config import Config
    versions = {}
    for name in ('gradio', 'spaces', 'torch'):
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    return {
        'app_version': '1.2.0', 'versions': versions,
        'zerogpu_environment': bool(Config.zero_gpu),
        'configured_backend': _backend(RenderConfig()),
        'gpu_backend_preference': os.environ.get('RENDER_BACKEND', 'auto'),
        'default_engine': 'EEVEE', 'default_device': 'AUTO',
        'default_allocates_gpu': True,
        'supported_engines': ['EEVEE', 'CYCLES'],
        'ffmpeg_available': bool(shutil.which('ffmpeg')),
        'gpu_verified': False,
        'message': '默认 GPU + Eevee。开始正式渲染前请运行 GPU / Eevee 自检；GPU 计算可用不等于图形接口可用。CPU 和 Cycles 仍可选。',
        'mcp_path': '/gradio_api/mcp/',
        'job_ttl_hours': int(os.environ.get('JOB_TTL_HOURS', '24')),
    }


def config_schema() -> dict:
    """Return the JSON Schema for supported render settings."""
    return json_schema()


def prepare_job(files: list[str], config_json: str, config_file: str | None = None) -> tuple[str, dict]:
    """Stage Blender files and validate configuration without a GPU allocation.

    Args:
        files: Uploaded .blend, resource ZIPs, textures and cache files.
        config_json: Render configuration as a JSON string.
        config_file: Optional uploaded JSON file; takes precedence over config_json.
    """
    if config_file:
        p = Path(config_file)
        if p.stat().st_size > 65536:
            raise ValueError('配置文件不得超过 64 KiB。')
        config_json = p.read_text(encoding='utf-8-sig')
    cfg = parse_config(config_json)
    backend = _backend(cfg)
    _cleanup()
    # 预留任务目录与数量检查是一个原子步骤，避免并行 API 越过任务上限。
    import fcntl
    with (DATA_ROOT / 'prepare.lock').open('a+') as reservation:
        fcntl.flock(reservation, fcntl.LOCK_EX)
        maximum = int(os.environ.get('MAX_JOBS', '32'))
        if sum(p.is_dir() for p in JOB_ROOT.iterdir()) >= maximum:
            raise ValueError(f'任务数量已到上限 {maximum}，请等待过期清理或由管理员清理空闲任务。')
        free = shutil.disk_usage(DATA_ROOT).free
        if free < int(os.environ.get('MIN_FREE_MB', '512')) * 1024**2:
            raise ValueError('磁盘剩余空间不足，请先清理任务。')
        job_id = secrets.token_hex(24)
        path = JOB_ROOT / job_id
        path.mkdir(mode=0o700)
    try:
        blend = stage_uploads(files, path / 'project', cfg.blend_file)
        (path / 'frames').mkdir()
        meta = {'job_id': job_id, 'status': 'prepared', 'config': cfg.model_dump(),
                'blend_relative': blend.relative_to(path / 'project').as_posix(),
                'created_at': time.time(), 'backend': backend, 'error': None}
        _save(path, meta)
        return job_id, _status(path)
    except Exception:
        shutil.rmtree(path, ignore_errors=True)
        raise


def _task(path: Path, meta: dict, action: str, frames: list[int] | None = None) -> dict:
    return {'action': action, 'project_root': str(path / 'project'),
            'blend_path': str(path / 'project' / meta['blend_relative']),
            'output_dir': str(path / 'frames'), 'result_path': str(path / f'{action}-result.json'),
            'frames': frames or [], 'config': meta['config']}


def inspect_job(job_id: str) -> dict:
    """Inspect scene, cameras and resources before rendering.

    Args:
        job_id: Private job identifier returned by prepare_job.
    """
    path = _job(job_id)
    ensure_blender()
    with _lock(path):
        meta = _load(path)
        try:
            result = run_task(_task(path, meta, 'inspect'), 120, path / 'inspect.log')
            if not result.get('ok', False):
                raise RuntimeError(result.get('error', '场景检查失败。'))
            meta.update(inspection=result, status='ready', error=None)
            _save(path, meta)
        except Exception as exc:
            meta.update(status='failed', error=str(exc)); _save(path, meta)
    return _status(path)


def _duration(task: dict, log_path: str, cancel_path: str, seconds: int) -> int:
    return seconds


def _allocated_render(task: dict, log_path: str, cancel_path: str, seconds: int) -> dict:
    # 租约内启动外部 Blender，保留平台给出的 CUDA_VISIBLE_DEVICES。
    # 给平台超时预留收尾时间；子进程 timeout 会终止整个进程组。
    diagnostics = _graphics_diagnostics() if task.get('action') == 'probe' else None
    try:
        result = run_task(task, max(5, seconds - 8), Path(log_path), Path(cancel_path) if cancel_path else None)
    except Exception as exc:
        # Return diagnostics across the ZeroGPU boundary instead of losing the
        # original exception to its generic remote RuntimeError wrapper.
        result = {'ok': False, 'error': str(exc), 'error_type': type(exc).__name__}
    if diagnostics is not None:
        result['graphics_diagnostics'] = diagnostics
    return result


def _graphics_diagnostics() -> dict:
    """Inspect only graphics-related runtime data within the allocated device lease."""
    import ctypes.util
    result = {'driver_capabilities': os.environ.get('NVIDIA_DRIVER_CAPABILITIES', ''),
              'egl_vendors': sorted(p.name for p in Path('/usr/share/glvnd/egl_vendor.d').glob('*.json')),
              'vulkan_icds': sorted(p.name for p in Path('/usr/share/vulkan/icd.d').glob('*.json')),
              'libraries': {name: ctypes.util.find_library(name) for name in ('EGL_nvidia', 'GLX_nvidia', 'vulkan')}}
    try:
        import torch
        result['cuda_device_name'] = torch.cuda.get_device_name(0)
    except Exception as exc:
        result['cuda_device_query_error'] = type(exc).__name__
    try:
        output = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True, timeout=5, check=False)
        result['allocated_devices'] = re.sub(r'UUID:[^)]*', 'UUID: redacted', output.stdout)[:4000]
        result['nvidia_smi_exit_code'] = output.returncode
    except (OSError, subprocess.TimeoutExpired) as exc:
        result['nvidia_smi_error'] = type(exc).__name__
    return result


# 只有 ZeroGPU 环境才构造装饰器，本地 CPU/GPU 无需安装 PyTorch。
from spaces.config import Config as _SpacesConfig


def _allocated_render_xlarge(task: dict, log_path: str, cancel_path: str, seconds: int) -> dict:
    # spaces caches decorators by function identity; sizes need distinct entry points.
    return _allocated_render(task, log_path, cancel_path, seconds)


_zero_render = spaces.GPU(duration=_duration)(_allocated_render) if _SpacesConfig.zero_gpu else None
_zero_render_xlarge = spaces.GPU(duration=_duration, size='xlarge')(_allocated_render_xlarge) if _SpacesConfig.zero_gpu else None


def _execute(task: dict, path: Path, backend: str, cfg: RenderConfig, log: Path) -> dict:
    if backend == 'zerogpu':
        if _zero_render is None:
            raise ValueError('ZeroGPU 装饰器未激活；请检查 Space 硬件配置并重启。')
        runner = _zero_render_xlarge if cfg.gpu_size == 'xlarge' else _zero_render
        if runner is None:
            raise ValueError('所选 ZeroGPU 规格未激活。')
        return runner(task, str(log), str(path / 'cancel.flag'), cfg.gpu_seconds)
    return run_task(task, cfg.cpu_timeout, log, path / 'cancel.flag')


def render_chunk(job_id: str) -> dict:
    """Render one bounded batch, preserving completed frames on failure.

    Args:
        job_id: Private job identifier returned by prepare_job.
    """
    path = _job(job_id)
    ensure_blender()  # 下载/安装绝不占用 GPU 租约。
    with _lock(path):
        meta = _load(path); cfg = _config(meta)
        if (path / 'cancel.flag').exists():
            raise ValueError('任务已暂停；先调用 resume_job 清除取消标记。')
        done = set(_completed(path, cfg, verify=True))
        for frame in cfg.frame_numbers():
            if frame not in done:
                _frame_path(path, frame, cfg).unlink(missing_ok=True)
        pending = [f for f in cfg.frame_numbers() if f not in done]
        if not pending:
            meta.update(status='complete', error=None); _save(path, meta)
            return _status(path)
        batch = pending[:cfg.batch_size]
        meta.update(status='rendering', error=None); _save(path, meta)
        try:
            result = _execute(_task(path, meta, 'render', batch), path, meta['backend'], cfg,
                              path / f'render-{batch[0]:06d}-{int(time.time())}.log')
            meta['last_render'] = result
            if not result.get('ok', False):
                raise RuntimeError(result.get('error', 'Blender 渲染失败。'))
            meta['status'] = 'complete' if len(_completed(path, cfg, verify=True)) == len(cfg.frame_numbers()) else 'partial'
        except Exception as exc:
            meta['status'] = 'cancelled' if (path / 'cancel.flag').exists() else 'partial' if _completed(path, cfg) else 'failed'
            meta['error'] = str(exc)
        _save(path, meta)
    return _status(path)


def job_status(job_id: str) -> dict:
    """Read progress without allocating a GPU.

    Args:
        job_id: Private job identifier returned by prepare_job.
    """
    return _status(_job(job_id))


def cancel_job(job_id: str) -> dict:
    """Request cancellation of the active Blender process, retaining finished frames.

    Args:
        job_id: Private job identifier returned by prepare_job.
    """
    path = _job(job_id)
    (path / 'cancel.flag').touch()
    return {**_status(path), 'message': '已请求暂停；正在执行的 Blender 会在进程监控轮询时结束。'}


def resume_job(job_id: str) -> dict:
    """Clear cancellation and validate saved frames; call render_chunk to continue.

    Args:
        job_id: Private job identifier returned by prepare_job.
    """
    path = _job(job_id)
    with _lock(path):
        (path / 'cancel.flag').unlink(missing_ok=True)
        meta = _load(path); cfg = _config(meta)
        done = _completed(path, cfg, verify=True)
        for frame in cfg.frame_numbers():
            if frame not in done:
                _frame_path(path, frame, cfg).unlink(missing_ok=True)
        meta.update(status='complete' if len(done) == len(cfg.frame_numbers()) else 'partial', error=None)
        _save(path, meta)
    return _status(path, verified_done=done)


def _export(path: Path, source: Path) -> str:
    target = PUBLIC_ROOT / path.name
    target.mkdir(parents=True, exist_ok=True)
    dest = target / source.name
    shutil.copy2(source, dest)
    return str(dest)


def _encode_video(path: Path, cfg: RenderConfig) -> Path:
    ffmpeg = shutil.which('ffmpeg')
    if not ffmpeg:
        raise RuntimeError('找不到 FFmpeg；请在 Space 的 packages.txt 安装 ffmpeg。')
    # 使用顺序硬链接避免 frame_step > 1 时 ffmpeg 将间隔误认为缺帧。
    sequence = path / 'video-sequence'
    sequence.mkdir(exist_ok=True)
    for index, frame in enumerate(cfg.frame_numbers()):
        dst = sequence / f'{index:06d}.png'
        if not dst.exists():
            try:
                os.link(_frame_path(path, frame, cfg), dst)
            except OSError:
                shutil.copy2(_frame_path(path, frame, cfg), dst)
    output = path / 'animation.mp4'
    temporary = path / 'animation.part.mp4'
    # step=2 表示隔帧采样，播放帧率同步除2以维持时间尺度。
    fps = f'{cfg.fps}/{cfg.frame_step}'
    command = [ffmpeg, '-hide_banner', '-nostdin', '-y', '-framerate', fps,
               '-i', str(sequence / '%06d.png'), '-c:v', 'libx264', '-preset', cfg.video_preset,
               '-crf', str(cfg.video_crf), '-pix_fmt', 'yuv420p', '-movflags', '+faststart',
               '-threads', os.environ.get('FFMPEG_THREADS', '4'), '-an', str(temporary)]
    with (path / 'ffmpeg.log').open('w') as log:
        try:
            subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True, timeout=cfg.cpu_timeout)
            os.replace(temporary, output)
        except Exception as exc:
            temporary.unlink(missing_ok=True)
            raise RuntimeError('视频合成失败；帧图片仍保留，详情见 ffmpeg.log。') from exc
        finally:
            shutil.rmtree(sequence, ignore_errors=True)
    return output


def finalize_job(job_id: str) -> tuple[dict, list[str], str | None, str | None]:
    """Export completed frames, manifest and logs; encode MP4 when every video frame exists.

    Args:
        job_id: Private job identifier returned by prepare_job.
    """
    path = _job(job_id)
    with _lock(path):
        meta = _load(path); cfg = _config(meta)
        done = _completed(path, cfg, verify=True)
        summary = _status(path)
        summary.update(completed_frames=done, completed_count=len(done),
                       remaining_frames=[f for f in cfg.frame_numbers() if f not in set(done)],
                       progress=round(len(done) / len(cfg.frame_numbers()), 4))
        if len(done) < len(cfg.frame_numbers()) and summary['status'] == 'complete':
            summary['status'] = 'partial'
        files = []; image = video = None
        if done and cfg.format == 'PNG':
            image = _export(path, _frame_path(path, done[0], cfg))
            files.append(image)
        if cfg.mode == 'video' and len(done) == len(cfg.frame_numbers()):
            try:
                video = _export(path, _encode_video(path, cfg))
                files.append(video)
            except Exception as exc:
                summary['export_error'] = str(exc)
        summary['export_complete'] = len(done) == len(cfg.frame_numbers()) and not summary.get('export_error')
        summary['video_playback_fps'] = cfg.fps / cfg.frame_step if cfg.mode == 'video' else None
        _atomic_json(path / 'manifest.json', summary)
        archive = path / 'render-output.zip'
        temporary = path / 'render-output.part.zip'
        with zipfile.ZipFile(temporary, 'w', compression=zipfile.ZIP_STORED) as z:
            z.write(path / 'manifest.json', 'manifest.json')
            for frame in done:
                p = _frame_path(path, frame, cfg); z.write(p, f'frames/{p.name}')
            for log in sorted(path.glob('*.log')):
                z.write(log, 'logs/' + log.name)
            for result in sorted(path.glob('*-result.json')):
                z.write(result, 'reports/' + result.name)
            if video:
                z.write(path / 'animation.mp4', 'animation.mp4')
        os.replace(temporary, archive)
        files += [_export(path, archive), _export(path, path / 'manifest.json')]
        return summary, files, image, video


def cpu_self_test() -> tuple[dict, list[str], str | None]:
    """Render a real Eevee frame on the CPU without requesting a GPU lease."""
    return _self_test(RenderConfig(engine='EEVEE', device='CPU'))


def gpu_self_test() -> tuple[dict, list[str], str | None]:
    """Explicitly test Cycles GPU; never invoked by CPU jobs."""
    return _self_test(RenderConfig(engine='CYCLES', device='AUTO'))


def eevee_gpu_self_test(gpu_size: str = 'large', graphics_backend: str = 'opengl') -> tuple[dict, list[str], str | None]:
    """Test actual GPU graphics rendering; CUDA availability alone is insufficient."""
    report, files, image = _self_test(RenderConfig(engine='EEVEE', device='AUTO', gpu_seconds=45,
                                                gpu_size=gpu_size, graphics_backend=graphics_backend))
    report.update(gpu_size=gpu_size, graphics_backend=graphics_backend)
    report['gpu_eevee_usable'] = bool(report.get('ok') and report.get('gpu_verified'))
    report['message'] = ('当前空间 GPU + Eevee 实际出图成功。' if report['gpu_eevee_usable'] else
                         '当前空间未通过 GPU + Eevee 出图验证。无需在你的电脑安装硬件；'
                         '需要服务器提供可用的 GPU 图形接口；更改 Eevee 选项或安装 CUDA 不能替代该接口。')
    return report, files, image


def _self_test(cfg: RenderConfig) -> tuple[dict, list[str], str | None]:
    backend = _backend(cfg)
    ensure_blender()
    token = secrets.token_hex(24)
    path = JOB_ROOT / ('probe-' + token)
    path.mkdir()
    task = {'action': 'probe', 'project_root': str(path), 'output_dir': str(path / 'frames'),
            'result_path': str(path / 'probe-result.json'), 'config': cfg.model_dump(), 'frames': [1]}
    try:
        result = _execute(task, path, backend, cfg, path / 'probe.log')
        result['deployment_backend'] = backend
        result['gpu_verified'] = bool(result.get('ok')) and backend != 'cpu'
    except Exception as exc:
        result = {'ok': False, 'gpu_verified': False, 'deployment_backend': backend,
                  'error': str(exc), 'message': '渲染自检失败，请查看日志。'}
    result = _public_report(result)
    _atomic_json(path / 'probe-result.json', result)
    files = [_export(path, path / 'probe-result.json')]
    if (path / 'probe.log').exists():
        files.append(_export(path, path / 'probe.log'))
    images = list((path / 'frames').glob('*.png'))
    preview = _export(path, images[0]) if images else None
    shutil.rmtree(path, ignore_errors=True)
    return result, files, preview