import argparse
import json
import shutil
import tempfile
import threading
import time
from pathlib import Path


IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = {
    ".mp4",
    ".mov",
    ".avi",
    ".mkv",
    ".webm",
    ".m4v",
    ".mts",
    ".mpg",
    ".mpeg",
}
SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")


def write_status(path, progress, message, error=None):
    payload = {"progress": progress, "message": message}
    if error:
        payload["error"] = error
    Path(path).write_text(
        json.dumps(payload, ensure_ascii=False),
        encoding="utf-8",
    )


def collect_images(directory):
    directory = Path(directory)
    return sorted(
        path
        for path in directory.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )


def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
    capture = cv2.VideoCapture(str(video_path))
    if not capture.isOpened():
        raise RuntimeError(f"Unable to open video: {video_path}")

    output_directory = Path(output_directory)
    output_directory.mkdir(parents=True, exist_ok=True)
    total_frames = int(capture.get(cv2.CAP_PROP_FRAME_COUNT))
    if total_frames <= 0:
        capture.release()
        raise RuntimeError("The video does not report a readable frame count")

    import numpy as np

    sample_count = min(maximum_frames, total_frames)
    frame_indices = np.linspace(
        0,
        total_frames - 1,
        sample_count,
        dtype=np.int64,
    )
    image_paths = []
    try:
        for output_index, frame_index in enumerate(frame_indices):
            capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
            success, frame = capture.read()
            if not success:
                raise RuntimeError(
                    f"Unable to read video frame {int(frame_index)}"
                )
            image_path = output_directory / f"frame-{output_index:04d}.jpg"
            if not cv2.imwrite(str(image_path), frame):
                raise RuntimeError(f"Unable to write video frame: {image_path}")
            image_paths.append(image_path)
    finally:
        capture.release()

    return image_paths


def prepare_input(input_path, video_frames, temporary_root):
    input_path = Path(input_path)
    suffix = input_path.suffix.lower()
    if suffix in IMAGE_SUFFIXES:
        return [input_path], None
    if suffix not in VIDEO_SUFFIXES:
        supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
        raise RuntimeError(f"Unsupported input type; expected one of: {supported}")

    import cv2

    temporary_directory = Path(
        tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
    )
    try:
        images = extract_video_frames(
            input_path,
            temporary_directory,
            video_frames,
            cv2,
        )
    except Exception:
        shutil.rmtree(temporary_directory, ignore_errors=True)
        raise
    return images, temporary_directory


class DownloadProgressReporter:
    def __init__(self, status_path, repository, clock=time.monotonic):
        self.status_path = status_path
        self.repository = repository
        self.clock = clock
        self.lock = threading.Lock()
        self.progress = 0.05
        self.last_write_time = 0.0

    def report(self, current, total):
        if not total or current <= 0:
            return

        factor = min(max(current / total, 0.0), 1.0)
        progress = 0.05 + factor * 0.90
        now = self.clock()
        with self.lock:
            if progress <= self.progress:
                return
            if now - self.last_write_time < 0.20 and progress < 0.95:
                return
            self.progress = progress
            self.last_write_time = now
            write_status(
                self.status_path,
                progress,
                f"Downloading {self.repository}",
            )


def create_download_progress_class(
    status_path,
    repository,
    tqdm_type=None,
    clock=time.monotonic,
):
    if tqdm_type is None:
        from tqdm.auto import tqdm as tqdm_type

    reporter = DownloadProgressReporter(status_path, repository, clock)

    class DownloadProgress(tqdm_type):
        def __init__(self, *args, **kwargs):
            description = str(kwargs.get("desc", ""))
            self.snapform_reports_progress = (
                description == "Downloading bytes"
                or description.startswith("Fetching ")
            )
            super().__init__(*args, **kwargs)

        def update(self, amount=1):
            result = super().update(amount)
            if self.snapform_reports_progress:
                reporter.report(self.n, self.total)
            return result

        def refresh(self, *args, **kwargs):
            result = super().refresh(*args, **kwargs)
            if self.snapform_reports_progress:
                reporter.report(self.n, self.total)
            return result

    return DownloadProgress


def download_model(args):
    from huggingface_hub import snapshot_download

    model_dir = Path(args.model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)
    write_status(args.status, 0.05, f"Downloading {args.repository}")
    progress_class = create_download_progress_class(
        args.status,
        args.repository,
    )
    snapshot_download(
        repo_id=args.repository,
        local_dir=model_dir,
        tqdm_class=progress_class,
    )
    write_status(args.status, 1.0, "Model is ready")


