"""Prompt caching utilities for cost-optimized LLM inference via OpenRouter.

OpenRouter provides prompt caching across multiple providers:

- **OpenAI / DeepSeek / Groq / Grok (automatic)**: Cache hits require exact
  prefix matches on prompts >= 1024 tokens. No configuration needed.

- **Anthropic (explicit or automatic)**: Place ``cache_control``
  breakpoints on content blocks (up to 4) for the most portable setup.
  Some Anthropic routes also accept top-level automatic caching, but that is
  route-dependent on OpenRouter. Minimum cacheable tokens vary by model
  (1024–4096). Default 5-minute TTL, optional 1-hour TTL.

- **Google Gemini 2.5 (implicit)**: Automatic caching, no setup required.

OpenRouter uses **provider sticky routing** to maximize cache hits — after a
cached request, subsequent requests for the same model/conversation are routed
to the same provider endpoint automatically. Conversations are identified by
hashing the first system + first non-system message.

The :class:`PromptBuilder` enforces the "stable prefix, append-only growth"
pattern so agents get cache hits on every provider automatically.

References:
    - https://openrouter.ai/docs/guides/best-practices/prompt-caching
    - https://developers.openai.com/api/docs/guides/prompt-caching/
    - https://openai.com/index/unrolling-the-codex-agent-loop/
"""

from __future__ import annotations

import copy
import logging
from typing import Any

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Cache metrics
# ---------------------------------------------------------------------------

class CacheMetrics:
    """Tracks prompt cache hit/miss statistics across an agent session.

    Feed it the ``usage`` dict from each OpenRouter/OpenAI API response to
    accumulate lifetime cache performance numbers.

    Usage::

        metrics = CacheMetrics()

        # After each API response:
        metrics.record(response_json.get("usage", {}))

        print(metrics.summary())
        # Cache: 85.2% hit rate | 14,208 cached / 16,672 total prompt tokens
        # Estimated savings: $0.0128 (at 50% cached discount)
    """

    def __init__(self) -> None:
        self.total_prompt_tokens: int = 0
        self.cached_tokens: int = 0
        self.cache_creation_tokens: int = 0
        self.total_completion_tokens: int = 0
        self.request_count: int = 0
        self.cache_hit_requests: int = 0

    def record(self, usage: dict[str, Any]) -> None:
        """Record usage from a single API response.

        Handles OpenRouter's normalized format and provider-native formats:

        OpenRouter normalized (all providers)::

            {"prompt_tokens": 2006, "prompt_tokens_details": {
                "cached_tokens": 1920, "cache_write_tokens": 0}}

        Anthropic native (direct API)::

            {"prompt_tokens": 2006, "cache_read_input_tokens": 1920,
             "cache_creation_input_tokens": 86}

        Logs a warning if a response has prompt tokens but no cache data at all,
        which may indicate the prompt is below the model's minimum cacheable size.
        """
        self.request_count += 1

        prompt_tokens = usage.get("prompt_tokens", 0)
        self.total_prompt_tokens += prompt_tokens
        self.total_completion_tokens += usage.get("completion_tokens", 0)

        # Primary: OpenRouter normalized format
        # usage.prompt_tokens_details.cached_tokens (cache reads)
        # usage.prompt_tokens_details.cache_write_tokens (cache writes)
        details = usage.get("prompt_tokens_details", {})
        cached = details.get("cached_tokens", 0) if details else 0
        creation = details.get("cache_write_tokens", 0) if details else 0

        # Fallback: Anthropic native format (direct API, not via OpenRouter)
        cached += usage.get("cache_read_input_tokens", 0)
        creation += usage.get("cache_creation_input_tokens", 0)

        self.cache_creation_tokens += creation

        # Warn if prompt tokens are present but no cache data at all.
        # This may mean the prompt is below the model's minimum cacheable size,
        # or caching is not enabled for this provider/model.
        has_cache_fields = (
            cached > 0
            or creation > 0
            or (details and ("cached_tokens" in details or "cache_write_tokens" in details))
            or "cache_read_input_tokens" in usage
            or "cache_creation_input_tokens" in usage
        )
        if prompt_tokens > 0 and not has_cache_fields:
            if self.request_count <= 3:
                # Only warn on early requests to avoid log spam
                logger.warning(
                    "API response has %d prompt tokens but no cache data. "
                    "This may mean the prompt is below the model's minimum "
                    "cacheable token count, or caching is not supported for "
                    "this provider/model.",
                    prompt_tokens,
                )

        if cached > 0:
            self.cached_tokens += cached
            self.cache_hit_requests += 1

    @property
    def hit_rate(self) -> float:
        """Fraction of prompt tokens served from cache (0.0 – 1.0)."""
        if self.total_prompt_tokens == 0:
            return 0.0
        return self.cached_tokens / self.total_prompt_tokens

    @property
    def request_hit_rate(self) -> float:
        """Fraction of requests that had any cache hit."""
        if self.request_count == 0:
            return 0.0
        return self.cache_hit_requests / self.request_count

    def estimated_savings_usd(
        self,
        prompt_cost_per_m: float,
        cache_discount: float = 0.5,
    ) -> float:
        """Estimate USD saved by caching.

        Args:
            prompt_cost_per_m: Cost per million prompt tokens (non-cached).
            cache_discount: Fraction of prompt cost for cached tokens.
                OpenAI: ~50% discount. Anthropic: ~90% discount on reads.
        """
        saved_per_token = (prompt_cost_per_m / 1_000_000) * (1.0 - cache_discount)
        return self.cached_tokens * saved_per_token

    def summary(self) -> str:
        """Human-readable cache performance summary."""
        if self.request_count == 0:
            return "Cache: no requests recorded"
        pct = self.hit_rate * 100
        return (
            f"Cache: {pct:.1f}% hit rate | "
            f"{self.cached_tokens:,} cached / {self.total_prompt_tokens:,} total prompt tokens | "
            f"{self.cache_hit_requests}/{self.request_count} requests hit"
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "total_prompt_tokens": self.total_prompt_tokens,
            "cached_tokens": self.cached_tokens,
            "cache_creation_tokens": self.cache_creation_tokens,
            "total_completion_tokens": self.total_completion_tokens,
            "request_count": self.request_count,
            "cache_hit_requests": self.cache_hit_requests,
            "hit_rate": round(self.hit_rate, 4),
            "request_hit_rate": round(self.request_hit_rate, 4),
        }


