Async Batch Processing Pipelines for Indoor Mapping & Wayfinding

import asyncio
import aiofiles
import logging
import uuid
import os
from pathlib import Path
from dataclasses import dataclass, field
from typing import AsyncIterator, Dict, Any
from concurrent.futures import ProcessPoolExecutor
import time

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s"
)
logger = logging.getLogger("indoor_pipeline")

# --- Data Structures ---
@dataclass
class TaskManifest:
    file_path: Path
    source_format: str
    floor_id: str
    trace_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    retry_count: int = 0

# --- CPU-Bound Geometry Kernel (Runs in separate process) ---
def process_geometry_kernel(file_path: str, source_format: str) -> Dict[str, Any]:
    """
    Simulates heavy CPU-bound geometry processing.
    Replace with actual Shapely/GeoPandas/IFCopenshell routines.
    """
    logger.info(f"[CPU Kernel] Processing {file_path} ({source_format})")
    time.sleep(0.5)  # Simulate polygonization, snapping, topology validation
    return {
        "status": "success",
        "geometry_count": 142,
        "topology_valid": True,
        "floor_id": Path(file_path).stem.split("_")[0]
    }

# --- Async Pipeline Components ---
class AsyncBatchPipeline:
    def __init__(self, max_workers: int = 4, queue_depth: int = 8):
        self.max_workers = max_workers
        self.queue_depth = queue_depth
        self.task_queue: asyncio.Queue[TaskManifest] = asyncio.Queue(maxsize=queue_depth)
        self.semaphore = asyncio.Semaphore(max_workers)
        self.executor = ProcessPoolExecutor(max_workers=max_workers)
        self.loop = asyncio.get_running_loop()
        self._shutdown_event = asyncio.Event()

    async def ingest_scanner(self, scan_dir: Path) -> None:
        """Coroutine-driven directory scanner with backpressure."""
        logger.info(f"Starting ingestion scan: {scan_dir}")
        for fp in scan_dir.rglob("*"):
            if fp.suffix.lower() in {".dwg", ".svg", ".pdf", ".ifc"}:
                manifest = TaskManifest(
                    file_path=fp,
                    source_format=fp.suffix.lstrip("."),
                    floor_id=fp.stem.split("_")[0]
                )
                # await blocks until queue has space (backpressure)
                await self.task_queue.put(manifest)
                logger.debug(f"Enqueued: {fp.name} | Queue size: {self.task_queue.qsize()}")
        await self.task_queue.put(None)  # Sentinel for graceful shutdown

    async def worker(self, worker_id: int) -> None:
        """Drains queue, respects semaphore, delegates CPU work."""
        while True:
            manifest = await self.task_queue.get()
            
            if manifest is None:
                self.task_queue.task_done()
                break

            async with self.semaphore:
                try:
                    logger.info(f"Worker-{worker_id} | Processing {manifest.file_path.name}")
                    
                    # Delegate CPU-bound work to process pool
                    result = await self.loop.run_in_executor(
                        self.executor,
                        process_geometry_kernel,
                        str(manifest.file_path),
                        manifest.source_format
                    )
                    
                    # Async export / DB write
                    await self._export_to_gis(result, manifest.trace_id)
                    logger.info(f"Worker-{worker_id} | Completed {manifest.file_path.name}")
                    
                except Exception as e:
                    logger.error(f"Worker-{worker_id} | Failed {manifest.file_path.name}: {e}")
                    if manifest.retry_count < 2:
                        manifest.retry_count += 1
                        await self.task_queue.put(manifest)
                finally:
                    self.task_queue.task_done()

    async def _export_to_gis(self, geo_result: Dict[str, Any], trace_id: str) -> None:
        """Simulate async GIS export / topology registration."""
        async with aiofiles.open(f"output_{trace_id}.json", "w") as f:
            await f.write(str(geo_result))

    async def run(self, scan_dir: Path) -> None:
        """Orchestrate pipeline lifecycle."""
        # Start workers
        workers = [asyncio.create_task(self.worker(i)) for i in range(self.max_workers)]
        
        # Start scanner
        scanner = asyncio.create_task(self.ingest_scanner(scan_dir))
        
        # Wait for scanner to finish
        await scanner
        
        # Wait for all queued tasks to complete
        await self.task_queue.join()
        
        # Cancel workers gracefully
        for w in workers:
            w.cancel()
        
        await asyncio.gather(*workers, return_exceptions=True)
        self.executor.shutdown(wait=True)
        logger.info("Pipeline completed successfully.")

# --- Entry Point ---
async def main():
    scan_directory = Path("./floor_plans_input")
    scan_directory.mkdir(exist_ok=True)
    
    # Create dummy files for demonstration
    for i in range(12):
        (scan_directory / f"level_{i}_A.dwg").touch()
        
    pipeline = AsyncBatchPipeline(max_workers=4, queue_depth=8)
    await pipeline.run(scan_directory)

if __name__ == "__main__":
    asyncio.run(main())