"""Clustering strategies for large-project decomposition.

Provides heuristics that partition a flat list of file paths into
bounded clusters suitable for independent subplan execution.  Three
complementary strategies are offered:

- **directory** — group by common directory prefix.
- **language** — group by file extension / programming language.
- **size** — group by cumulative estimated token count.

A deterministic sort ensures that clusters and their contents are
always produced in the same order regardless of filesystem iteration
order.

Based on Forgejo issue #205.
"""

from __future__ import annotations

import os
from collections import defaultdict

# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------

#: Rough bytes-per-token ratio (conservative).  Used when the caller
#: does not supply pre-computed token counts.
_BYTES_PER_TOKEN: int = 4


def estimate_tokens_for_path(file_path: str) -> int:
    """Return a rough token estimate for *file_path*.

    Falls back to ``0`` when the file cannot be read (e.g. binary or
    missing).  This function is intentionally cheap — it reads the
    file size from the OS, not the actual content.
    """
    try:
        size = os.path.getsize(file_path)
    except OSError:
        return 0
    return max(size // _BYTES_PER_TOKEN, 1)


# ---------------------------------------------------------------------------
# Clustering helpers
# ---------------------------------------------------------------------------


def _extension_of(path: str) -> str:
    """Return the lowercased file extension (e.g. ``".py"``)."""
    _, ext = os.path.splitext(path)
    return ext.lower()


def _directory_key(path: str, depth: int = 2) -> str:
    """Return the first *depth* path components as a grouping key.

    For ``"src/cleveragents/foo/bar.py"`` with *depth=2* the key is
    ``"src/cleveragents"``.
    """
    parts = path.replace("\\", "/").split("/")
    return "/".join(parts[:depth]) if len(parts) > depth else "/".join(parts[:-1])


# ---------------------------------------------------------------------------
# Public clustering API
# ---------------------------------------------------------------------------


class ClusteringStrategy:
    """Stateless helper that partitions files into bounded clusters.

    All methods return ``list[list[str]]`` — an ordered sequence of
    clusters where each inner list contains file paths.
    """

    @staticmethod
    def cluster_by_directory(
        files: list[str],
        max_per_cluster: int,
        *,
        depth: int = 2,
    ) -> list[list[str]]:
        """Group *files* by directory prefix.

        Args:
            files: File paths to partition.
            max_per_cluster: Maximum files per cluster.
            depth: Number of leading path components for the grouping key.

        Returns:
            Ordered list of clusters.
        """
        buckets: dict[str, list[str]] = defaultdict(list)
        for f in sorted(files):
            buckets[_directory_key(f, depth=depth)].append(f)

        clusters: list[list[str]] = []
        for key in sorted(buckets):
            bucket = buckets[key]
            # Split oversized buckets
            for i in range(0, len(bucket), max_per_cluster):
                clusters.append(bucket[i : i + max_per_cluster])
        return clusters

    @staticmethod
    def cluster_by_language(
        files: list[str],
        max_per_cluster: int,
    ) -> list[list[str]]:
        """Group *files* by programming language (file extension).

        Args:
            files: File paths to partition.
            max_per_cluster: Maximum files per cluster.

        Returns:
            Ordered list of clusters.
        """
        buckets: dict[str, list[str]] = defaultdict(list)
        for f in sorted(files):
            buckets[_extension_of(f)].append(f)

        clusters: list[list[str]] = []
        for ext in sorted(buckets):
            bucket = buckets[ext]
            for i in range(0, len(bucket), max_per_cluster):
                clusters.append(bucket[i : i + max_per_cluster])
        return clusters

    @staticmethod
    def cluster_by_size(
        files: list[str],
        max_tokens: int,
        *,
        token_map: dict[str, int] | None = None,
    ) -> list[list[str]]:
        """Group *files* so each cluster stays within *max_tokens*.

        Args:
            files: File paths to partition.
            max_tokens: Maximum estimated tokens per cluster.
            token_map: Optional pre-computed mapping of path → tokens.
                When ``None``, :func:`estimate_tokens_for_path` is used.

        Returns:
            Ordered list of clusters.
        """
        sorted_files = sorted(files)
        clusters: list[list[str]] = []
        current: list[str] = []
        current_tokens = 0

        for f in sorted_files:
            tokens = (
                token_map[f]
                if token_map and f in token_map
                else estimate_tokens_for_path(f)
            )
            if current and current_tokens + tokens > max_tokens:
                clusters.append(current)
                current = []
                current_tokens = 0
            current.append(f)
            current_tokens += tokens

        if current:
            clusters.append(current)
        return clusters

    @staticmethod
    def deterministic_sort(
        clusters: list[list[str]],
        language_priority: list[str] | None = None,
    ) -> list[list[str]]:
        """Sort *clusters* for deterministic execution ordering.

        Ordering criteria (most significant first):

        1. Average directory depth (shallower first).
        2. Dominant language position in *language_priority*.
        3. Total estimated file size (smaller first).
        4. First file path (lexicographic tiebreak).

        Args:
            clusters: Clusters to sort.
            language_priority: Ordered extensions for language ranking.

        Returns:
            New list with deterministic ordering.
        """
        priority = language_priority or []
        priority_map: dict[str, int] = {ext: idx for idx, ext in enumerate(priority)}
        max_priority = len(priority)

        def _sort_key(cluster: list[str]) -> tuple[float, int, int, str]:
            avg_depth = sum(f.count("/") for f in cluster) / max(len(cluster), 1)
            # Dominant extension = most frequent
            ext_counts: dict[str, int] = defaultdict(int)
            for f in cluster:
                ext_counts[_extension_of(f)] += 1
            dominant = max(ext_counts, key=ext_counts.get, default="")  # type: ignore[arg-type]
            lang_rank = priority_map.get(dominant, max_priority)
            total_size = len(cluster)
            first = cluster[0] if cluster else ""
            return (avg_depth, lang_rank, total_size, first)

        return sorted(clusters, key=_sort_key)


__all__: list[str] = [
    "ClusteringStrategy",
    "estimate_tokens_for_path",
]
