"""
Patch for Thrift-generated exception classes to make them compatible with Airflow's secret masking.
"""

import logging
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

# --- Python Path Setup ---
project_root = Path(__file__).parent.absolute()
# Add project root to sys.path (needed for the 'pangramia' symlink)
if str(project_root) not in sys.path: sys.path.insert(0, str(project_root))
# --- End Python Path Setup ---

logger = logging.getLogger(__name__)

def patch_thrift_exceptions():
    """
    Patch Thrift-generated exception classes to make them compatible with Airflow's secret masking.
    """
    try:
        from pangramia.yt.exceptions.ttypes import PBServiceException, PBUserException
        
        # Save original __setattr__ methods
        original_service_setattr = PBServiceException.__setattr__
        original_user_setattr = PBUserException.__setattr__
        
        # Define a new __setattr__ method that allows modifying any attribute
        def new_service_setattr(self, name, value):
            logger.debug(f"Setting attribute {name} on PBServiceException")
            object.__setattr__(self, name, value)
                
        def new_user_setattr(self, name, value):
            logger.debug(f"Setting attribute {name} on PBUserException")
            object.__setattr__(self, name, value)
        
        # Apply the patch to both exception classes
        PBServiceException.__setattr__ = new_service_setattr
        PBUserException.__setattr__ = new_user_setattr
        
        logger.info("Successfully patched Thrift exception classes for Airflow compatibility")
        
        # Verify the patch
        try:
            test_exception = PBServiceException(message="Test")
            test_exception.args = ("Test",)  # Try to modify an attribute
            logger.info("Verified Thrift exception patch is working correctly")
        except Exception as e:
            logger.error(f"Thrift exception patch verification failed: {e}")
    except ImportError as e:
        logger.warning(f"Could not import Thrift exception classes: {e}")
        logger.warning("Airflow error handling may not work properly with Thrift exceptions")
    except Exception as e:
        logger.error(f"Error patching Thrift exception classes: {e}")

# Apply the patch when this module is imported
patch_thrift_exceptions()