"""Input preparation, checkpoint selection, inference, and portable result bundles."""
from dataclasses import dataclass
import gc
import json
import math
import os
from pathlib import Path
import secrets
import shutil
import tempfile
import threading
import time
import warnings
import zipfile

import numpy as np
from PIL import Image

from src.utils.camera_io import IMAGE_EXTENSIONS, list_image_files, load_cameras_json

ROOT = Path(__file__).resolve().parents[1]
RGB_MODE = 'RGB-only'
POSE_MODE = 'Pose-conditioned'


@dataclass(frozen=True)
class Options:
    mode: str = RGB_MODE
    width: int = 896
    height: int = 896
    opacity: float = 0.05
    sky: bool = False
    align_cameras: bool = True
    max_images: int | None = None
    voxel_size: float = 0.0018
    max_scale_pct: float = 99.5
    filter_flying_points_infer: bool = False

    def validate(self):
        if not math.isfinite(self.max_scale_pct) or not 0 < self.max_scale_pct <= 100:
            raise ValueError('max_scale_pct must be greater than 0 and at most 100.')
        if not math.isfinite(self.voxel_size) or not 0.0015 <= self.voxel_size <= 0.0020:
            raise ValueError('Voxel size must be between 0.0015 and 0.0020.')
        if self.max_images is not None and (not isinstance(self.max_images, int) or self.max_images < 0):
            raise ValueError('Max images must be 0 (all) or a positive integer.')
        if self.mode not in (RGB_MODE, POSE_MODE):
            raise ValueError('Choose RGB-only or Pose-conditioned inference.')
        for value in (self.width, self.height):
            if not isinstance(value, int) or value <= 0 or value % 14:
                raise ValueError('Width and height must be positive multiples of 14.')
        if not math.isfinite(self.opacity) or not 0 <= self.opacity <= 1:
            raise ValueError('Opacity must be between 0 and 1.')


def ordered_images(files, minimum=1):
    paths = sorted((Path(p) for p in files or []), key=lambda p: p.name)
    if len(paths) < minimum:
        raise ValueError('Upload at least one image of your scene.')
    if any(not p.is_file() or p.suffix.lower() not in IMAGE_EXTENSIONS for p in paths):
        raise ValueError('Use supported image files: PNG, JPEG, WebP, BMP, or TIFF.')
    stems = [p.stem for p in paths]
    if len(stems) != len(set(stems)):
        raise ValueError('Image filename stems must be unique, including across file extensions.')
    return paths


def uploaded_resolution(files):
    """Default to 896 square without exceeding any uploaded image dimension."""
    width = height = 896
    for path in ordered_images(files, minimum=0):
        with Image.open(path) as image:
            width, height = min(width, image.width), min(height, image.height)
    width, height = width // 14 * 14, height // 14 * 14
    if min(width, height) < 14:
        raise ValueError('Each uploaded image must be at least 14 pixels wide and high.')
    return width, height


def stage_inputs(files, camera, options, destination):
    """Preserve filenames so camera IDs and exported image IDs remain consistent."""
    options.validate()
    paths = ordered_images(files)
    if options.max_images:
        paths = paths[:options.max_images]
    destination = Path(destination)
    images = destination / 'images'
    images.mkdir(parents=True, exist_ok=False)
    if options.mode == POSE_MODE:
        if not camera:
            raise ValueError('Pose-conditioned inference requires prior_camera.json.')
        extrinsics, intrinsics = load_cameras_json(camera)
        for path in paths:
            if path.stem not in extrinsics or path.stem not in intrinsics:
                raise ValueError(f'Camera JSON is missing poses or intrinsics for {path.stem}.')
            with Image.open(path) as image:
                cx, cy = intrinsics[path.stem][:2, 2]
                if not (0 < cx < image.width and 0 < cy < image.height):
                    raise ValueError(f'Camera principal point lies outside image {path.name}.')
        camera_path = destination / 'prior_camera.json'
        shutil.copyfile(camera, camera_path)
    else:
        # Switching back to RGB-only never silently uses a previously uploaded prior.
        camera_path = None
    for path in paths:
        with Image.open(path) as image:
            image.verify()
        shutil.copyfile(path, images / path.name)
    return images, camera_path