def select_device(requested, torch):
    if requested == "cpu":
        return torch.device("cpu")
    if requested == "cuda":
        if not torch.cuda.is_available():
            raise RuntimeError("CUDA was requested but is not available")
        return torch.device("cuda")
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def exported_paths(output_directory, export_format):
    output_directory = Path(output_directory)
    if export_format == "glb":
        scene_path = output_directory / SCENE_FILENAME
        return [scene_path] if scene_path.is_file() else []
    if export_format == "gs_ply":
        return sorted((output_directory / "gs_ply").glob("*.ply"))
    if export_format == "depth_vis":
        return sorted((output_directory / "depth_vis").glob("*.jpg"))
    if export_format == "mini_npz":
        result_path = output_directory / "exports" / "mini_npz" / "results.npz"
        return [result_path] if result_path.is_file() else []
    return []


def wait_for_exported_paths(output_directory, export_format, timeout=30.0):
    # 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        result_paths = exported_paths(output_directory, export_format)
        if result_paths:
            return result_paths
        time.sleep(0.05)
    return []


def reconstruct(args):
    write_status(args.status, 0.02, "Starting DA3 runtime")

    import torch
    from depth_anything_3.api import DepthAnything3
    from depth_anything_3.utils.export import export

    output_directory = Path(args.output_dir)
    output_directory.mkdir(parents=True, exist_ok=False)
    write_status(args.status, 0.08, "Preparing input")
    images, temporary_directory = prepare_input(
        args.input,
        args.video_frames,
        output_directory.parent,
    )

    try:
        device = select_device(args.device, torch)
        write_status(args.status, 0.15, f"Loading model on {device.type}")
        model = DepthAnything3.from_pretrained(args.model_dir)
        model = model.to(device=device)
        model.eval()

        write_status(
            args.status,
            0.25,
            f"Estimating depth for {len(images)} frame(s)",
        )
        export_kwargs = {}
        if args.export_format == "glb":
            export_kwargs["glb"] = {
                "conf_thresh_percentile": args.conf_thresh_percentile,
                "num_max_points": args.num_max_points,
                "show_cameras": bool(args.show_cameras),
                "export_depth_vis": False,
            }
        if args.export_format == "gs_ply":
            export_kwargs["gs_ply"] = {
                "gs_views_interval": args.gs_views_interval,
            }
        with torch.inference_mode():
            prediction = model.inference(
                [str(path) for path in images],
                infer_gs=args.export_format == "gs_ply",
                process_res=args.process_res,
                process_res_method=args.process_res_method,
                use_ray_pose=bool(args.use_ray_pose),
                ref_view_strategy=args.ref_view_strategy,
            )
        write_status(args.status, 0.88, f"Exporting {args.export_format}")
        export(
            prediction,
            args.export_format,
            str(output_directory),
            **export_kwargs,
        )

        result_paths = wait_for_exported_paths(
            output_directory,
            args.export_format,
        )
        if not result_paths:
            raise RuntimeError(
                f"DA3 did not export any {args.export_format} files"
            )
        write_status(
            args.status,
            1.0,
            f"Exported {len(result_paths)} {args.export_format} file(s)",
        )
    finally:
        if temporary_directory is not None:
            shutil.rmtree(temporary_directory, ignore_errors=True)


def parser():
    root = argparse.ArgumentParser()
    commands = root.add_subparsers(dest="command", required=True)

    download = commands.add_parser("download")
    download.add_argument("--repository", required=True)
    download.add_argument("--model-dir", required=True)
    download.add_argument("--status", required=True)

    reconstruction = commands.add_parser("reconstruct")
    reconstruction.add_argument("--input", required=True)
    reconstruction.add_argument("--model-dir", required=True)
    reconstruction.add_argument("--output-dir", required=True)
    reconstruction.add_argument(
        "--export-format",
        choices=EXPORT_FORMATS,
        required=True,
    )
    reconstruction.add_argument("--status", required=True)
    reconstruction.add_argument("--process-res", type=int, required=True)
    reconstruction.add_argument(
        "--process-res-method",
        choices=("upper_bound_resize", "lower_bound_resize"),
        required=True,
    )
    reconstruction.add_argument(
        "--conf-thresh-percentile",
        type=float,
        required=True,
    )
    reconstruction.add_argument("--num-max-points", type=int, required=True)
    reconstruction.add_argument(
        "--show-cameras",
        type=int,
        choices=(0, 1),
        required=True,
    )
    reconstruction.add_argument(
        "--gs-views-interval",
        type=int,
        required=True,
    )
    reconstruction.add_argument("--video-frames", type=int, required=True)
    reconstruction.add_argument(
        "--device",
        choices=("auto", "cuda", "cpu"),
        required=True,
    )
    reconstruction.add_argument(
        "--use-ray-pose",
        type=int,
        choices=(0, 1),
        required=True,
    )
    reconstruction.add_argument(
        "--ref-view-strategy",
        choices=("first", "middle", "saddle_balanced", "saddle_sim_range"),
        required=True,
    )
    return root


def main():
    args = parser().parse_args()
    try:
        if args.command == "download":
            download_model(args)
        else:
            reconstruct(args)
    except Exception as error:
        write_status(args.status, 1.0, "DA3 task failed", str(error))
        raise


if __name__ == "__main__":
    main()