#!/usr/bin/env python3
from __future__ import annotations

import argparse
from collections import defaultdict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
MATERIALS = ROOT / 'materials'
OUTPUT = ROOT / 'docs/internal/material-inventory.md'


def human_size(size: int) -> str:
    units = ['B', 'KB', 'MB', 'GB']
    value = float(size)
    for unit in units:
        if value < 1024.0 or unit == units[-1]:
            if unit == 'B':
                return f'{int(value)} {unit}'
            return f'{value:.1f} {unit}'
        value /= 1024.0
    return f'{size} B'


def group_for(path: Path) -> str:
    rel = path.relative_to(MATERIALS)
    parts = rel.parts
    if parts[0] == 'benchmark' and len(parts) > 1:
        return f'benchmark/{parts[1]}'
    if parts[0] == 'reference' and len(parts) > 1:
        return f'reference/{parts[1]}'
    return parts[0]


def collect() -> tuple[dict[str, list[Path]], int]:
    groups: dict[str, list[Path]] = defaultdict(list)
    total = 0
    for path in sorted(
        p for p in MATERIALS.rglob('*')
        if p.is_file() and not p.name.startswith('.')
    ):
        groups[group_for(path)].append(path)
        total += path.stat().st_size
    return groups, total


def render(groups: dict[str, list[Path]], total: int) -> str:
    lines = [
        '# Material inventory',
        '',
        f'- total files: {sum(len(paths) for paths in groups.values())}',
        f'- total size: {human_size(total)}',
        '',
    ]
    for group in sorted(groups):
        lines.append(f'## {group}')
        lines.append('')
        lines.append('| Path | Size |')
        lines.append('| --- | ---: |')
        for path in groups[group]:
            rel = path.relative_to(ROOT)
            lines.append(f'| `{rel}` | {human_size(path.stat().st_size)} |')
        lines.append('')
    return '\n'.join(lines)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument('--check', action='store_true')
    args = parser.parse_args()

    if not MATERIALS.exists():
        raise SystemExit('materials/ does not exist')

    groups, total = collect()
    if not groups:
        raise SystemExit('No material files found')

    content = render(groups, total)
    if args.check:
        print(f'OK: {sum(len(paths) for paths in groups.values())} files, {human_size(total)}')
        return

    OUTPUT.parent.mkdir(parents=True, exist_ok=True)
    OUTPUT.write_text(content + '\n', encoding='utf-8')
    print(f'Wrote {OUTPUT}')


if __name__ == '__main__':
    main()