"""Skeleton compressor service for ACMS context inheritance.

The ``SkeletonCompressorService`` takes a collection of context
fragments from a parent plan and produces a compressed representation
suitable for propagation to child plans.  The compression is governed
by ``skeleton_ratio``:

- **0.0** — no compression; all fragments pass through unchanged.
- **1.0** — maximum compression; only the single highest-relevance
  fragment is kept (with minimal content).

Fragments are sorted by ``relevance_score`` in **descending** order
(highest first).  A stable secondary sort on ``fragment_id`` ensures
deterministic output for equal-relevance fragments.

Based on ``docs/specification.md`` ACMS Skeleton section.
"""

from __future__ import annotations

from dataclasses import dataclass

from cleveragents.domain.models.core.context_fragment import ContextFragment
from cleveragents.domain.models.core.skeleton_metadata import SkeletonMetadata

# ---------------------------------------------------------------------------
# Public data structures
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class CompressionResult:
    """Output of a skeleton compression pass.

    Attributes:
        fragments: The compressed (filtered/truncated) fragments,
            ordered by relevance descending then fragment_id ascending.
        metadata: Auditable metadata for the compression pass.
    """

    fragments: tuple[ContextFragment, ...]
    metadata: SkeletonMetadata


# ---------------------------------------------------------------------------
# Service
# ---------------------------------------------------------------------------

# Default skeleton_ratio when the caller does not specify one.
DEFAULT_SKELETON_RATIO: float = 0.15


class SkeletonCompressorService:
    """Compress context fragments for subplan context inheritance.

    The service is stateless; all state lives in the arguments and
    the returned ``CompressionResult``.
    """

    def compress(
        self,
        fragments: list[ContextFragment],
        skeleton_ratio: float | None = None,
    ) -> CompressionResult:
        """Compress *fragments* according to *skeleton_ratio*.

        Args:
            fragments: Context fragments to compress.  Each must
                have ``token_count >= 0`` and ``relevance_score`` in
                ``[0.0, 1.0]``.
            skeleton_ratio: Compression ratio in ``[0.0, 1.0]``.
                ``None`` falls back to ``DEFAULT_SKELETON_RATIO``.

        Returns:
            A ``CompressionResult`` containing the filtered fragments
            and associated ``SkeletonMetadata``.

        Raises:
            ValueError: If *skeleton_ratio* is outside ``[0.0, 1.0]``
                or any fragment has invalid fields.
            TypeError: If *fragments* is not a list.
        """
        # -- argument validation ------------------------------------------
        if not isinstance(fragments, list):
            raise TypeError(f"fragments must be a list, got {type(fragments).__name__}")

        effective_ratio = (
            skeleton_ratio if skeleton_ratio is not None else DEFAULT_SKELETON_RATIO
        )

        if not isinstance(effective_ratio, (int, float)):
            raise TypeError(
                f"skeleton_ratio must be a float, got {type(effective_ratio).__name__}"
            )

        if effective_ratio < 0.0 or effective_ratio > 1.0:
            raise ValueError(
                f"skeleton_ratio must be in [0.0, 1.0], got {effective_ratio}"
            )

        self._validate_fragments(fragments)

        # -- compute totals -----------------------------------------------
        original_tokens = sum(f.token_count for f in fragments)

        # -- stable sort: relevance desc, fragment_id asc -----------------
        sorted_fragments = sorted(
            fragments,
            key=lambda f: (-f.relevance_score, f.fragment_id),
        )

        # -- select fragments within budget -------------------------------
        kept = self._select_fragments(sorted_fragments, effective_ratio)

        compressed_tokens = sum(f.token_count for f in kept)

        source_ids = tuple(
            f.metadata["source_decision_id"]
            for f in kept
            if "source_decision_id" in f.metadata
        )

        metadata = SkeletonMetadata(
            ratio=effective_ratio,
            original_tokens=original_tokens,
            compressed_tokens=compressed_tokens,
            source_decision_ids=source_ids,
        )

        return CompressionResult(
            fragments=tuple(kept),
            metadata=metadata,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _validate_fragments(fragments: list[ContextFragment]) -> None:
        """Validate every fragment in the list.

        Raises:
            TypeError: If an element is not a ``ContextFragment``.
            ValueError: If a fragment has invalid field values.
        """
        for idx, frag in enumerate(fragments):
            if not isinstance(frag, ContextFragment):
                raise TypeError(
                    f"fragments[{idx}] must be a ContextFragment, "
                    f"got {type(frag).__name__}"
                )
            if frag.token_count < 0:
                raise ValueError(
                    f"fragments[{idx}].token_count must be >= 0, got {frag.token_count}"
                )
            if frag.relevance_score < 0.0 or frag.relevance_score > 1.0:
                raise ValueError(
                    f"fragments[{idx}].relevance_score must be in [0.0, 1.0], "
                    f"got {frag.relevance_score}"
                )
            if not frag.fragment_id:
                raise ValueError(f"fragments[{idx}].fragment_id must be non-empty")

    @staticmethod
    def _select_fragments(
        sorted_fragments: list[ContextFragment],
        ratio: float,
    ) -> list[ContextFragment]:
        """Select which fragments to keep given the compression ratio.

        When *ratio* is 0.0 every fragment is kept.  When *ratio* is
        1.0 only the single highest-relevance fragment survives (or
        none if the input is empty).  For intermediate values the
        token budget is ``original_tokens * (1 - ratio)``; fragments
        are added in relevance order until the budget is exhausted.
        """
        if not sorted_fragments:
            return []

        if ratio == 0.0:
            return list(sorted_fragments)

        original_tokens = sum(f.token_count for f in sorted_fragments)

        # Budget: fraction of tokens to *keep*
        budget = int(original_tokens * (1.0 - ratio))

        # At maximum compression keep at most one fragment
        if ratio == 1.0:
            budget = 0

        kept: list[ContextFragment] = []
        used = 0
        for frag in sorted_fragments:
            if used + frag.token_count > budget and kept:
                # Already have at least one fragment and budget exceeded
                break
            kept.append(frag)
            used += frag.token_count
            if used >= budget and budget > 0:
                break

        # At ratio 1.0, keep exactly the top fragment
        if ratio == 1.0 and sorted_fragments:
            return [sorted_fragments[0]]

        return kept
