from __future__ import annotations

import os
import shutil
import threading
import time
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Iterable


DEVICE_STORAGE_CACHE_SECONDS = 300.0


def _normalized_roots(paths: Iterable[Path]) -> list[Path]:
    roots: list[Path] = []
    for candidate in sorted(
        {Path(path).resolve() for path in paths},
        key=lambda path: len(path.parts),
    ):
        if any(candidate == root or candidate.is_relative_to(root) for root in roots):
            continue
        roots.append(candidate)
    return roots


def directory_bytes(paths: Iterable[Path]) -> int:
    total = 0
    pending = _normalized_roots(paths)
    while pending:
        path = pending.pop()
        if not path.exists():
            continue
        if not path.is_dir():
            total += path.stat(follow_symlinks=False).st_size
            continue
        with os.scandir(path) as entries:
            for entry in entries:
                if entry.is_dir(follow_symlinks=False):
                    pending.append(Path(entry.path))
                else:
                    total += entry.stat(follow_symlinks=False).st_size
    return total


class DeviceStorageMonitor:
    def __init__(
        self,
        program_root: Path,
        data_root: Path,
        *,
        cache_seconds: float = DEVICE_STORAGE_CACHE_SECONDS,
        clock: Callable[[], float] = time.monotonic,
        usage_reader: Callable[[Path], shutil._ntuple_diskusage] = shutil.disk_usage,
        directory_reader: Callable[[Iterable[Path]], int] = directory_bytes,
    ) -> None:
        self.program_root = Path(program_root).resolve()
        self.data_root = Path(data_root).resolve()
        self.cache_seconds = float(cache_seconds)
        self._clock = clock
        self._usage_reader = usage_reader
        self._directory_reader = directory_reader
        self._lock = threading.RLock()
        self._snapshot: dict | None = None
        self._sampled_monotonic = float("-inf")

    def invalidate(self) -> None:
        with self._lock:
            self._snapshot = None
            self._sampled_monotonic = float("-inf")

    def get_status(self) -> dict:
        with self._lock:
            now = self._clock()
            if (
                self._snapshot is not None
                and now - self._sampled_monotonic < self.cache_seconds
            ):
                return dict(self._snapshot)

            usage = self._usage_reader(self.program_root)
            snapshot = {
                "sampled_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
                "device_total_bytes": usage.total,
                "device_free_bytes": usage.free,
                "software_bytes": self._directory_reader(
                    [self.program_root, self.data_root]
                ),
                "templates_bytes": self._directory_reader(
                    [self.data_root / "templates", self.data_root / "animations"]
                ),
            }
            self._snapshot = snapshot
            self._sampled_monotonic = now
            return dict(snapshot)