# -*- coding: utf-8 -*-
"""Location: ./plugins/cached_tool_result/cached_tool_result.py
Copyright 2025
SPDX-License-Identifier: Apache-2.0
Authors: Mihai Criveti

Cached Tool Result Plugin.
Stores idempotent tool results in an in-memory cache keyed by tool name and
selected argument fields. Reads are advisory (metadata) due to framework
constraints; writes occur in tool_post_invoke.
"""

# Future
from __future__ import annotations

# Standard
from dataclasses import dataclass
import hashlib
import time
from typing import Any, Dict, List, Optional

# Third-Party
import orjson
from pydantic import BaseModel, Field

# Third-Party
from cpex.framework import (
    Plugin,
    PluginConfig,
    PluginContext,
    ToolPostInvokePayload,
    ToolPostInvokeResult,
    ToolPreInvokePayload,
    ToolPreInvokeResult,
)


class CacheConfig(BaseModel):
    """Configuration for cached tool result plugin.

    Attributes:
        cacheable_tools: List of tool names that should be cached.
        ttl: Time-to-live in seconds for cached results.
        key_fields: Optional mapping of tool names to specific argument fields to use for cache keys.
    """

    cacheable_tools: List[str] = Field(default_factory=list)
    ttl: int = 300
    key_fields: Optional[Dict[str, List[str]]] = None  # {tool: [fields...]}


@dataclass
class _Entry:
    """Cache entry containing a value and expiration timestamp.

    Attributes:
        value: Cached tool result.
        expires_at: Unix timestamp when the cached value expires.
    """

    value: Any
    expires_at: float


_CACHE: Dict[str, _Entry] = {}


def _make_key(tool: str, args: dict | None, fields: Optional[List[str]]) -> str:
    """Generate a cache key hash from tool name and selected argument fields.

    Args:
        tool: Tool name.
        args: Tool arguments dictionary.
        fields: Optional list of specific argument fields to include in the key.

    Returns:
        SHA256 hex digest cache key.
    """
    base = {"tool": tool, "args": {}}
    if args:
        if fields:
            base["args"] = {k: args.get(k) for k in fields}
        else:
            base["args"] = args
    raw = orjson.dumps(base, default=str, option=orjson.OPT_SORT_KEYS)
    return hashlib.sha256(raw).hexdigest()


class CachedToolResultPlugin(Plugin):
    """Cache idempotent tool results (write-through)."""

    def __init__(self, config: PluginConfig) -> None:
        """Initialize the cached tool result plugin.

        Args:
            config: Plugin configuration.
        """
        super().__init__(config)
        self._cfg = CacheConfig(**(config.config or {}))

    async def tool_pre_invoke(self, payload: ToolPreInvokePayload, context: PluginContext) -> ToolPreInvokeResult:
        """Check cache before tool invocation and store cache key in context.

        Args:
            payload: Tool invocation payload.
            context: Plugin execution context.

        Returns:
            Result with cache hit/miss metadata.
        """
        tool = payload.name
        if tool not in self._cfg.cacheable_tools:
            return ToolPreInvokeResult(continue_processing=True)
        fields = (self._cfg.key_fields or {}).get(tool)
        key = _make_key(tool, payload.args or {}, fields)
        # Persist key for post-invoke
        context.set_state("cache_key", key)
        context.set_state("cache_tool", tool)
        ent = _CACHE.get(key)
        now = time.time()
        if ent and ent.expires_at > now:
            # Advisory metadata; actual short-circuiting is not supported here
            return ToolPreInvokeResult(metadata={"cache_hit": True, "key": key})
        return ToolPreInvokeResult(metadata={"cache_hit": False, "key": key})

    async def tool_post_invoke(self, payload: ToolPostInvokePayload, context: PluginContext) -> ToolPostInvokeResult:
        """Store tool result in cache after invocation.

        Args:
            payload: Tool invocation result payload.
            context: Plugin execution context.

        Returns:
            Result with cache storage metadata.
        """
        tool = payload.name
        # Persist only for configured tools
        if tool not in self._cfg.cacheable_tools:
            return ToolPostInvokeResult(continue_processing=True)
        # Read key from context
        key = context.get_state("cache_key") if context else None
        if not key:
            # Fallback to a coarse key when args are unknown
            key = _make_key(tool, None, None)
        ttl = max(1, int(self._cfg.ttl))
        _CACHE[key] = _Entry(value=payload.result, expires_at=time.time() + ttl)
        return ToolPostInvokeResult(metadata={"cache_stored": True, "key": key, "ttl": ttl})
