# -*- mode: python ; coding: utf-8 -*-

import os
from PyInstaller.utils.hooks import (
    collect_data_files,
    collect_dynamic_libs,
    collect_submodules,
)


def _safe_collect_submodules(name: str):
    try:
        return collect_submodules(name)
    except Exception:
        return []


def _safe_collect_data_files(name: str):
    try:
        return collect_data_files(name)
    except Exception:
        return []


def _safe_collect_dynamic_libs(name: str):
    try:
        return collect_dynamic_libs(name)
    except Exception:
        return []


project_root = r"e:\code\vismatch-main\vismatch-main"
test_dir = os.path.join(project_root, "test")
script_path = os.path.join(test_dir, "StripStitch.py")

# Find the actual vismatch installation location (usually in site-packages)
import vismatch as _vismatch_pkg
vismatch_sitepkg_root = os.path.dirname(_vismatch_pkg.__file__)


hiddenimports = []
hiddenimports += _safe_collect_submodules("vismatch")
hiddenimports += _safe_collect_submodules("rasterio")
hiddenimports += _safe_collect_submodules("rasterio._base")
hiddenimports += _safe_collect_submodules("rasterio._io")
hiddenimports += _safe_collect_submodules("affine")
hiddenimports += _safe_collect_submodules("cv2")
hiddenimports += ["tif_caijain"]

hiddenimports += _safe_collect_submodules("pyproj")
hiddenimports += _safe_collect_submodules("scipy")
hiddenimports += _safe_collect_submodules("skimage")
hiddenimports += _safe_collect_submodules("SimpleITK")
hiddenimports += _safe_collect_submodules("pirt")
hiddenimports += _safe_collect_submodules("loguru")

# MatchAnything's src module - need to collect from third_party directory
# Add the src module path to pathex so PyInstaller can find it during analysis
matchanything_src_dir = os.path.join(vismatch_sitepkg_root, "third_party", "MatchAnything", "imcui", "third_party", "MatchAnything")
if os.path.isdir(matchanything_src_dir):
    # Add to pathex for analysis
    pass  # Will be added to pathex in Analysis
    # Also try to collect src submodules if they exist
    try:
        hiddenimports += _safe_collect_submodules("src")
    except Exception:
        pass


datas = []
datas += _safe_collect_data_files("vismatch")
datas += _safe_collect_data_files("rasterio")
datas += _safe_collect_data_files("pyproj")

# vismatch 的 third_party 下包含大量运行时动态 add_to_path 的源码（例如 matchanything 依赖的 src.*）。
# 这些 .py 文件不会被 collect_data_files 收集，需作为 datas 复制到 dist 里，并保持相对路径结构，
# 使得运行时 THIRD_PARTY_DIR = Path(vismatch.__file__).parent/"third_party" 可找到它们.
# Use the actual vismatch installation location (site-packages) instead of project root
# NOTE: Analysis() expects hook-style 2-tuples (source_dir, dest_dir), not Tree() TOC 3-tuples.
third_party_dir = os.path.join(vismatch_sitepkg_root, "third_party")
if os.path.isdir(third_party_dir):
    datas.append((third_party_dir, os.path.join("vismatch", "third_party")))

# Include HuggingFace model weights for offline use (~/.cache/huggingface/hub/)
hf_cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "hub")
if os.path.isdir(hf_cache_dir):
    for model_dir in os.listdir(hf_cache_dir):
        if "vismatch" in model_dir.lower():
            full_model_path = os.path.join(hf_cache_dir, model_dir)
            if os.path.isdir(full_model_path):
                datas.append((full_model_path, os.path.join("hub", model_dir)))


binaries = []
binaries += _safe_collect_dynamic_libs("rasterio")
binaries += _safe_collect_dynamic_libs("pyproj")
binaries += _safe_collect_dynamic_libs("cv2")

# Build pathex - include MatchAnything src directory for proper import analysis
pathex = [project_root, test_dir]
matchanything_src_dir = os.path.join(vismatch_sitepkg_root, "third_party", "MatchAnything", "imcui", "third_party", "MatchAnything")
if os.path.isdir(matchanything_src_dir):
    pathex.append(matchanything_src_dir)
    # Also add ROMA if it exists
    roma_dir = os.path.join(vismatch_sitepkg_root, "third_party", "MatchAnything", "imcui", "third_party", "MatchAnything", "third_party", "ROMA")
    if os.path.isdir(roma_dir):
        pathex.append(roma_dir)

a = Analysis(
    [script_path],
    pathex=pathex,
    binaries=binaries,
    datas=datas,
    hiddenimports=hiddenimports,
    hookspath=[],
    hooksconfig={},
    runtime_hooks=[],
    excludes=[],
    noarchive=False,
)

pyz = PYZ(a.pure)

exe = EXE(
    pyz,
    a.scripts,
    [],
    exclude_binaries=True,
    name="StripStitch",
    debug=False,
    bootloader_ignore_signals=False,
    strip=False,
    upx=False,
    upx_exclude=[],
    runtime_tmpdir=None,
    console=False,
    disable_windowed_traceback=False,
    argv_emulation=False,
    target_arch=None,
    codesign_identity=None,
    entitlements_file=None,
)

coll = COLLECT(
    exe,
    a.binaries,
    a.zipfiles,
    a.datas,
    strip=False,
    upx=False,
    upx_exclude=[],
    name="StripStitch",
)