def extract_video_frames(video, count, destination):
    import cv2
    if not video or not Path(video).is_file():
        raise ValueError('Upload a video first.')
    if not 2 <= int(count) <= 32:
        raise ValueError('Extract between 2 and 32 frames.')
    capture = cv2.VideoCapture(str(video))
    destination = Path(destination)
    destination.mkdir(parents=True, exist_ok=False)
    try:
        total = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
        if total < 2:
            raise ValueError('The video could not be decoded or has fewer than two frames.')
        paths = []
        for index in np.linspace(0, total - 1, min(int(count), total), dtype=int):
            capture.set(cv2.CAP_PROP_POS_FRAMES, int(index))
            ok, frame = capture.read()
            if not ok:
                continue
            path = destination / f'frame_{index:08d}.jpg'
            if cv2.imwrite(str(path), frame):
                paths.append(str(path))
        if len(paths) < 2:
            raise ValueError('Fewer than two frames could be decoded; try a different video.')
        return paths
    finally:
        capture.release()


def inspect_panorama(path):
    """Validate the same ERP convention consumed by scripts/convert_panorama.py."""
    if not path or not Path(path).is_file():
        raise ValueError('Upload one full 360° panorama first.')
    try:
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', Image.DecompressionBombWarning)
            with Image.open(path) as image:
                if image.format not in ('JPEG', 'PNG', 'WEBP'):
                    raise ValueError('Use a JPG, PNG, or WebP panorama.')
                if image.width != 2 * image.height:
                    raise ValueError('Use a full 360° × 180° equirectangular panorama with a 2:1 aspect ratio.')
                if image.getexif().get(274, 1) != 1:
                    raise ValueError('Apply EXIF orientation to the panorama before uploading it.')
                size = image.size
            # Reading PNG EXIF may load the image; verify a freshly opened stream.
            with Image.open(path) as image:
                image.verify()
            return size
    except (OSError, Image.DecompressionBombError) as error:
        raise ValueError('The panorama could not be decoded. Upload a valid JPG, PNG, or WebP image.') from error


def resolve_weights(config_path, student=None, geometry=None):
    from omegaconf import OmegaConf
    cfg = OmegaConf.load(config_path)
    repo = os.environ.get('SPLATANYTHING_MODEL_REPO')
    revision = os.environ.get('SPLATANYTHING_MODEL_REVISION', 'main')
    choices = [
        ('SplatAnything_pretrained', student or os.environ.get('SPLATANYTHING_CKPT'),
         'SPLATANYTHING_MODEL_FILE'),
        ('worldmirror2_pretrained', geometry or os.environ.get('WORLDMIRROR2_CKPT'),
         'SPLATANYTHING_GEOMETRY_FILE'),
    ]
    for key, override, file_var in choices:
        path = Path(override or cfg.wrapper[key]).expanduser()
        if not path.is_absolute():
            path = (Path.cwd() if override else ROOT) / path
        if not path.is_file() and repo and not override:
            from huggingface_hub import hf_hub_download
            filename = os.environ.get(file_var, Path(cfg.wrapper[key]).name)
            path = Path(hf_hub_download(repo, filename, revision=revision))
        if not path.is_file():
            raise FileNotFoundError(
                f'{key} not found: {path}. Configure local checkpoints, '
                'SPLATANYTHING_CKPT / WORLDMIRROR2_CKPT, or SPLATANYTHING_MODEL_REPO.')
        # Hub snapshot links point to extensionless blobs. Keep the logical
        # filename so safetensors is decoded with the correct loader.
        cfg.wrapper[key] = str(path.absolute())
    return cfg


