"""CSA/HCA compressor — Python API bridge.

Wraps the CuTeDSL compressor kernels with the interface that
AttentionSubBlock expects. The compressor itself is CuTeDSL because
it doesn't have the FMHA pipeline constraints — pure elementwise
softmax over m entries, no tensor cores needed.

The long-term path is raw CUDA C++ per doctrine, but the CuTeDSL
compressor is already working and correct. Rewrite only if MLIR
compilation becomes a blocker.
"""
from dsv4.kernels.compressor.csa_hca import (
    launch_csa_compress_projected,
    launch_hca_compress_projected,
)


def csa_compress_and_store(
    kv_raw: "torch.Tensor",      # (T, 4 * head_dim) BF16 — (Ca, Cb, Za, Zb) interleaved
    cache: "LayerCacheHandle",   # writes compressed entries to paged pool
    positions: "torch.Tensor",   # (T,) int64
    compression_ratio: int = 4,  # m=4
) -> None:
    """CSA: compress KV entries and store into the classical paged cache.

    The compressor reads the (Ca, Cb, Za, Zb) streams from kv_raw,
    runs token-level softmax compression (paper eq. 11-12), and writes
    the compressed entries to the cache's paged pool.

    The b-stream from the previous flush is read from the state cache's
    tail buffer (tail_kb, tail_zb).
    """
    # TODO: implement the full CSA compression + store path.
    # For now, this is a placeholder that writes raw KV to the tail buffer.
    # The full path needs:
    # 1. Read prev b-stream from state cache
    # 2. Run CuTeDSL compression kernel
    # 3. Write compressed output to paged pool via flush kernel
    # 4. Update tail buffer (a-stream becomes next b-stream)
    raise NotImplementedError(
        "CSA compress_and_store requires the full flush pipeline. "
        "See dsv4/kernels/cuda/flush_write.cu and dsv4/cache/flush.py"
    )


def hca_compress_and_store(
    kv_raw: "torch.Tensor",      # (T, 2 * head_dim) BF16 — (C, Z) interleaved
    cache: "LayerCacheHandle",   # writes compressed entries to paged pool
    positions: "torch.Tensor",   # (T,) int64
    compression_ratio: int = 128,  # m'=128
) -> None:
    """HCA: compress KV entries and store into the classical paged cache.

    Same structure as CSA but no b-stream, no overlap, and m'=128
    means compression only fires once per 128 tokens.
    """
    raise NotImplementedError(
        "HCA compress_and_store requires the full flush pipeline. "
        "See dsv4/kernels/cuda/flush_write.cu and dsv4/cache/flush.py"
    )