# Copyright (c) Microsoft. All rights reserved.

import asyncio
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import Annotated
from uuid import uuid4

from semantic_kernel import Kernel
from semantic_kernel.connectors.ai.open_ai import OpenAIChatCompletion, OpenAITextEmbedding
from semantic_kernel.connectors.in_memory import InMemoryStore
from semantic_kernel.data.vector import VectorStore, VectorStoreCollection, VectorStoreField, vectorstoremodel
from semantic_kernel.filters import FilterTypes, FunctionInvocationContext, PromptRenderContext
from semantic_kernel.functions import FunctionResult

COLLECTION_NAME = "llm_responses"
RECORD_ID_KEY = "cache_record_id"


# Define a simple data model to store, the prompt and the result
# we annotate the prompt field as the vector field, the prompt itself will not be stored.
# and if you use `include_vectors` in the search, it will return the vector, but not the prompt.
@vectorstoremodel(collection_name=COLLECTION_NAME)
@dataclass
class CacheRecord:
    result: Annotated[str, VectorStoreField("data", is_full_text_indexed=True)]
    prompt: Annotated[str | None, VectorStoreField("vector", dimensions=1536)] = None
    id: Annotated[str, VectorStoreField("key")] = field(default_factory=lambda: str(uuid4()))


# Define the filters, one for caching the results and one for using the cache.
class PromptCacheFilter:
    """A filter to cache the results of the prompt rendering and function invocation."""

    def __init__(
        self,
        vector_store: VectorStore,
        score_threshold: float = 0.2,
    ):
        if vector_store.embedding_generator is None:
            raise ValueError("The vector store must have an embedding generator.")
        self.vector_store = vector_store
        self.collection: VectorStoreCollection[str, CacheRecord] = vector_store.get_collection(record_type=CacheRecord)
        self.score_threshold = score_threshold

    async def on_prompt_render(
        self, context: PromptRenderContext, next: Callable[[PromptRenderContext], Awaitable[None]]
    ):
        """Filter to cache the rendered prompt and the result of the function.

        It uses the score threshold to determine if the result should be cached.
        The direction of the comparison is based on the default distance metric for
        the in memory vector store, which is cosine distance, so the closer to 0 the
        closer the match.
        """
        await next(context)
        await self.collection.ensure_collection_exists()
        results = await self.collection.search(context.rendered_prompt, vector_property_name="prompt", top=1)
        async for result in results.results:
            if result.score and result.score < self.score_threshold:
                context.function_result = FunctionResult(
                    function=context.function.metadata,
                    value=result.record.result,
                    rendered_prompt=context.rendered_prompt,
                    metadata={RECORD_ID_KEY: result.record.id},
                )

    async def on_function_invocation(
        self, context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
    ):
        """Filter to store the result in the cache if it is new."""
        await next(context)
        result = context.result
        if result and result.rendered_prompt and RECORD_ID_KEY not in result.metadata:
            cache_record = CacheRecord(prompt=result.rendered_prompt, result=str(result))
            await self.collection.ensure_collection_exists()
            await self.collection.upsert(cache_record)


async def execute_async(kernel: Kernel, title: str, prompt: str):
    """Helper method to execute and log time."""
    print(f"{title}: {prompt}")
    start = time.time()
    result = await kernel.invoke_prompt(prompt)
    elapsed = time.time() - start
    print(f"\tElapsed Time: {elapsed:.3f}")
    return result


async def main():
    # create the kernel and add the chat service and the embedding service
    kernel = Kernel()
    chat = OpenAIChatCompletion(service_id="default")
    embedding = OpenAITextEmbedding(service_id="embedder")
    kernel.add_service(chat)
    # create the in-memory vector store
    vector_store = InMemoryStore(embedding_generator=embedding)
    # create the cache filter and add the filters to the kernel
    cache = PromptCacheFilter(vector_store=vector_store)
    kernel.add_filter(FilterTypes.PROMPT_RENDERING, cache.on_prompt_render)
    kernel.add_filter(FilterTypes.FUNCTION_INVOCATION, cache.on_function_invocation)

    # Run the sample
    print("\nIn-memory cache sample:")
    r1 = await execute_async(kernel, "First run", "What's the tallest building in New York?")
    print(f"\tResult 1: {r1}")
    r2 = await execute_async(kernel, "Second run", "How are you today?")
    print(f"\tResult 2: {r2}")
    r3 = await execute_async(kernel, "Third run", "What is the highest building in New York City?")
    print(f"\tResult 3: {r3}")


if __name__ == "__main__":
    asyncio.run(main())