import os
from contextlib import contextmanager
from typing import Dict, Optional
from .lakefs import setup_lakefs

@contextmanager
def sandbox(repository: str, service_account: Optional[str] = None):
    """
    Context manager that provisions a scoped LakeFS environment and 
    temporarily overrides environment variables.
    """
    # 1. Provision the environment
    creds = setup_lakefs(repository, service_account)
    
    # 2. Save original env state
    original_env = {k: os.environ.get(k) for k in creds.keys()}
    
    # 3. Apply scoped credentials
    try:
        for k, v in creds.items():
            os.environ[k] = v
        yield creds
    finally:
        # 4. Restore original environment
        for k, v in original_env.items():
            if v is None:
                del os.environ[k]
            else:
                os.environ[k] = v