# ---------------------------------------------------------------------------
# Prompt builder
# ---------------------------------------------------------------------------

class PromptBuilder:
    """Builds cache-optimized prompts for OpenRouter chat/completions.

    Enforces the "stable prefix, append-only growth" pattern from the Codex
    agent loop architecture. The prompt is assembled in this order:

    1. **System instructions** — static, cached aggressively
    2. **Tool definitions** — static, same order every request
    3. **Environment context** — semi-static (working dir, config)
    4. **Conversation history** — append-only, never mutated

    For Anthropic models, two caching modes are supported:

    - **Explicit** (default, recommended): Places ``cache_control`` on
      individual content blocks
      (system message + second-to-last user message). Max 4 breakpoints.
      Use for fine-grained control over what gets cached.
    - **Automatic**: Adds top-level ``cache_control`` to the request.
      This works only on Anthropic routes that advertise automatic caching.
      Plain string messages.

    For OpenAI/DeepSeek/Gemini, caching is automatic — no extra config needed.

    Usage::

        builder = PromptBuilder(
            system="You are a helpful coding assistant.",
            tools=[{"type": "function", "function": {"name": "bash", ...}}],
        )

        # First turn
        builder.add_user("Fix the bug in main.py")
        resp = httpx.post(OR_CHAT, json=builder.build(model="anthropic/claude-sonnet-4"))

        # Record cache metrics from OpenRouter's normalized usage format
        builder.metrics.record(resp.json().get("usage", {}))

        # Agent loop: append assistant + tool results, then next user turn
        builder.add_assistant(resp.json()["choices"][0]["message"]["content"])
        builder.add_user("Now add tests for it")
        resp = httpx.post(OR_CHAT, json=builder.build(model="anthropic/claude-sonnet-4"))
    """

    def __init__(
        self,
        system: str,
        tools: list[dict[str, Any]] | None = None,
        environment: str | None = None,
        provider: str = "auto",
        anthropic_mode: str = "explicit",
        cache_ttl: str | None = None,
    ) -> None:
        """
        Args:
            system: System prompt (static across the session).
            tools: Tool definitions list (static, order-preserved).
            environment: Optional environment context appended to system.
                Avoid timestamps or request IDs here — they break caching.
            provider: ``"anthropic"``, ``"openai"``, or ``"auto"`` (infers
                from model ID at build time).
            anthropic_mode: Caching mode for Anthropic models.
                ``"explicit"`` (default, recommended): per-block
                ``cache_control`` breakpoints on system message and
                second-to-last user message.
                ``"auto"``: top-level ``cache_control`` for Anthropic routes
                that support automatic caching.
            cache_ttl: Optional cache TTL. ``None`` (default) = 5 minutes,
                ``"1h"`` = 1 hour (Anthropic only, higher write cost).
        """
        self._system = system
        self._tools = tools or []
        self._environment = environment
        self._provider = provider
        self._anthropic_mode = anthropic_mode
        self._cache_ttl = cache_ttl
        self._messages: list[dict[str, Any]] = []
        self.metrics = CacheMetrics()

    # -- Message management (append-only) -----------------------------------

    def add_user(self, content: str) -> list[dict[str, Any]]:
        """Append a user message. Returns the current messages list."""
        self._messages.append({"role": "user", "content": content})
        return self._messages

    def add_assistant(self, content: str) -> list[dict[str, Any]]:
        """Append an assistant message. Returns the current messages list."""
        self._messages.append({"role": "assistant", "content": content})
        return self._messages

    def add_tool_result(self, tool_call_id: str, content: str) -> list[dict[str, Any]]:
        """Append a tool result message. Returns the current messages list."""
        self._messages.append({
            "role": "tool",
            "tool_call_id": tool_call_id,
            "content": content,
        })
        return self._messages

    def add_message(self, message: dict[str, Any]) -> list[dict[str, Any]]:
        """Append an arbitrary message dict. Returns the current messages list."""
        self._messages.append(message)
        return self._messages

    @property
    def turn_count(self) -> int:
        """Number of user messages so far."""
        return sum(1 for m in self._messages if m.get("role") == "user")

    # -- Build request payload ----------------------------------------------

    def build(
        self,
        model: str,
        max_tokens: int | None = None,
        **extra: Any,
    ) -> dict[str, Any]:
        """Build a complete chat/completions request payload.

        The payload is structured for maximum cache efficiency:

        - System message is always first (stable prefix)
        - Tools are in fixed order (part of cached prefix)
        - Messages are append-only
        - For Anthropic: ``cache_control`` breakpoints are set
        - For OpenAI: structure naturally supports automatic caching

        Args:
            model: OpenRouter model ID (e.g. ``"anthropic/claude-sonnet-4"``).
            max_tokens: Optional max completion tokens.
            **extra: Additional params forwarded to the API.

        Returns:
            Dict ready to POST to ``/chat/completions``.
        """
        provider = self._resolve_provider(model)
        messages = self._build_messages(provider)

        payload: dict[str, Any] = {
            "model": model,
            "messages": messages,
        }

        if self._tools:
            payload["tools"] = self._tools

        if max_tokens is not None:
            payload["max_tokens"] = max_tokens

        # Anthropic automatic caching: top-level cache_control
        # OpenRouter places the breakpoint on the last cacheable block
        # and advances it as the conversation grows.
        if provider == "anthropic" and self._anthropic_mode == "auto":
            cc: dict[str, str] = {"type": "ephemeral"}
            if self._cache_ttl:
                cc["ttl"] = self._cache_ttl
            payload["cache_control"] = cc

        payload.update(extra)

        return payload

    def _resolve_provider(self, model: str) -> str:
        """Determine provider from model ID or explicit setting."""
        if self._provider != "auto":
            return self._provider
        model_lower = model.lower()
        if "anthropic/" in model_lower or "claude" in model_lower:
            return "anthropic"
        return "openai"

    def _build_messages(self, provider: str) -> list[dict[str, Any]]:
        """Assemble messages with provider-appropriate caching hints."""
        messages: list[dict[str, Any]] = []

        # System message — always first, always cached
        system_content = self._system
        if self._environment:
            system_content += "\n\n" + self._environment

        use_explicit = (provider == "anthropic" and self._anthropic_mode == "explicit")

        if use_explicit:
            # Explicit mode: per-block cache_control breakpoints
            messages.append(self._anthropic_system(system_content, ttl=self._cache_ttl))
        else:
            # Auto mode (or non-Anthropic): plain string system message.
            # For Anthropic auto mode, top-level cache_control is added in build().
            messages.append({"role": "system", "content": system_content})

        # Conversation messages — append-only, deep-copied to avoid mutation
        conv = [copy.deepcopy(m) for m in self._messages]

        if use_explicit and len(conv) >= 2:
            # Explicit mode: place cache_control on the second-to-last user
            # message so earlier conversation is reused from cache.
            for i in range(len(conv) - 2, -1, -1):
                if conv[i].get("role") == "user":
                    conv[i] = self._anthropic_cache_message(conv[i], ttl=self._cache_ttl)
                    break

        messages.extend(conv)
        return messages

    # -- Anthropic-specific cache_control helpers ---------------------------

    @staticmethod
    def _cache_control(ttl: str | None = None) -> dict[str, str]:
        """Build a cache_control object with optional TTL."""
        cc: dict[str, str] = {"type": "ephemeral"}
        if ttl:
            cc["ttl"] = ttl
        return cc

    @staticmethod
    def _anthropic_system(content: str, ttl: str | None = None) -> dict[str, Any]:
        """Build a system message with Anthropic cache_control breakpoint."""
        cc = PromptBuilder._cache_control(ttl)
        return {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": content,
                    "cache_control": cc,
                }
            ],
        }

    @staticmethod
    def _anthropic_cache_message(msg: dict[str, Any], ttl: str | None = None) -> dict[str, Any]:
        """Add cache_control to a message's content for Anthropic."""
        cc = PromptBuilder._cache_control(ttl)
        msg = copy.deepcopy(msg)
        content = msg.get("content", "")
        if isinstance(content, str):
            msg["content"] = [
                {
                    "type": "text",
                    "text": content,
                    "cache_control": cc,
                }
            ]
        elif isinstance(content, list) and content:
            # Add cache_control to the last content block
            content[-1] = {**content[-1], "cache_control": cc}
            msg["content"] = content
        return msg

    # -- Convenience: request-level caching hints ---------------------------

    @staticmethod
    def openai_cache_params(retention: str = "in_memory") -> dict[str, Any]:
        """Return OpenAI-specific caching parameters.

        Args:
            retention: ``"in_memory"`` (default, 5-10 min TTL) or
                ``"24h"`` (extended, up to 24h, select models only).
        """
        params: dict[str, Any] = {}
        if retention == "24h":
            params["prompt_cache_retention"] = "24h"
        return params


