from datetime import datetime
import pandas as pd
import pyarrow as pa
from deltalake import write_deltalake, DeltaTable
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_deltalake_azure_write():
    # Configuration
    storage_account = "o11yprodjapaneast"
    container = "o11y-prod-shared-japaneast"
    delta_table_path = "delta/test_table"
    
    # Full Azure Blob Storage path
    delta_uri = f"abfss://{container}@{storage_account}.dfs.core.windows.net/{delta_table_path}"
    
    # Sample data
    data = {
        "id": [1, 2, 3],
        "name": ["Alice", "Bob", "Charlie"],
        "value": [10.5, 20.3, 15.7],
        "timestamp": [datetime.now().isoformat()] * 3
    }
    df = pd.DataFrame(data)
    
    try:
        logger.info("Setting up Azure credentials...")
        
        # Convert pandas DataFrame to PyArrow Table
        table = pa.Table.from_pandas(df)
        
        # Configure storage options
        storage_options = {
            "AZURE_STORAGE_ACCOUNT_NAME": storage_account,
            "AZURE_STORAGE_USE_EMULATOR": "false"
        }
        
        logger.info(f"Writing Delta table to: {delta_uri}")
        
        # Write the Delta table
        write_deltalake(
            delta_uri,
            table,  # Using PyArrow Table instead of pandas DataFrame
            mode="overwrite",
            storage_options=storage_options,
            partition_by=["timestamp"]
        )
        
        logger.info("Successfully wrote Delta table to Azure Blob Storage")
        
        # Read back to verify
        logger.info("Reading back the Delta table to verify...")
        dt = DeltaTable(delta_uri, storage_options=storage_options)
        df_read = dt.to_pandas()
        logger.info("Read back data:")
        print(df_read)
        
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        raise

if __name__ == "__main__":
    test_deltalake_azure_write()