#!/usr/bin/env python3
import base64
import argparse
import sys
import os

def extract_signature(base64_data, output_path):
    """
    Extract and save signature image from base64 data
    Handles both data URL format and raw base64 strings
    """
    try:
        # Check if the data is a data URL (starts with data:image)
        if base64_data.startswith('data:image'):
            # Extract just the base64 part by splitting on the comma
            base64_part = base64_data.split(',', 1)[1]
        else:
            # Assume it's already just the base64 data
            base64_part = base64_data
        
        # Decode the base64 data
        image_data = base64.b64decode(base64_part)
        
        # Write to file
        with open(output_path, 'wb') as f:
            f.write(image_data)
        
        print(f"Signature saved to {output_path}")
        return True
    
    except Exception as e:
        print(f"Error extracting signature: {e}")
        return False

def main():
    parser = argparse.ArgumentParser(description='Extract signature image from base64 data')
    parser.add_argument('--output', '-o', required=True, help='Output file path (.png)')
    parser.add_argument('--data', '-d', help='Base64 image data (if not provided, will read from clipboard or stdin)')
    
    args = parser.parse_args()
    
    # Get the base64 data
    base64_data = args.data
    
    # If no data provided via command line, try to get from clipboard
    if not base64_data:
        try:
            # Try to use clipboard (platform specific)
            import pyperclip
            base64_data = pyperclip.paste()
            print("Using data from clipboard")
        except ImportError:
            print("pyperclip not installed. Please install it with: pip install pyperclip")
            print("Trying to read from stdin instead...")
            base64_data = sys.stdin.read().strip()
    
    if not base64_data:
        print("Error: No base64 data provided")
        return
    
    # Extract and save the signature
    extract_signature(base64_data, args.output)

if __name__ == "__main__":
    main()