# ---------------------------------------------------------------------------
# Helpers for agents running inside containers
# ---------------------------------------------------------------------------

# ---------------------------------------------------------------------------
# Anthropic minimum cacheable token counts (per OpenRouter docs)
# ---------------------------------------------------------------------------

ANTHROPIC_MIN_CACHE_TOKENS: dict[str, int] = {
    "anthropic/claude-opus-4.6": 4096,
    "anthropic/claude-opus-4-6": 4096,
    "anthropic/claude-opus-4.5": 4096,
    "anthropic/claude-opus-4-5": 4096,
    "anthropic/claude-haiku-4.5": 4096,
    "anthropic/claude-haiku-4-5": 4096,
    "anthropic/claude-sonnet-4.6": 2048,
    "anthropic/claude-sonnet-4-6": 2048,
    "anthropic/claude-haiku-3.5": 2048,
    "anthropic/claude-haiku-3-5": 2048,
    "anthropic/claude-sonnet-4.5": 1024,
    "anthropic/claude-sonnet-4-5": 1024,
    "anthropic/claude-opus-4.1": 1024,
    "anthropic/claude-opus-4-1": 1024,
    "anthropic/claude-opus-4": 1024,
    "anthropic/claude-sonnet-4": 1024,
    "anthropic/claude-sonnet-3.7": 1024,
    "anthropic/claude-sonnet-3-7": 1024,
}
"""Minimum prompt tokens required for caching per Anthropic model."""

