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
    delta_table_path = "delta/test_table"
    
    # Full Azure Blob Storage path
    delta_uri = f"s3://o11y-prod-shared-ap-southeast-1/{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:        
        # Convert pandas DataFrame to PyArrow Table
        table = pa.Table.from_pandas(df)
        
        # Configure storage options
        storage_options = {
            "AWS_ENDPOINT_URL": "https://oss-ap-southeast-1.aliyuncs.com",
            "AWS_REGION": "ap-southeast-1",
            "AWS_COPY_IF_NOT_EXISTS": "header-with-status:x-oss-forbid-overwrite:true:409"
        }
        
        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()