"""Space quotas, including logical descendants, trash, and pending uploads.

A child can be constrained by both its own allocation and its ancestors'.
Call assert_space_capacity() inside transaction.atomic() before reserving/writing.
"""
from contextlib import contextmanager
from django.db import transaction
from django.db.models import Sum
from .models import DriveItem, DriveSpace, UploadSession

MAX_QUOTA_BYTES = 9_007_199_254_740_991  # exact integer in JS too
PRIVATE_SPACE_QUOTA_BYTES = 50 * 1024 ** 3  # Maximum privé, non réservé dans le quota attribué


class QuotaExceeded(Exception):
    pass


def parse_quota(value):
    """None = unlimited. Accept only an integer number of bytes (0 is valid)."""
    if value is None:
        return None
    if isinstance(value, bool) or isinstance(value, float):
        raise ValueError('Le quota doit être un entier en octets, ou null pour illimité.')
    if isinstance(value, str) and (not value.strip().isdigit()):
        raise ValueError('Le quota doit être un entier en octets, ou null pour illimité.')
    try:
        quota = int(value)
    except (TypeError, ValueError, OverflowError):
        raise ValueError('Quota invalide.') from None
    if quota < 0 or quota > MAX_QUOTA_BYTES:
        raise ValueError('Quota hors limites (0 à 9 Po environ).')
    return quota


def ancestry(space):
    result = []
    seen = set()
    current = space
    while current is not None:
        if current.pk in seen:
            raise ValueError('Hiérarchie des espaces invalide.')
        seen.add(current.pk)
        result.append(current)
        current = current.parent
    return result


def subtree_space_ids(space):
    """Resolve the logical quota subtree, not the filesystem path.

    Private accounts store files below storage/privates/ but are children of
    their assigned spaces *for quota accounting*. Physical path-prefix
    queries would silently omit their real disk usage.
    """
    by_parent = {}
    for pk, parent_id in DriveSpace.objects.values_list('pk', 'parent_id'):
        by_parent.setdefault(parent_id, []).append(pk)
    result = set()
    to_visit = [space.pk]
    while to_visit:
        pk = to_visit.pop()
        if pk in result:
            raise ValueError('Hiérarchie des espaces invalide.')
        result.add(pk)
        to_visit.extend(by_parent.get(pk, []))
    return result


def occupancy(space):
    """Count ALL files, even in the trash. Pending sessions reserve their full size."""
    ids = subtree_space_ids(space)
    files = DriveItem.objects.filter(space_id__in=ids, is_folder=False).aggregate(n=Sum('size'))['n'] or 0
    pending = UploadSession.objects.filter(space_id__in=ids).aggregate(n=Sum('total_size'))['n'] or 0
    return int(files), int(pending)


def trash_occupancy(space):
    """Bytes still on disk in the trash, including descendants of trashed folders.

    Ordinary (not trashed) folders require no traversal. Only traverse the
    descendants of explicitly trashed folders; avoid counting trashed files twice.
    """
    ids = subtree_space_ids(space)
    bytes_in_trash = DriveItem.objects.filter(
        space_id__in=ids, is_folder=False, trashed_at__isnull=False
    ).aggregate(n=Sum('size'))['n'] or 0
    folder_ids = list(DriveItem.objects.filter(
        space_id__in=ids, is_folder=True, trashed_at__isnull=False
    ).values_list('pk', flat=True))
    visited = set(folder_ids)
    while folder_ids:
        next_folders = []
        for start in range(0, len(folder_ids), 300):
            for pk, is_folder, size, trashed in DriveItem.objects.filter(
                space_id__in=ids, parent_id__in=folder_ids[start:start + 300]
            ).values_list('pk', 'is_folder', 'size', 'trashed_at'):
                if is_folder:
                    if pk not in visited:
                        visited.add(pk)
                        next_folders.append(pk)
                elif trashed is None:
                    bytes_in_trash += size
        folder_ids = next_folders
    return int(bytes_in_trash)


