"""CachingMixin — response cache lookup and storage."""
from __future__ import annotations

from datetime import datetime, UTC

from logos_adk.types import Message, Response, content_text, normalize_content_parts


class CachingMixin:
    def _build_cache_key_data(self, provider_messages: list[Message]) -> dict:
        """Build data dict for cache key generation."""
        model_name = self._resolve_model_name()
        try:
            candidates = self._provider_candidates()
            decision = self._routing_decision_for_candidates(candidates)
            if decision is not None:
                _, _, model_name = decision.selected_model_key.partition(":")
            elif candidates:
                model_name = self._provider_model_name(candidates[0])
        except Exception:
            model_name = self._resolve_model_name()
        return {
            "model": model_name,
            "system_prompt": self._config.system_prompt,
            "messages": [
                {
                    "role": m.role,
                    "content": m.content,
                    "content_parts": [
                        {
                            "type": part.type,
                            "text": part.text,
                            "uri": part.uri,
                            "data": part.data,
                            "mime_type": part.mime_type,
                            "name": part.name,
                            "metadata": dict(part.metadata),
                        }
                        for part in normalize_content_parts(m.content_parts)
                    ],
                }
                for m in provider_messages
            ],
            "tools": self._build_tool_schemas(),
        }

    async def _cache_lookup(
        self,
        provider_messages: list[Message],
        trace_id: str,
        parent_span_id: str,
        session_id: str | None,
    ) -> Response | None:
        """Check cache for a matching response."""
        backend = self._config.cache_backend
        config = self._config.cache_config
        if backend is None or config is None:
            return None

        from logos_adk.cache.key import CacheKeyBuilder

        builder = CacheKeyBuilder()
        key_data = self._build_cache_key_data(provider_messages)

        started_at = datetime.now(UTC)
        entry = None

        if config.semantic and config.embed_fn:
            prompt = (
                content_text(
                    provider_messages[-1].content,
                    provider_messages[-1].content_parts,
                )
                if provider_messages
                else ""
            )
            vector = await config.embed_fn(prompt)
            scope = builder.scope_hash(
                model=key_data["model"],
                system_prompt=key_data["system_prompt"],
                tools=key_data["tools"],
            )
            results = await backend.search_similar(
                vector,
                config.similarity_threshold,
                limit=1,
                scope=scope,
            )
            if results:
                entry = results[0]

        if entry is None:
            key = builder.exact(**key_data)
            entry = await backend.get(key)

        hit = entry is not None
        await self._emit_span(
            name="Cache.lookup",
            trace_id=trace_id,
            parent_span_id=parent_span_id,
            started_at=started_at,
            session_id=session_id,
            attributes={
                "hit": hit,
                "backend": type(backend).__name__,
                "strategy": "semantic" if config.semantic else "exact",
            },
        )

        return entry.to_response() if entry else None

    async def _cache_store(
        self,
        provider_messages: list[Message],
        response: Response,
        trace_id: str,
        parent_span_id: str,
        session_id: str | None,
    ) -> None:
        """Store a response in the cache."""
        backend = self._config.cache_backend
        config = self._config.cache_config
        if backend is None or config is None:
            return
        if response.tool_calls:
            return

        from logos_adk.cache import CacheEntry
        from logos_adk.cache.key import CacheKeyBuilder

        builder = CacheKeyBuilder()
        key_data = self._build_cache_key_data(provider_messages)
        key = builder.exact(**key_data)

        started_at = datetime.now(UTC)
        entry = CacheEntry.from_response(response)

        if config.semantic and config.embed_fn:
            prompt = (
                content_text(
                    provider_messages[-1].content,
                    provider_messages[-1].content_parts,
                )
                if provider_messages
                else ""
            )
            entry.vector = await config.embed_fn(prompt)
            entry.scope = builder.scope_hash(
                model=key_data["model"],
                system_prompt=key_data["system_prompt"],
                tools=key_data["tools"],
            )

        await backend.put(key, entry, config.ttl)

        await self._emit_span(
            name="Cache.store",
            trace_id=trace_id,
            parent_span_id=parent_span_id,
            started_at=started_at,
            session_id=session_id,
            attributes={"key": key[:16]},
        )
