#!/usr/bin/env python3
"""Read-only Proxmox storage/reference inventory; no third-party dependencies."""
import argparse
import csv
from datetime import datetime
import json
import math
import os
import shutil
import textwrap
from pathlib import Path
import re
import subprocess
import sys

VOLUME = re.compile(r"(?<![\w.-])([\w.-]+:[^\s,;\"']+)")
FIELDS = ['node', 'storage', 'volume', 'size_bytes', 'status', 'detail', 'size_is_approximate']


def api(path, *options):
    result = subprocess.run(['pvesh', 'get', path, *options, '--output-format', 'json'],
                            capture_output=True, text=True, timeout=90)
    if result.returncode:
        raise RuntimeError(result.stderr.strip() or 'pvesh failed')
    data = json.loads(result.stdout)
    if not isinstance(data, list):
        raise ValueError('expected an API array')
    return data


def enabled(value):
    return value in (1, True, '1')


def references(root, nodes, guests):
    """Read pmxcfs on this host, including offline guests and all config sections."""
    refs, errors, seen, contents = {}, [], set(), {}
    guest_names = {(guest['type'], str(guest['vmid'])): guest.get('name')
                   for guest in guests}
    for node in nodes:
        for kind, directory in [('qemu', 'qemu-server'), ('lxc', 'lxc')]:
            folder = root / node['node'] / directory
            try:
                paths = list(folder.iterdir())
            except OSError as exc:
                errors.append(f'Cannot read {folder}: {exc}')
                continue
            for path in paths:
                if path.suffix != '.conf':
                    continue
                try:
                    content = path.read_text()
                except OSError as exc:
                    errors.append(f'Cannot read {path}: {exc}')
                    continue
                contents[str(path)] = content
                seen.add((node['node'], kind, path.stem))
                name = guest_names.get((kind, path.stem))
                name_key = 'name' if kind == 'qemu' else 'hostname'
                for line in content.splitlines():
                    line = line.strip()
                    if line.startswith('['):
                        break
                    key, sep, value = line.partition(':')
                    if sep and key == name_key and value.strip():
                        name = value.strip()
                        break
                guest_id = f"{'VM' if kind == 'qemu' else 'LXC'} {path.stem}"
                guest_label = f'{name} ({guest_id})' if name else f'{guest_id} (name unavailable)'
                section = 'current'
                for line in content.splitlines():
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue
                    if line.startswith('['):
                        section = line
                        continue
                    key, sep, value = line.partition(':')
                    if not sep or key in ('description', 'name'):
                        continue
                    for volume in VOLUME.findall(value):
                        refs.setdefault(volume, []).append({
                            'node': node['node'], 'guest': guest_label,
                            'unused': bool(re.fullmatch(r'unused\d+', key)) and section == 'current',
                            'section': section,
                        })
    for guest in guests:
        identity = (guest['node'], guest['type'], str(guest['vmid']))
        if identity not in seen:
            errors.append(f'Missing guest configuration: {identity}')
    return refs, errors, contents


def classify(volume, node, shared, refs, incomplete):
    all_refs = refs.get(volume['volid'], [])
    local_refs = [r for r in all_refs if shared or r['node'] == node]
    if any(not r['unused'] for r in local_refs):
        return None
    if incomplete:
        return 'Needs review', 'Guest reference inventory incomplete'
    if local_refs:
        return 'Unused disk', ', '.join(sorted({r['guest'] for r in local_refs}))
    if all_refs:
        return 'Needs review', 'Referenced on another node; possible replication copy: ' + ', '.join(
            sorted({r['node'] + '/' + r['guest'] for r in all_refs}))
    if re.search(r'(^|[/ :])base[-_]', volume['volid']) or volume.get('parent'):
        return 'Needs review', 'Possible template/clone backing relationship'
    return 'Orphan candidate', 'No guest configuration reference; verify replication and backing dependencies'