def to_cpu(value):
    import torch
    if isinstance(value, torch.Tensor):
        return value.detach().cpu()
    if isinstance(value, dict):
        return {key: to_cpu(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return type(value)(to_cpu(item) for item in value)
    return value


class Pipeline:
    def __init__(self, config_path=None, device='cuda', student=None, geometry=None, sky_path=None):
        self.config_path = Path(config_path or ROOT / 'configs/inference.yaml')
        self.device = device
        self.student_path, self.geometry_path = student, geometry
        self._explicit_sky_path = sky_path or os.environ.get('SPLATANYTHING_SKY_CKPT')
        self.sky_path = self._explicit_sky_path or str(ROOT / 'checkpoints/skyseg.onnx')
        self.model = self.teacher = self.segmenter = self.config = None
        self.lock = threading.RLock()

    def initialize(self):
        with self.lock:
            self._initialize()

    def _initialize(self):
        """Called before GPU scheduling in ZeroGPU and lazily for ordinary local runs."""
        if self.model is not None:
            return
        import torch
        from omegaconf import OmegaConf
        from src.models.models.gaussian_worldmirror import GaussianWorldMirror
        from src.models.models.geometry_worldmirror import GeometryWorldMirror
        from src.utils.checkpoints import load_model_weights
        cfg = resolve_weights(self.config_path, self.student_path, self.geometry_path)
        if str(self.device).startswith('cuda') and not torch.cuda.is_available():
            raise RuntimeError('CUDA is unavailable. Select a GPU runtime or use --device cpu for a small smoke test.')
        # Keep the conditioning layers resident, while selecting the exact CLI token
        # layout per request. A shared lock protects mode and renderer state.
        model = GaussianWorldMirror(**OmegaConf.to_container(cfg.wrapper.model), enable_cond=True)
        load_model_weights(model, cfg.wrapper.SplatAnything_pretrained)
        model.to(self.device).eval().requires_grad_(False)
        teacher = GeometryWorldMirror(enable_cond=True)
        load_model_weights(teacher, cfg.wrapper.worldmirror2_pretrained)
        teacher.to(self.device).eval().requires_grad_(False)
        self.model, self.teacher, self.config = model, teacher, cfg
        gc.collect()

    def _set_mode(self, enabled):
        for model in (self.model, self.teacher):
            model.enable_cond = enabled
            encoder = model.visual_geometry_transformer
            if not hasattr(encoder, '_demo_conditioned_patch_start'):
                encoder._demo_conditioned_patch_start = encoder.patch_start_idx
            encoder.enable_cond = enabled
            encoder.patch_start_idx = encoder._demo_conditioned_patch_start - (0 if enabled else 2)
        self.model.gs_renderer.use_gt_camera = enabled

    def prepare_inputs(self, images, camera, options):
        """Run CPU image preparation and sky segmentation before GPU scheduling."""
        from src.utils.scene_io import build_views
        from src.utils.sky_segmentation import SkySegmenter
        options.validate()
        with self.lock:
            if options.sky and self.segmenter is None:
                path = Path(self.sky_path).expanduser()
                if not path.is_file():
                    if self._explicit_sky_path:
                        raise ValueError('Sky segmentation model not found at the configured SPLATANYTHING_SKY_CKPT path.')
                    from huggingface_hub import hf_hub_download
                    model_repo = os.environ.get('SPLATANYTHING_MODEL_REPO')
                    repo = os.environ.get('SPLATANYTHING_SKY_REPO', model_repo or 'ritianyu/SplatAnything')
                    revision = os.environ.get('SPLATANYTHING_SKY_REVISION') or (
                        os.environ.get('SPLATANYTHING_MODEL_REVISION', 'main') if repo == model_repo else 'main')
                    filename = os.environ.get('SPLATANYTHING_SKY_FILE', 'skyseg.onnx')
                    try:
                        path = Path(hf_hub_download(repo, filename, revision=revision))
                    except Exception as error:
                        raise ValueError(f'Could not download the sky segmentation model from {repo}. Check the sky model settings.') from error
                self.sky_path = str(path.absolute())
                self.segmenter = SkySegmenter(self.sky_path, device='cpu')
            if options.sky:
                self.segmenter._ensure_session()
            views, _, paths = build_views(
                str(images), str(camera) if options.mode == POSE_MODE else None,
                (options.width, options.height), 'cpu', cam_align=options.align_cameras,
                sky_segmenter=self.segmenter if options.sky else None)
            return views, paths

    def predict(self, images, camera, options, prepared=None):
        import torch
        from src.utils.scene_io import build_views
        with self.lock:
            self.initialize()
            enabled = options.mode == POSE_MODE
            self._set_mode(enabled)
            options.validate()
            self.model.gs_renderer.voxel_size = options.voxel_size
            self.model.gs_renderer.max_scale_pct = options.max_scale_pct
            self.model.gs_renderer.filter_flying_points_infer = options.filter_flying_points_infer
            if options.sky and prepared is None:
                prepared = self.prepare_inputs(images, camera, options)
            if prepared is not None:
                cpu_views, paths = prepared
                views = {key: value.to(self.device) for key, value in cpu_views.items()}
            else:
                views, _, paths = build_views(
                    str(images), str(camera) if enabled else None,
                    (options.width, options.height), self.device,
                    cam_align=options.align_cameras)
            flags = [1, 0, 1] if enabled else [0, 0, 0]
            self.model.gs_renderer.filter_sky_infer = options.sky
            try:
                with torch.inference_mode():
                    geometry = self.teacher(views, cond_flags=flags)
                    predictions = self.model(views, geometry, cond_flags=flags)
                result = to_cpu(predictions), to_cpu(views), paths
                del geometry, predictions, views
                return result
            finally:
                # Token/rotary caches are device-local and must not retain CUDA
                # allocations across a ZeroGPU reservation or a mode switch.
                for model in (self.model, self.teacher):
                    encoder = model.visual_geometry_transformer
                    if hasattr(encoder, 'pos_getter'):
                        encoder.pos_getter.position_cache.clear()
                    for module in model.modules():
                        if hasattr(module, 'frequency_cache'):
                            module.frequency_cache.clear()
                if str(self.device).startswith('cuda'):
                    torch.cuda.empty_cache()


def export_preview(predictions, paths, output, options, config):
    """Write only the Gaussians and small metadata needed before showing the viewer."""
    from omegaconf import OmegaConf
    from src.utils.scene_io import export_gaussians
    output = Path(output)
    output.mkdir(parents=True, exist_ok=True)
    splats = predictions['splats']
    if not len(splats['means'][0]) or not (splats['opacities'][0] >= options.opacity).any():
        raise ValueError('No Gaussians remain after filtering. Lower opacity or use more overlapping views.')
    export_gaussians(predictions, output, options.opacity, None)
    cameras = {'extrinsics': [], 'intrinsics': []}
    for i, image in enumerate(paths):
        for name, key in [('extrinsics', 'export_camera_poses'), ('intrinsics', 'export_camera_intrs')]:
            cameras[name].append({'camera_id': Path(image).stem,
                                  'matrix': predictions[key][0, i].tolist()})
    (output / 'cameras.json').write_text(json.dumps(cameras, indent=2) + '\n')
    # Downloaded metadata contains portable filenames rather than server paths.
    public_config = OmegaConf.create(OmegaConf.to_container(config))
    for key in ('worldmirror2_pretrained', 'SplatAnything_pretrained'):
        public_config.wrapper[key] = Path(public_config.wrapper[key]).name
    public_config.wrapper.enable_cond = options.mode == POSE_MODE
    public_config.wrapper.model.use_gt_camera = options.mode == POSE_MODE
    public_config.wrapper.infer_opacity_thresh = options.opacity
    public_config.wrapper.model.gs_voxel_size = options.voxel_size
    public_config.wrapper.model.max_scale_pct = options.max_scale_pct
    public_config.wrapper.model.filter_flying_points_infer = options.filter_flying_points_infer
    OmegaConf.save(public_config, output / 'config.yaml')
    metadata = dict(mode=options.mode, resolution=[options.width, options.height],
                    images=[Path(p).name for p in paths], opacity=options.opacity,
                    sky_seg=options.sky, camera_alignment=options.align_cameras,
                    max_images=options.max_images, voxel_size=options.voxel_size,
                    max_scale_pct=options.max_scale_pct,
                    filter_flying_points_infer=options.filter_flying_points_infer)
    (output / 'run.json').write_text(json.dumps(metadata, indent=2) + '\n')
    return output / 'gaussians.ply'


def export_auxiliary(predictions, views, paths, output):
    """Export full-resolution raw points, depth and RGB after the viewer is ready.

    Paths are used only for image names: staged input folders may already be gone.
    All tensors passed here must be CPU tensors, independent of the GPU request.
    """
    from src.utils.scene_io import export_points, export_depth_rgb
    export_points(predictions, Path(output))
    export_depth_rgb(predictions, views, paths, Path(output))


def export_result(predictions, views, paths, output, options, config):
    """Synchronous export for callers that explicitly need the complete result."""
    ply = export_preview(predictions, paths, output, options, config)
    export_auxiliary(predictions, views, paths, output)
    return ply


def package_result(output):
    output = Path(output)
    archive = output / 'result.zip'
    with tempfile.NamedTemporaryFile(dir=output.parent, suffix='.zip', delete=False) as temporary:
        staged = Path(temporary.name)
    try:
        with zipfile.ZipFile(staged, 'w', zipfile.ZIP_DEFLATED, compresslevel=1) as handle:
            for path in sorted(output.rglob('*')):
                if path.is_file() and path != archive and path.name not in {'preview.compressed.ply', 'viewer-settings.json'}:
                    compressed = path.suffix.lower() in {'.png', '.jpg', '.jpeg', '.webp'}
                    handle.write(path, path.relative_to(output),
                                 compress_type=zipfile.ZIP_STORED if compressed else zipfile.ZIP_DEFLATED)
        staged.replace(archive)
    finally:
        staged.unlink(missing_ok=True)
    return archive


class ResultStore:
    """Unpredictable per-result links, with a bounded lifetime and no arbitrary paths."""
    def __init__(self, root=None, lifetime=24 * 3600):
        self.root = Path(root or tempfile.mkdtemp(prefix='splatanything-results-')).resolve()
        self.root.mkdir(parents=True, exist_ok=True)
        self.lifetime = lifetime
        self.records = {}
        self.exports = {}
        self.lock = threading.Lock()

    def reserve(self):
        with self.lock:
            now = time.time()
            for token, (path, expiry) in list(self.records.items()):
                task = self.exports.get(token)
                if expiry < now and (task is None or task['done'].is_set()):
                    shutil.rmtree(path, ignore_errors=True)
                    del self.records[token]
                    self.exports.pop(token, None)
            token = secrets.token_urlsafe(32)
            path = self.root / secrets.token_hex(16)
            path.mkdir()
            self.records[token] = path, now + self.lifetime
            return token, path

    def begin_export(self, token):
        with self.lock:
            self.exports[token] = {'done': threading.Event(), 'html_done': threading.Event(), 'error': False}

    def finish_html(self, token):
        with self.lock:
            self.exports[token]['html_done'].set()

    def wait_for_html(self, token):
        with self.lock:
            task = self.exports.get(token)
        if task is not None:
            task['html_done'].wait()

    def finish_export(self, token, failed=False):
        with self.lock:
            task = self.exports[token]
            task['error'] = failed
            task['html_done'].set()
            task['done'].set()

    def export_status(self, token):
        with self.lock:
            record = self.records.get(token)
            if record is None or record[1] < time.time():
                return None
            task = self.exports.get(token)
            if task is None:
                return 'ready'
            if not task['done'].is_set():
                return 'preparing'
            return 'failed' if task['error'] else 'ready'

    def wait_for_export(self, token):
        with self.lock:
            task = self.exports.get(token)
        if task is not None:
            task['done'].wait()
        return self.export_status(token)

    def resolve(self, token, filename):
        if filename not in {'scene.html', 'gaussians.ply', 'result.zip', 'cameras.json',
                            'preview.compressed.ply', 'viewer-settings.json'}:
            return None
        with self.lock:
            record = self.records.get(token)
            if record is None or record[1] < time.time():
                return None
            path = record[0] / filename
            return path if path.is_file() else None