#!/usr/bin/env python3
"""Structure prompts for maximum prefix-cache hits across providers.

Demonstrates the core rule: static content first, dynamic content last,
with provider-specific cache breakpoints (Anthropic) vs auto-caching (OpenAI).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


# ── Canonical prompt structure (static-first) ──────────────────────────────

@dataclass
class Prompt:
    system: str
    tools: list[dict] = field(default_factory=list)
    examples: list[dict[str, str]] = field(default_factory=list)
    user_input: str = ""

    def static_block(self) -> str:
        """Everything that rarely changes, ordered first for cache reuse."""
        parts = [self.system]
        if self.tools:
            parts.append("\n## Tools\n" + "\n".join(
                f"- {t['name']}: {t['description']}" for t in self.tools
            ))
        if self.examples:
            parts.append("\n## Examples")
            for ex in self.examples:
                parts.append(f"Q: {ex['q']}\nA: {ex['a']}")
        return "\n".join(parts)

    def dynamic_block(self) -> str:
        """Per-request content, always appended LAST."""
        return self.user_input


# ── Anthropic: explicit cache_control breakpoints ─────────────────────────

def anthropic_messages(prompt: Prompt) -> list[dict[str, Any]]:
    """Build an Anthropic messages payload with cache_control breakpoints."""
    messages: list[dict[str, Any]] = []

    # Static system prompt gets a cache breakpoint
    messages.append({
        "role": "system",
        "content": [{
            "type": "text",
            "text": prompt.static_block(),
            "cache_control": {"type": "ephemeral"},  # mark stable block as cacheable
        }],
    })

    # Dynamic user content: NO breakpoint (it changes every request)
    messages.append({
        "role": "user",
        "content": [{"type": "text", "text": prompt.dynamic_block()}],
    })
    return messages


# ── OpenAI: automatic caching (just structure correctly) ──────────────────

def openai_messages(prompt: Prompt) -> list[dict[str, Any]]:
    """OpenAI caches automatically; keep stable prefix identical across calls."""
    return [
        {"role": "system", "content": prompt.static_block()},
        {"role": "user", "content": prompt.dynamic_block()},
    ]


# ── Cache-hit simulation: detect prefix divergence ────────────────────────

def estimate_prefix_overlap(prompts: list[Prompt]) -> dict[str, int]:
    """Count how many tokens a stable prefix is shared (rough proxy)."""
    statics = [p.static_block() for p in prompts]
    if not statics:
        return {}
    base = statics[0]
    shared_chars = 0
    for s in statics[1:]:
        n = 0
        for a, b in zip(base, s):
            if a == b:
                n += 1
            else:
                break
        shared_chars += n
    return {
        "base_static_chars": len(base),
        "total_shared_across_requests": shared_chars,
        "note": "A timestamp/random id ANYWHERE in the prefix destroys overlap",
    }


# ── Demo ───────────────────────────────────────────────────────────────────

def main() -> None:
    shared_tools = [{"name": "get_weather", "description": "Fetch weather by city"}]
    base_system = (
        "You are a helpful assistant. Respond concisely. "
        "If unsure, say UNKNOWN. Do not invent data."
    )

    prompts = [
        Prompt(system=base_system, tools=shared_tools, user_input="Weather in Paris?"),
        Prompt(system=base_system, tools=shared_tools, user_input="Weather in Tokyo?"),
        # BAD: injecting a timestamp into the static block breaks the cache
        Prompt(
            system=base_system + "\nRequest time: 2025-08-15T10:00:00Z",
            tools=shared_tools,
            user_input="Weather in Berlin?",
        ),
    ]

    print("=== Anthropic-style (with breakpoint) ===")
    print(anthropic_messages(prompts[0])[0]["content"][0]["cache_control"])

    print("\n=== Prefix overlap report ===")
    report = estimate_prefix_overlap(prompts)
    for k, v in report.items():
        print(f"  {k}: {v}")

    print("\nRules:")
    print("  1. Static blocks (system, tools, examples) go FIRST")
    print("  2. Dynamic content (user query) goes LAST")
    print("  3. Anthropic: add cache_control breakpoints to stable blocks only")
    print("  4. OpenAI: automatic - just keep the prefix byte-identical")


if __name__ == "__main__":
    main()