# Other providers
OPENAI_MIN_CACHE_TOKENS = 1024
"""OpenAI requires >= 1024 tokens for automatic caching."""

GOOGLE_MIN_CACHE_TOKENS = 4096
"""Gemini models typically require >= 4096 tokens for caching."""


def min_cache_tokens(model_id: str) -> int:
    """Return the minimum cacheable token count for a model.

    Falls back to 1024 if the model is not in the known list.
    """
    if model_id in ANTHROPIC_MIN_CACHE_TOKENS:
        return ANTHROPIC_MIN_CACHE_TOKENS[model_id]
    provider = detect_provider(model_id)
    if provider == "openai":
        return OPENAI_MIN_CACHE_TOKENS
    if provider == "google":
        return GOOGLE_MIN_CACHE_TOKENS
    return 1024


def detect_provider(model_id: str) -> str:
    """Infer the caching strategy from an OpenRouter model ID.

    Returns ``"anthropic"``, ``"openai"``, ``"google"``, or ``"other"``.
    """
    model_lower = model_id.lower()
    if model_lower.startswith("anthropic/") or "claude" in model_lower:
        return "anthropic"
    if model_lower.startswith("openai/") or "gpt-" in model_lower:
        return "openai"
    if model_lower.startswith("google/") or "gemini" in model_lower:
        return "google"
    return "other"