def _budget_tree(space):
    """One DB snapshot of all descendant spaces and their direct physical usage.

    Shared descendants retain the existing reservation rules for assigned
    quotas. A private space's 50-Go cap is NEVER preallocated: its parent counts
    only its real files and pending uploads (including its own descendants).
    """
    ids=list(subtree_space_ids(space))
    nodes={node.pk:node for node in DriveSpace.objects.filter(pk__in=ids)}
    children={pk:[] for pk in nodes}
    for node in nodes.values():
        if node.parent_id in children:children[node.parent_id].append(node.pk)
    direct={pk:0 for pk in nodes}
    for pk,amount in DriveItem.objects.filter(space_id__in=ids,is_folder=False).values('space_id').annotate(total=Sum('size')).values_list('space_id','total'):
        direct[pk]+=int(amount or 0)
    for pk,amount in UploadSession.objects.filter(space_id__in=ids).values('space_id').annotate(total=Sum('total_size')).values_list('space_id','total'):
        direct[pk]+=int(amount or 0)
    memo={}
    def committed(pk):
        if pk not in memo:
            memo[pk] = direct[pk] + sum(
                committed(child) if nodes[child].private_owner_id
                else max(nodes[child].quota_bytes or 0, committed(child))
                for child in children[pk]
            )
        return memo[pk]
    committed(space.pk)
    return memo


def _committed(space):
    """Allocated bytes inside an entire subtree, including nested quota reservations."""
    return _budget_tree(space)[space.pk]


def available_for_growth(space):
    """Bytes that can be added directly to space without breaking any ancestor allocation.

    Quotas of normal assigned children can absorb growth already reserved in
    their ancestors. For private children, the 50-Go cap is NOT reserved: every
    newly uploaded byte immediately reduces the parent's available budget.
    """
    maximum=MAX_QUOTA_BYTES
    absorbed=0
    chain=ancestry(space)
    committed_tree=_budget_tree(chain[-1])
    for ancestor in chain:
        committed=committed_tree[ancestor.pk]
        if ancestor.quota_bytes is not None:
            maximum=min(maximum,ancestor.quota_bytes-committed+absorbed)
            # Private caps constrain uploads, but never absorb parent capacity.
            if not ancestor.private_owner_id:
                absorbed+=max(0,ancestor.quota_bytes-committed)
    return None if maximum==MAX_QUOTA_BYTES else max(0,int(maximum))


def assert_new_child_quota(parent, quota):
    """A new nested space cannot reserve more than its parent's unallocated capacity."""
    if quota is None:return
    remaining=available_for_growth(parent)
    if remaining is not None and quota>remaining:
        raise ValueError(f'Le quota dépasse les {remaining} octets disponibles dans l’espace parent.')


def assert_updated_quota(space, quota):
    """Validate nested allocations, including sibling limits and actual disk use.

    The old reservation is replaced atomically with the requested value. Parent
    quotas must be large enough for their direct content and all child quotas.
    """
    chain=ancestry(space)
    committed_tree=_budget_tree(chain[-1])
    committed=committed_tree[space.pk]
    if quota is not None and quota<committed:
        raise ValueError('Le quota est inférieur aux fichiers et sous-espaces déjà attribués.')
    current=space.quota_bytes
    if current==quota:return
    # Compute the changed reservation at every ancestor: each parent's
    # committed usage has this subtree's reserved/actual allocation substituted.
    old_child=max(current or 0,committed)
    new_child=max(quota or 0,committed)
    for ancestor in chain[1:]:
        amount=committed_tree[ancestor.pk]-old_child+new_child
        if ancestor.quota_bytes is not None and amount>ancestor.quota_bytes:
            raise ValueError(f'Quota impossible : capacité de « {ancestor.display_name} » dépassée.')
        old_child=max(ancestor.quota_bytes or 0,committed_tree[ancestor.pk])
        new_child=max(ancestor.quota_bytes or 0,amount)


def quota_state(space):
    used,reserved=occupancy(space)
    trashed=trash_occupancy(space)
    available=available_for_growth(space)
    return {
        'quota_bytes':space.quota_bytes,
        'used_bytes':used,
        'active_bytes':max(0,used-trashed),
        'trashed_bytes':trashed,
        'reserved_bytes':reserved,
        'available_bytes':available,
        'effective_limit_bytes':None if available is None else used+reserved+available,
    }


@contextmanager
def locked_quota_tree(space):
    """Lock ancestors in deterministic order so uploads and limit changes serialize."""
    if not transaction.get_connection().in_atomic_block:
        raise RuntimeError('locked_quota_tree must run inside transaction.atomic()')
    ids = sorted(a.pk for a in ancestry(space))
    list(DriveSpace.objects.select_for_update().filter(pk__in=ids).order_by('pk'))
    yield


def assert_space_capacity(space, additional_bytes):
    """Called with ancestor locks; private uploads count toward both caps."""
    available=available_for_growth(space)
    if available is not None and additional_bytes>available:
        raise QuotaExceeded(
            f'Espace insuffisant : {available} octets encore disponibles après les quotas des sous-espaces.'
        )