"""
Example: Dataset encryption workflow.
"""

from fhenomai import FHEnomClient, FHEnomConfig
import logging

logging.basicConfig(level=logging.INFO)

def main():
    # Load configuration
    config = FHEnomConfig.from_file()
    
    with FHEnomClient(config) as client:
        
        print("=" * 60)
        print("Dataset Encryption Workflow")
        print("=" * 60)
        
        # Dataset parameters
        dataset_file = "./datasets/my_dataset.parquet"
        dataset_name = "my_dataset"
        text_fields = ["text", "summary", "prompt"]  # Fields to encrypt
        
        # Upload and encrypt dataset
        encrypted_dataset_id = client.encrypt_and_upload_dataset(
            local_file=dataset_file,
            dataset_name=dataset_name,
            text_fields=text_fields,
            block_size=1024,
            wait=True,
            show_progress=True
        )
        
        print(f"\\n✓ Dataset encrypted: {encrypted_dataset_id}")
        
        # Download encrypted dataset (optional)
        output_file = f"./datasets/{encrypted_dataset_id}.parquet"
        client.download_dataset(
            dataset_name=encrypted_dataset_id,
            local_file=output_file
        )
        
        print(f"✓ Encrypted dataset downloaded to: {output_file}")
        
        # Verify encrypted dataset
        try:
            import pandas as pd
            
            print("\\n" + "=" * 60)
            print("Verifying Encrypted Dataset")
            print("=" * 60)
            
            # Load and inspect
            df = pd.read_parquet(output_file)
            
            print(f"\\nDataset shape: {df.shape}")
            print(f"Columns: {list(df.columns)}")
            print(f"\\nFirst few rows:")
            print(df.head())
            
            # Check encryption
            for field in text_fields:
                if field in df.columns:
                    sample = df[field].iloc[0]
                    print(f"\\n{field} (encrypted): {sample[:100]}...")
        
        except ImportError:
            print("\\nPandas not installed. Install with: pip install pandas pyarrow")
        except Exception as e:
            print(f"\\nVerification error: {e}")
        
        print("\\n" + "=" * 60)
        print("Dataset workflow completed!")
        print("=" * 60)

if __name__ == '__main__':
    main()
