"""HPC file list generation and splitting for distributed downloads."""

from __future__ import annotations

from pathlib import Path

import structlog

from xc_dl.core.manifest import Manifest

logger = structlog.get_logger()


def generate_file_lists(
    data_dir: Path,
    num_nodes: int,
    only_pending: bool = True,
) -> list[Path]:
    """Generate file lists split across N nodes.

    Returns list of paths to the generated file list files.
    """
    manifest = Manifest(data_dir)
    entries = manifest.read_all()

    if only_pending:
        ids = [xc_id for xc_id, e in entries.items() if e.status != "verified"]
    else:
        ids = list(entries.keys())

    ids.sort()

    file_lists_dir = data_dir / "file-lists"
    file_lists_dir.mkdir(parents=True, exist_ok=True)

    # Write full list
    full_list = file_lists_dir / "full-list.txt"
    full_list.write_text("\n".join(ids) + "\n")

    # Split into chunks
    chunk_size = max(1, len(ids) // num_nodes)
    paths: list[Path] = []

    for i in range(num_nodes):
        start = i * chunk_size
        if i == num_nodes - 1:
            chunk = ids[start:]
        else:
            chunk = ids[start : start + chunk_size]

        node_file = file_lists_dir / f"node-{i}.txt"
        node_file.write_text("\n".join(chunk) + "\n")
        paths.append(node_file)

    logger.info(
        "file_lists_generated",
        total_ids=len(ids),
        num_nodes=num_nodes,
        directory=str(file_lists_dir),
    )

    return paths


def load_file_list(path: Path) -> set[str]:
    """Load a file list (one XC ID per line) into a set."""
    ids: set[str] = set()
    with open(path) as f:
        for line in f:
            line = line.strip()
            if line:
                ids.add(line)
    return ids