def stable_system_prompt(*sections: str) -> str:
    """Join system prompt sections with double newlines.

    Use this to build a system prompt from stable, reusable parts::

        system = stable_system_prompt(
            "You are a coding assistant.",
            "Available tools: bash, read_file, write_file.",
            "Rules: always explain before editing.",
        )

    Each section should be static text — avoid timestamps, request IDs,
    or user-specific data that would invalidate the cache prefix.
    """
    return "\n\n".join(s.strip() for s in sections if s and s.strip())


def format_cache_report(metrics: CacheMetrics, model_id: str = "") -> str:
    """Format a cache performance report for logging or budget.txt injection.

    Example output::

        Prompt Cache Performance (anthropic/claude-sonnet-4):
          Hit rate:     85.2% (14,208 / 16,672 prompt tokens cached)
          Requests:     12 total, 10 with cache hits
          Cache writes: 2,464 tokens (creation cost)
    """
    lines = []
    header = "Prompt Cache Performance"
    if model_id:
        header += f" ({model_id})"
    lines.append(f"{header}:")

    if metrics.request_count == 0:
        lines.append("  (no requests recorded)")
        return "\n".join(lines)

    lines.append(
        f"  Hit rate:     {metrics.hit_rate * 100:.1f}% "
        f"({metrics.cached_tokens:,} / {metrics.total_prompt_tokens:,} prompt tokens cached)"
    )
    lines.append(
        f"  Requests:     {metrics.request_count} total, "
        f"{metrics.cache_hit_requests} with cache hits"
    )
    if metrics.cache_creation_tokens > 0:
        lines.append(
            f"  Cache writes: {metrics.cache_creation_tokens:,} tokens (creation cost)"
        )
    return "\n".join(lines)
