#!/usr/bin/env python3

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

# 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. Skips cleanly if ZFS is unavailable."""
    metrics = {"entities": []}

    # Check binary availability
    zpool_path = shutil.which('zpool')
    zfs_path = shutil.which('zfs')
    if not zpool_path or not zfs_path:
        # Skip gracefully when binaries are missing
        return {"entities": []}

    # Check device node (required for libzfs operations)
    if not os.path.exists('/dev/zfs'):
        # Skip if device not present in container/host
        return {"entities": []}
    
    pools = get_zfs_pools()
    
    def fmt_size(bytes_value: int):
        tb = 1024**4
        gb = 1024**3
        if bytes_value >= tb:
            return round(bytes_value / tb, 2), 'TB'
        return round(bytes_value / gb, 2), 'GB'

    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-check',
            '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'])
        size_value, size_unit = fmt_size(size_bytes)
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_size',
            'name': f'ZFS Pool {pool["name"]} Size',
            'value': str(size_value),
            'state_class': 'measurement',
            'unit_of_measurement': size_unit,
            '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'])
        alloc_value, alloc_unit = fmt_size(alloc_bytes)
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_allocated',
            'name': f'ZFS Pool {pool["name"]} Allocated',
            'value': str(alloc_value),
            'state_class': 'measurement',
            'unit_of_measurement': alloc_unit,
            'device_class': 'data_size',
            'icon': 'mdi:database-minus',
            'attributes': {
                'friendly_name': f'ZFS Pool {pool["name"]} Allocated Space'
            }
        })
        
        # Pool free space
        free_bytes = convert_size_to_bytes(pool['free'])
        free_value, free_unit = fmt_size(free_bytes)
        metrics['entities'].append({
            'sensor_id': f'zfs_pool_{pool["name"]}_free',
            'name': f'ZFS Pool {pool["name"]} Free',
            'value': str(free_value),
            'state_class': 'measurement',
            'unit_of_measurement': free_unit,
            'device_class': 'data_size',
            'icon': 'mdi:database-plus',
            '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:chart-donut',
            '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))