"""Short-lived in-memory RDP agent bundles for WinRM clients to download."""

from __future__ import annotations

import io
import secrets
import time
import zipfile
from pathlib import Path
from threading import Lock

TTL_SECONDS = 600
_UTF8_BOM = b"\xef\xbb\xbf"
_BOM_SUFFIXES = {".ps1", ".txt"}

_lock = Lock()
_entries: dict[str, tuple[bytes, float]] = {}


def _prune_expired(now: float) -> None:
    expired = [token for token, (_, expires_at) in _entries.items() if expires_at <= now]
    for token in expired:
        _entries.pop(token, None)


def _bundle_file_bytes(path: Path) -> bytes:
    data = path.read_bytes()
    if path.suffix.lower() in _BOM_SUFFIXES and not data.startswith(_UTF8_BOM):
        return _UTF8_BOM + data
    return data


def build_rdp_bundle_zip(repo_dir: Path, filenames: tuple[str, ...]) -> bytes:
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive:
        for name in filenames:
            path = repo_dir / name
            if path.is_file():
                info = zipfile.ZipInfo(name)
                info.compress_type = zipfile.ZIP_DEFLATED
                info.flag_bits |= 0x800
                archive.writestr(info, _bundle_file_bytes(path))
    return buffer.getvalue()


def register_rdp_bundle_zip(zip_bytes: bytes) -> str:
    token = secrets.token_urlsafe(32)
    expires_at = time.time() + TTL_SECONDS
    with _lock:
        _prune_expired(time.time())
        _entries[token] = (zip_bytes, expires_at)
    return token


def get_rdp_bundle_zip(token: str) -> bytes | None:
    text = (token or "").strip()
    if not text:
        return None
    now = time.time()
    with _lock:
        _prune_expired(now)
        entry = _entries.get(text)
        if entry is None:
            return None
        data, expires_at = entry
        if expires_at <= now:
            _entries.pop(text, None)
            return None
        return data