import argparse
import json
import os
import shutil
import tempfile
import threading
import time
import urllib.error
import urllib.request
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")
HF_ENDPOINT = "https://huggingface.co/"
HF_MIRRORS = ("https://hf-mirror.com",)
HF_PROBE_TIMEOUT = 5.0


class TaskCancelled(RuntimeError):
    pass


def check_cancelled(cancel_check):
    if cancel_check is not None and cancel_check():
        raise TaskCancelled("Task cancelled")


def reachable(url, timeout=HF_PROBE_TIMEOUT):
    request = urllib.request.Request(
        url,
        method="HEAD",
        headers={"User-Agent": "SnapForm-Runtime-Installer"},
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout):
            return True
    except urllib.error.HTTPError:
        return True
    except Exception:
        return False


def resolve_hf_endpoint():
    """Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror)."""
    configured = os.environ.get("HF_ENDPOINT")
    if configured:
        return configured, False
    if reachable(HF_ENDPOINT):
        return HF_ENDPOINT, False
    for mirror in HF_MIRRORS:
        if reachable(mirror):
            return mirror, True
    return HF_ENDPOINT, False


def write_status(
    path,
    progress,
    message,
    error=None,
    state=None,
    stage=None,
    stages=None,
    stage_label=None,
):
    payload = {"progress": progress, "message": message}
    if error:
        payload["error"] = error
    if stage is not None:
        payload["stage"] = stage
        payload["stages"] = stages
        payload["stage_label"] = stage_label
    if state is None:
        if error:
            state = "failed"
        elif progress >= 1.0:
            state = "succeeded"
        else:
            state = "running"
    payload["state"] = state
    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,
    cancel_check=None,
):
    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):
            check_cancelled(cancel_check)
            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, cancel_check=None):
    check_cancelled(cancel_check)
    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,
            cancel_check=cancel_check,
        )
    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,
        start=0.05,
        span=0.90,
        stage=None,
        stages=None,
        stage_label=None,
        label=None,
    ):
        self.status_path = status_path
        self.repository = repository
        self.clock = clock
        self.lock = threading.Lock()
        self.start = start
        self.span = span
        self.stage = stage
        self.stages = stages
        self.stage_label = stage_label
        self.label = label
        self.progress = start
        self.last_write_time = 0.0

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

        factor = min(max(current / total, 0.0), 1.0)
        progress = self.start + factor * self.span
        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
            speed = format_rate(rate) if rate else ""
            if self.stage is not None:
                message = f"downloading {speed}" if speed else "downloading"
            else:
                prefix = self.label or self.repository
                message = f"downloading {prefix}"
                if speed:
                    message = f"{message} {speed}"
            write_status(
                self.status_path,
                progress,
                message,
                stage=self.stage,
                stages=self.stages,
                stage_label=self.stage_label,
            )


def create_download_progress_class(
    status_path,
    repository,
    tqdm_type=None,
    clock=time.monotonic,
    start=0.05,
    span=0.90,
    stage=None,
    stages=None,
    stage_label=None,
    label=None,
):
    if tqdm_type is None:
        from tqdm.auto import tqdm as tqdm_type

    reporter = DownloadProgressReporter(
        status_path,
        repository,
        clock,
        start,
        span,
        stage,
        stages,
        stage_label,
        label,
    )

    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, self._download_rate())
            return result

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

        def _download_rate(self):
            try:
                return getattr(self, "format_dict", {}).get("rate")
            except Exception:
                return None

    return DownloadProgress


def progress_range(args, staged):
    start = getattr(args, "progress_start", None)
    end = getattr(args, "progress_end", None)
    if start is not None and end is not None:
        return float(start), float(end) - float(start)
    if staged:
        return 0.0, 1.0
    return 0.05, 0.90


