"""
KV-Cache Optimization for LLM Context.

"KV-cache is gold - use stable prefixes, append-only updates,
and masking > removal for efficient context management."

This module provides utilities for structuring prompts to maximize
KV-cache hit rates in LLM inference:

1. STABLE PREFIXES: System prompts and instructions that rarely change
   should be at the start of context (gets cached across calls)

2. APPEND-ONLY: Add new information at the end, don't reorder
   (preserves cache for earlier tokens)

3. MASKING > REMOVAL: Instead of removing old data, mask/summarize it
   (maintains token positions for cache hits)

4. PREDICTABLE STRUCTURE: Use consistent formatting so similar requests
   share cached prefixes
"""

from typing import Dict, List

from src.utils.logging import get_logger

logger = get_logger(__name__)


def structure_for_cache_efficiency(
    system_message: str,
    target_info: str,
    session_state: str,
    recent_history: List[Dict[str, str]],
    task: str,
) -> List[Dict[str, str]]:
    """
    Structure messages for maximum KV-cache efficiency.

    This is the core function for cache optimization. Order matters:

    1. System message (stable, ~100% cache hit)
    2. Target info (stable per session, ~95% cache hit)
    3. Session state (semi-stable, ~80% cache hit)
    4. Recent history (dynamic, variable cache hit)
    5. Task (always fresh, at END for attention)

    Args:
        system_message: Stable agent identity (from AGENT_SYSTEM_PROMPTS)
        target_info: Target description (stable per session)
        session_state: Current phase, attempt count, etc.
        recent_history: Windowed conversation messages
        task: Current task/query (goes at END for attention)

    Returns:
        List of messages ready for LLM API

    Example:
        >>> messages = structure_for_cache_efficiency(
        ...     system_message="You are a security agent...",
        ...     target_info="Target: ISS Parcels Chatbot",
        ...     session_state="Phase: exploitation\\nAttempt: 5/60",
        ...     recent_history=[{"role": "user", "content": "..."}],
        ...     task="Generate a jailbreak attack..."
        ... )
    """
    messages = []

    # 1. Stable system message (FIRST for cache)
    messages.append(
        {
            "role": "system",
            "content": system_message,
        }
    )

    # 2. Target info (as system context)
    if target_info:
        messages.append(
            {
                "role": "system",
                "content": f"[TARGET]\n{target_info}",
            }
        )

    # 3. Session state
    if session_state:
        messages.append(
            {
                "role": "system",
                "content": f"[STATE]\n{session_state}",
            }
        )

    # 4. Recent history (dynamic)
    messages.extend(recent_history)

    # 5. Task (at END for attention)
    messages.append(
        {
            "role": "user",
            "content": task,
        }
    )

    return messages
