"""
Wrapper around COLMAP CLI for feature extraction, matching, and reconstruction
"""
import os
import subprocess

class COLMAPWrapper:
    def __init__(self, colmap_bin='colmap'):
        self.colmap_bin = colmap_bin

    def run_full_pipeline(self, image_dir, workspace_dir):
        # Paths for database and outputs
        db_path = os.path.join(workspace_dir, 'database.db')
        sparse_dir = os.path.join(workspace_dir, 'sparse')
        dense_dir = os.path.join(workspace_dir, 'dense')
        os.makedirs(sparse_dir, exist_ok=True)
        os.makedirs(dense_dir, exist_ok=True)

        # Feature extraction
        subprocess.check_call([
            self.colmap_bin, 'feature_extractor',
            '--database_path', db_path,
            '--image_path', image_dir
        ])

        # Feature matching
        subprocess.check_call([
            self.colmap_bin, 'exhaustive_matcher',
            '--database_path', db_path
        ])

        # Sparse reconstruction (mapping)
        subprocess.check_call([
            self.colmap_bin, 'mapper',
            '--database_path', db_path,
            '--image_path', image_dir,
            '--output_path', sparse_dir
        ])

        # Image undistortion and dense stereo
        subprocess.check_call([
            self.colmap_bin, 'image_undistorter',
            '--image_path', image_dir,
            '--input_path', sparse_dir,
            '--output_path', dense_dir,
            '--output_type', 'COLMAP'
        ])

        subprocess.check_call([
            self.colmap_bin, 'patch_match_stereo',
            '--workspace_path', dense_dir,
            '--workspace_format', 'COLMAP',
            '--PatchMatchStereo.geom_consistency', 'true'
        ])

        # Fuse dense point cloud into mesh
        fused_path = os.path.join(dense_dir, 'fused.ply')
        subprocess.check_call([
            self.colmap_bin, 'stereo_fusion',
            '--workspace_path', dense_dir,
            '--workspace_format', 'COLMAP',
            '--input_type', 'geometric',
            '--output_path', fused_path
        ])
        return fused_path