def download_model(args):
    label = getattr(args, "label", None) or args.repository
    model_dir = Path(args.model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)
    staged = getattr(args, "stage", None) is not None
    start, span = progress_range(args, staged)
    stage_kwargs = {
        "stage": getattr(args, "stage", None),
        "stages": getattr(args, "stages", None),
        "stage_label": getattr(args, "stage_label", None),
    }
    write_status(
        args.status,
        start,
        "checking download source",
        **stage_kwargs,
    )
    endpoint, used_mirror = resolve_hf_endpoint()
    if endpoint != os.environ.get("HF_ENDPOINT"):
        os.environ["HF_ENDPOINT"] = endpoint
    from huggingface_hub import snapshot_download

    message = f"downloading {label}"
    if used_mirror:
        message = f"{message} (mirror)"
    write_status(args.status, start, message, **stage_kwargs)
    progress_class = create_download_progress_class(
        args.status,
        args.repository,
        start=start,
        span=span,
        label=label,
        **stage_kwargs,
    )
    snapshot_download(
        repo_id=args.repository,
        local_dir=model_dir,
        tqdm_class=progress_class,
        allow_patterns=args.allow_pattern or None,
    )
    write_status(
        args.status,
        start + span,
        "model ready",
        **stage_kwargs,
    )


def format_rate(rate):
    size = float(rate)
    if size <= 0:
        return ""
    units = ("B/s", "KB/s", "MB/s", "GB/s")
    unit_index = 0
    while size >= 1024 and unit_index < len(units) - 1:
        size /= 1024
        unit_index += 1
    if unit_index == 0:
        return f"{size:.0f} {units[unit_index]}"
    return f"{size:.1f} {units[unit_index]}"


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")
    if requested == "mps":
        if not torch.backends.mps.is_available():
            raise RuntimeError("MPS was requested but is not available")
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def infer_foreground(
    input_path,
    model_dir,
    device,
    torch,
    model_cache=None,
    cancel_check=None,
):
    from PIL import Image

    input_path = Path(input_path)
    check_cancelled(cancel_check)
    if input_path.suffix.lower() not in IMAGE_SUFFIXES:
        raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
    owns_model = model_cache is None
    if owns_model:
        from ben2 import BEN_Base

        load_started = time.perf_counter()
        model = BEN_Base.from_pretrained(model_dir)
        model = model.to(device=device)
        model.eval()
        load_ms = (time.perf_counter() - load_started) * 1000.0
    else:
        model, load_ms = model_cache.get_ben2(model_dir, device)

    check_cancelled(cancel_check)
    with Image.open(input_path) as opened_image:
        image = opened_image.convert("RGB")
    if device.type == "cuda":
        torch.cuda.synchronize()
    inference_started = time.perf_counter()
    with torch.inference_mode():
        foreground = model.inference(
            image,
            refine_foreground=False,
        )
    if device.type == "cuda":
        torch.cuda.synchronize()
    inference_ms = (time.perf_counter() - inference_started) * 1000.0
    if owns_model:
        del model
        if device.type == "cuda":
            torch.cuda.empty_cache()
    check_cancelled(cancel_check)
    return foreground, load_ms, inference_ms


def remove_background(args, model_cache=None, cancel_check=None):
    total_started = time.perf_counter()
    write_status(args.status, 0.02, "Starting BEN2 Runtime")
    check_cancelled(cancel_check)

    if model_cache is None:
        import torch
    else:
        torch = model_cache.torch

    output_directory = Path(args.output_dir)
    output_directory.mkdir(parents=True, exist_ok=False)
    device = select_device(args.device, torch)

    write_status(args.status, 0.12, f"Loading BEN2 on {device.type}")
    write_status(args.status, 0.35, "Removing background")
    foreground, load_ms, inference_ms = infer_foreground(
        args.input,
        args.model_dir,
        device,
        torch,
        model_cache=model_cache,
        cancel_check=cancel_check,
    )

    check_cancelled(cancel_check)
    write_status(args.status, 0.92, "Writing transparent PNG")
    output_path = output_directory / "foreground.png"
    foreground.convert("RGBA").save(
        output_path,
        format="PNG",
        compress_level=1,
    )
    if not output_path.is_file():
        raise RuntimeError("BEN2 did not write a foreground image")

    check_cancelled(cancel_check)
    total_ms = (time.perf_counter() - total_started) * 1000.0
    message = (
        f"Background removed in {inference_ms:.0f} ms "
        f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
    )
    print(message)
    write_status(args.status, 1.0, message)


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 npz_is_complete(path):
    """Return True only when the npz archive can be fully read."""
    try:
        import numpy as np

        with np.load(path, allow_pickle=False) as archive:
            return "depth" in archive.files
    except Exception:
        return False