def scan(root=Path('/etc/pve/nodes'), get=api):
    warnings, rows, coverage = [], [], []
    nodes = get('/nodes')
    if not nodes:
        raise RuntimeError('No nodes returned; cannot inventory cluster')
    guests = get('/cluster/resources', '--type', 'vm')
    refs, errors, before = references(root, nodes, guests)
    warnings.extend(errors)
    shared_done = set()
    for node in sorted(nodes, key=lambda item: item['node']):
        name = node['node']
        if node.get('status') != 'online':
            warnings.append(f'{name}: offline; storage not scanned')
            continue
        ReportDisplay(sys.stderr).line(f'Scanning {name} storage...', 'muted')
        try:
            storages = get(f'/nodes/{name}/storage')
        except (RuntimeError, ValueError, OSError, subprocess.TimeoutExpired) as exc:
            warnings.append(f'{name}: storage discovery failed: {exc}')
            continue
        for storage in sorted(storages, key=lambda item: item['storage']):
            sid = storage['storage']
            if not {'images', 'rootdir'} & set(storage.get('content', '').split(',')):
                continue
            shared = enabled(storage.get('shared', 0))
            if shared and sid in shared_done:
                continue
            if not enabled(storage.get('active', 0)) or not enabled(storage.get('enabled', 1)):
                warnings.append(f'{name}/{sid}: inactive or disabled; not scanned')
                continue
            try:
                volumes = get(f'/nodes/{name}/storage/{sid}/content')
            except (RuntimeError, ValueError, OSError, subprocess.TimeoutExpired) as exc:
                warnings.append(f'{name}/{sid}: content listing failed: {exc}')
                continue
            if shared:
                shared_done.add(sid)
            label = 'shared (via ' + name + ')' if shared else name
            disks = [v for v in volumes if v.get('content') in ('images', 'rootdir')]
            coverage.append(f'{label}/{sid}: {len(disks)} disk volume(s) scanned')
            for volume in disks:
                result = classify(volume, name, shared, refs, bool(errors))
                if result:
                    size = volume.get('size')
                    approximate = False
                    if format_size(size) == 'Unknown':
                        size = volume.get('approximate-size')
                        approximate = format_size(size) != 'Unknown'
                    rows.append(dict(zip(FIELDS, [label, sid, volume['volid'], size, *result, approximate])))
    # Detect migrations/config changes during this non-atomic inventory.
    _, after_errors, after = references(root, nodes, guests)
    if after != before or after_errors:
        warnings.append('Guest configurations changed or became unreadable during scan; rerun report')
        for row in rows:
            row.update(status='Needs review', detail='Guest configurations changed during scan')
    return rows, warnings, coverage


def format_size(size_bytes, approximate=False):
    try:
        size = float(size_bytes)
    except (TypeError, ValueError):
        return 'Unknown'
    if not math.isfinite(size) or size < 0:
        return 'Unknown'
    for unit in ('B', 'KiB', 'MiB', 'GiB', 'TiB', 'PiB'):
        if size < 1024 or unit == 'PiB':
            return f"{'~' if approximate else ''}{size:.2f} {unit}"
        size /= 1024


class ReportDisplay:
    """Keep redirected reports plain and wrap text before applying ANSI styling."""
    def __init__(self, stream=None):
        self.stream = stream if stream is not None else sys.stdout
        self.color = (self.stream.isatty() and os.environ.get('TERM', 'dumb') != 'dumb'
                      and 'NO_COLOR' not in os.environ)
        self.width = max(20, min(75, shutil.get_terminal_size((80, 24)).columns))

    def line(self, text='', style='', indent=0):
        codes = {'bold': '1', 'muted': '90', 'accent': '96',
                 'success': '32', 'warning': '33', 'error': '31'}
        prefix = '\033[' + codes[style] + 'm' if self.color and style else ''
        reset = '\033[0m' if prefix else ''
        for line in textwrap.wrap(str(text), width=self.width - indent,
                                  replace_whitespace=False) or ['']:
            print(' ' * indent + prefix + line + reset, file=self.stream, flush=True)


