#!/usr/bin/env python3

import subprocess
from typing import Dict, Any, List
import json

# Default update interval in seconds
DEFAULT_INTERVAL = 300  # 5 minutes

def get_zfs_pools() -> List[Dict[str, Any]]:
    """Get information about ZFS pools."""
    try:
        # Get list of pools
        pools = subprocess.check_output(['zpool', 'list', '-H', '-o', 'name,size,alloc,free,health,readonly,dedup,altroot']).decode().strip().split('\n')
        
        pool_info = []
        for pool in pools:
            if not pool:  # Skip empty lines
                continue
                
            name, size, alloc, free, health, readonly, dedup, altroot = pool.split('\t')
            
            # Get detailed pool status
            status = subprocess.check_output(['zpool', 'status', name]).decode()
            
            # Get pool properties
            properties = subprocess.check_output(['zpool', 'get', 'all', name]).decode()
            
            pool_info.append({
                'name': name,
                'size': size,
                'allocated': alloc,
                'free': free,
                'health': health,
                'readonly': readonly == 'on',
                'dedup': dedup,
                'altroot': altroot,
                'status': status,
                'properties': properties
            })
            
        return pool_info
    except subprocess.SubprocessError as e:
        print(f"Error getting ZFS pool information: {e}")
        return []

def convert_size_to_bytes(size_str: str) -> int:
    """Convert ZFS size string to bytes."""
    units = {
        'B': 1,
        'K': 1024,
        'M': 1024**2,
        'G': 1024**3,
        'T': 1024**4,
        'P': 1024**5
    }
    
    try:
        number = float(size_str[:-1])
        unit = size_str[-1].upper()
        return int(number * units[unit])
    except (ValueError, KeyError):
        return 0

def collect_metrics() -> Dict[str, Any]:
    """Collect ZFS pool metrics."""
    metrics = {
        "entities": []
    }
    
    pools = get_zfs_pools()
    
    for pool in pools:
        # Pool health status
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_health',
            'name': f'ZFS Pool {pool["name"]} Health',
            'value': pool['health'],
            'state_class': 'measurement',
            'unit_of_measurement': '',
            'device_class': 'enum',
            'icon': 'mdi:database',
            'attributes': {
                'friendly_name': f'ZFS Pool {pool["name"]} Health Status',
                'readonly': pool['readonly'],
                'dedup': pool['dedup'],
                'altroot': pool['altroot']
            }
        })
        
        # Pool size
        size_bytes = convert_size_to_bytes(pool['size'])
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_size',
            'name': f'ZFS Pool {pool["name"]} Size',
            'value': str(round(size_bytes / (1024**3), 2)),  # Convert to GB
            'state_class': 'measurement',
            'unit_of_measurement': 'GB',
            'device_class': 'data_size',
            'icon': 'mdi:database',
            'attributes': {
                'friendly_name': f'ZFS Pool {pool["name"]} Total Size'
            }
        })
        
        # Pool allocated space
        alloc_bytes = convert_size_to_bytes(pool['allocated'])
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_allocated',
            'name': f'ZFS Pool {pool["name"]} Allocated',
            'value': str(round(alloc_bytes / (1024**3), 2)),  # Convert to GB
            'state_class': 'measurement',
            'unit_of_measurement': 'GB',
            'device_class': 'data_size',
            'icon': 'mdi:database',
            'attributes': {
                'friendly_name': f'ZFS Pool {pool["name"]} Allocated Space'
            }
        })
        
        # Pool free space
        free_bytes = convert_size_to_bytes(pool['free'])
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_free',
            'name': f'ZFS Pool {pool["name"]} Free',
            'value': str(round(free_bytes / (1024**3), 2)),  # Convert to GB
            'state_class': 'measurement',
            'unit_of_measurement': 'GB',
            'device_class': 'data_size',
            'icon': 'mdi:database',
            'attributes': {
                'friendly_name': f'ZFS Pool {pool["name"]} Free Space'
            }
        })
        
        # Pool usage percentage
        usage_percent = (alloc_bytes / size_bytes * 100) if size_bytes > 0 else 0
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_usage',
            'name': f'ZFS Pool {pool["name"]} Usage',
            'value': str(round(usage_percent, 1)),
            'state_class': 'measurement',
            'unit_of_measurement': '%',
            'device_class': 'power_factor',
            'icon': 'mdi:database',
            'attributes': {
                'friendly_name': f'ZFS Pool {pool["name"]} Usage Percentage'
            }
        })
    
    return metrics

if __name__ == "__main__":
    # Example usage
    metrics = collect_metrics()
    print(json.dumps(metrics, indent=2))