def wait_for_exported_paths(
    output_directory,
    export_format,
    timeout=30.0,
    cancel_check=None,
):
    # DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
    # 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        check_cancelled(cancel_check)
        result_paths = exported_paths(output_directory, export_format)
        if result_paths:
            if (
                export_format in ("mini_npz", "npz")
                and not npz_is_complete(result_paths[0])
            ):
                time.sleep(0.1)
                continue
            return result_paths
        time.sleep(0.05)
    return []


def reconstruct(args, model_cache=None, cancel_check=None):
    write_status(args.status, 0.02, "Starting AI runtime")
    check_cancelled(cancel_check)

    if model_cache is None:
        import torch
        from depth_anything_3.api import DepthAnything3
    else:
        torch = model_cache.torch
    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,
        cancel_check=cancel_check,
    )

    try:
        check_cancelled(cancel_check)
        device = select_device(args.device, torch)
        remove_background_model_dir = getattr(
            args,
            "remove_background_model_dir",
            "",
        )
        if remove_background_model_dir:
            write_status(args.status, 0.12, "Removing image background")
            foreground, load_ms, inference_ms = infer_foreground(
                images[0],
                remove_background_model_dir,
                device,
                torch,
                model_cache=model_cache,
                cancel_check=cancel_check,
            )
            check_cancelled(cancel_check)
            foreground.convert("RGBA").save(
                output_directory / "foreground.png",
                format="PNG",
                compress_level=1,
            )
            if not (output_directory / "foreground.png").is_file():
                raise RuntimeError("BEN2 did not write a foreground image")
            print(
                f"BEN2 inference {inference_ms:.0f} ms; "
                f"model load {load_ms:.0f} ms"
            )
        write_status(args.status, 0.15, f"Loading model on {device.type}")
        if model_cache is None:
            model = DepthAnything3.from_pretrained(args.model_dir)
            model = model.to(device=device)
            model.eval()
        else:
            model = model_cache.get_da3(args.model_dir, device)

        check_cancelled(cancel_check)
        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,
            )
        check_cancelled(cancel_check)
        write_status(args.status, 0.88, f"Exporting {args.export_format}")
        export(
            prediction,
            args.export_format,
            str(output_directory),
            **export_kwargs,
        )
        check_cancelled(cancel_check)

        result_paths = wait_for_exported_paths(
            output_directory,
            args.export_format,
            cancel_check=cancel_check,
        )
        if not result_paths:
            raise RuntimeError(
                f"No {args.export_format} files were exported"
            )
        check_cancelled(cancel_check)
        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)
    download.add_argument("--allow-pattern", action="append", default=[])
    download.add_argument("--stage", type=int, default=None)
    download.add_argument("--stages", type=int, default=None)
    download.add_argument("--stage-label", default=None)
    download.add_argument("--label", default=None)
    download.add_argument("--progress-start", type=float, default=None)
    download.add_argument("--progress-end", type=float, default=None)

    background = commands.add_parser("remove-background")
    background.add_argument("--input", required=True)
    background.add_argument("--model-dir", required=True)
    background.add_argument("--output-dir", required=True)
    background.add_argument("--status", required=True)
    background.add_argument(
        "--device",
        choices=("auto", "cuda", "mps", "cpu"),
        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("--remove-background-model-dir", default="")
    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", "mps", "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)
        elif args.command == "reconstruct":
            reconstruct(args)
        else:
            remove_background(args)
    except Exception as error:
        write_status(args.status, 1.0, "Task failed", str(error))
        raise


if __name__ == "__main__":
    main()