def render(rows, warnings, coverage):
    display = ReportDisplay()
    display.line()
    display.line('Scan finished — incomplete coverage' if warnings else 'Scan finished',
                 'warning' if warnings else 'success')
    display.line(f'{len(coverage)} storage location(s) scanned.', 'muted')
    for warning in warnings:
        if warning.endswith(': inactive or disabled; not scanned'):
            continue
        display.line('WARNING: ' + warning, 'warning')
    display.line()
    display.line('Results', 'bold')
    counts = {status: sum(row['status'] == status for row in rows)
              for status in ('Unused disk', 'Orphan candidate', 'Needs review')}
    display.line(f"{counts['Unused disk']} unused disk(s)    "
                 f"{counts['Orphan candidate']} orphan candidate(s)    "
                 f"{counts['Needs review']} needing review", indent=2)
    group = None
    for row in sorted(rows, key=lambda r: (r['node'], r['storage'], r['volume'])):
        if group != (row['node'], row['storage']):
            group = (row['node'], row['storage'])
            display.line()
            display.line(row['storage'], 'bold')
            node = row['node']
            context = ('Shared storage · scanned via ' + node[len('shared (via '):-1]
                       if node.startswith('shared (via ') else 'Node: ' + node)
            display.line(context, 'muted')
        display.line()
        if row['status'] == 'Unused disk':
            display.line(row['detail'] + '  [Unused disk]', 'bold', indent=2)
        else:
            display.line('[' + row['status'] + ']', 'warning', indent=2)
            display.line(row['detail'], 'muted', indent=2)
        display.line('Volume:    ' + row['volume'], indent=2)
        size = format_size(row['size_bytes'], row.get('size_is_approximate', False))
        display.line('Disk size: ' + size, indent=2)
    display.line()
    display.line(f'{len(rows)} volume(s) reported.', 'bold')
    if not rows:
        display.line('No unreferenced or unused volumes found in the successfully scanned storage.')
    if any(row.get('size_is_approximate') for row in rows):
        display.line('~ indicates an approximate size.', 'muted')
    display.line('Review candidates before deletion. No disks were changed.', 'muted')
    display.line('Storage-managed volumes only; excludes backups, ISOs and passthrough disks.', 'muted')


def export_csv(path, rows):
    destination = Path(path).expanduser()
    folder = destination.is_dir() or str(path).endswith('/')
    if folder:
        destination.mkdir(parents=True, exist_ok=True)
        destination /= datetime.now().strftime('tapm-orphan-disks-%Y-%m-%d-%H%M%S.csv')
    elif destination.suffix.lower() != '.csv':
        destination = Path(str(destination) + '.csv')
    # Exclusive creation avoids overwriting existing files, including symlink targets.
    base = destination
    suffix = 0
    while True:
        try:
            handle = destination.open('x', newline='')
            break
        except FileExistsError:
            if not folder:
                raise
            suffix += 1
            destination = base.with_name(f'{base.stem}-{suffix}.csv')
    with handle:
        export_fields = [field for field in FIELDS if field != 'size_is_approximate']
        writer = csv.DictWriter(handle, fieldnames=export_fields + ['disk_size'])
        writer.writeheader()
        for row in rows:
            exported = {field: row.get(field) for field in export_fields}
            exported['disk_size'] = format_size(row['size_bytes'], row.get('size_is_approximate', False))
            writer.writerow(exported)
    return destination


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--csv', metavar='PATH', help='export to a CSV file or folder (.csv is added if missing)')
    parser.add_argument('--interactive', action='store_true', help='offer CSV export after the report')
    # The TAPM preloader pads its forwarded arguments with empty strings.
    args = parser.parse_args([arg for arg in sys.argv[1:] if arg])
    try:
        rows, warnings, coverage = scan()
        render(rows, warnings, coverage)
        destination = args.csv
        if args.interactive and sys.stdin.isatty() and not destination:
            display = ReportDisplay()
            display.line()
            display.line('CSV export file or folder:', 'accent')
            display.line('Enter a path, or press Enter to skip.', 'muted')
            destination = input('> ').strip()
        if destination:
            destination = export_csv(destination, rows)
            display = ReportDisplay()
            display.line('Report saved', 'success')
            display.line(str(destination))
        return 2 if warnings else 0
    except (OSError, RuntimeError, ValueError, subprocess.TimeoutExpired) as exc:
        print(f'Disk report failed: {exc}', file=sys.stderr)
        return 1


if __name__ == '__main__':
    sys.exit(main())