import unittest

import openai
import requests

from sglang.srt.utils import kill_process_tree
from sglang.test.test_utils import (
    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
    DEFAULT_URL_FOR_TEST,
    CustomTestCase,
    popen_launch_server,
)


class TestCacheReport(CustomTestCase):
    @classmethod
    def setUpClass(cls):
        cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
        cls.base_url = DEFAULT_URL_FOR_TEST
        cls.min_cached = 5
        cls.process = popen_launch_server(
            cls.model,
            cls.base_url,
            timeout=300,
            other_args=[
                "--chunked-prefill-size=40",
                "--enable-cache-report",
            ],
        )
        cls.client = openai.Client(api_key="EMPTY", base_url=f"{cls.base_url}/v1")
        cls.aclient = openai.AsyncClient(api_key="EMPTY", base_url=f"{cls.base_url}/v1")

        usage = cls.run_openai(cls, "1").usage
        # we can assume that our request is of size 1, plus the total template size
        # ideally we would like to know the begin size / end size of the template to be more precise
        total_template_size = usage.prompt_tokens - 1
        print(f"template size: {total_template_size}")
        usage2 = cls.run_openai(cls, "2").usage
        assert usage2.prompt_tokens_details.cached_tokens <= total_template_size
        cls.min_cached = max(
            usage2.prompt_tokens_details.cached_tokens,
            total_template_size - usage2.prompt_tokens_details.cached_tokens,
        )

    @classmethod
    def tearDownClass(cls):
        kill_process_tree(cls.process.pid)

    def run_decode(self, return_logprob=False, top_logprobs_num=0, n=1):
        response = requests.post(
            self.base_url + "/generate",
            # we use an uncommon start to minimise the chance that the cache is hit by chance
            json={
                "text": "_ The capital of France is",
                "sampling_params": {
                    "temperature": 0 if n == 1 else 0.5,
                    "max_new_tokens": 128,
                    "n": n,
                    "stop_token_ids": [119690],
                },
                "stream": False,
                "return_logprob": return_logprob,
                "top_logprobs_num": top_logprobs_num,
                "logprob_start_len": 0,
            },
        )
        return response

    def run_openai(self, message):
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                # {"role": "system", "content": "You are a helpful AI assistant"},
                {"role": "user", "content": message},
            ],
            temperature=0,
            max_tokens=100,
        )
        return response

    async def run_openai_async(self, message):
        response = await self.aclient.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "user", "content": message},
            ],
            temperature=0,
            max_tokens=100,
        )
        return response

    def cache_report_openai(self, message):
        response = self.run_openai(message)
        print(
            f"openai first request cached_tokens: {int(response.usage.prompt_tokens_details.cached_tokens)}"
        )
        first_cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
        # assert int(response.usage.cached_tokens) == 0
        assert first_cached_tokens <= self.min_cached
        response = self.run_openai(message)
        cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
        print(f"openai second request cached_tokens: {cached_tokens}")
        assert cached_tokens > 0
        assert cached_tokens == int(response.usage.prompt_tokens) - 1
        return first_cached_tokens

    async def cache_report_openai_async(self, message):
        response = await self.run_openai_async(message)
        cached_tokens = int(response.usage.prompt_tokens_details.cached_tokens)
        prompt_tokens = int(response.usage.prompt_tokens)
        return cached_tokens, prompt_tokens

    def test_generate(self):
        print("=" * 100)
        response = self.run_decode()
        # print(response.json())
        cached_tokens = int(response.json()["meta_info"]["cached_tokens"])
        print(f"sglang first request cached_tokens: {cached_tokens}")
        print(
            f"sglang first request prompt_tokens: {int(response.json()['meta_info']['prompt_tokens'])}"
        )
        # can't assure to be 0: depends on the initialisation request / if a template is used with the model
        assert cached_tokens < self.min_cached
        response = self.run_decode()
        cached_tokens = int(response.json()["meta_info"]["cached_tokens"])
        print(f"sglang second request cached_tokens: {cached_tokens}")
        print(
            f"sglang second request prompt_tokens: {int(response.json()['meta_info']['prompt_tokens'])}"
        )
        assert cached_tokens == int(response.json()["meta_info"]["prompt_tokens"]) - 1

    def test_cache_split_prefill_openai(self):
        print("=" * 100)
        self.cache_report_openai(
            "€ This is a very long and unique text that should not be already cached, the twist is"
            " that it should be longer than the chunked-prefill-size, so it should be split among"
            " several prefill requests. Still, it shouldn't be cached"
        )

    def test_cache_report_openai(self):
        print("=" * 100)
        # warm up the cache, for the template
        self.run_openai("Introduce the capital of France.")

        first_cached_tokens_1 = self.run_openai(
            "How many sparrow do you need to lift a coconut?"
        ).usage.prompt_tokens_details.cached_tokens

        usage_2 = self.run_openai("* sing something about cats").usage
        first_cached_tokens_2 = usage_2.prompt_tokens_details.cached_tokens
        # first request may not have 0 cached tokens, but if they only have the template in common they
        # should be the same once the cache is warmed up
        assert first_cached_tokens_1 == first_cached_tokens_2

        resp = self.run_openai("* sing something about cats and dogs")
        print(resp.usage)

        resp = self.run_openai("* sing something about cats, please")
        print(resp.usage)
        assert (
            resp.usage.prompt_tokens_details.cached_tokens
            >= usage_2.prompt_tokens - self.min_cached
        )

    # TODO: flaky test
    # def test_cache_report_openai_async(self):
    #     print("=" * 100)

    #     async def run_test():
    #         task0 = asyncio.create_task(
    #             self.cache_report_openai_async(
    #                 "first request, to start the inference and let the next two request be started in the same batch"
    #             )
    #         )
    #         await asyncio.sleep(1)  # to force the first request to be started first
    #         task1 = asyncio.create_task(
    #             self.cache_report_openai_async(
    #                 "> can the same batch parallel request use the cache?"
    #             )
    #         )
    #         task2 = asyncio.create_task(
    #             self.cache_report_openai_async(
    #                 "> can the same batch parallel request use the cache?"
    #             )
    #         )
    #         result0, result1, result2 = await asyncio.gather(task0, task1, task2)

    #         cached_tokens0, prompt_tokens0 = result0
    #         cached_tokens1, prompt_tokens1 = result1
    #         cached_tokens2, prompt_tokens2 = result2

    #         print(
    #             f"Async request 0 - Cached tokens: {cached_tokens0}, Prompt tokens: {prompt_tokens0}"
    #         )
    #         print(
    #             f"Async request 1 - Cached tokens: {cached_tokens1}, Prompt tokens: {prompt_tokens1}"
    #         )
    #         print(
    #             f"Async request 2 - Cached tokens: {cached_tokens2}, Prompt tokens: {prompt_tokens2}"
    #         )

    #         # Assert that no requests used the cache (because first is alone, and the next two are in the same batch)
    #         # If a new optimisation limiting starting request with same prefix at the same time was added
    #         # to maximise the cache hit, this would not be true
    #         assert cached_tokens1 == cached_tokens2 == cached_tokens0

    #     asyncio.run(run_test())

    def test_cache_salt_effectiveness(self):
        print("=" * 100)
        print("Testing cache_salt effectiveness")

        # Use a unique message to avoid interference with other tests
        test_message = "What is the capital of Japan?"

        # First request with cache_salt "salt1"
        response1 = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": test_message}],
            temperature=0,
            max_tokens=10,
            extra_body={"cache_salt": "salt1"},
        )
        cached_tokens_1_first = int(response1.usage.prompt_tokens_details.cached_tokens)
        prompt_tokens_1 = int(response1.usage.prompt_tokens)
        print(
            f"First request with salt1 - cached_tokens: {cached_tokens_1_first}, prompt_tokens: {prompt_tokens_1}"
        )

        # Second request with same cache_salt "salt1" - should get cache hit
        response2 = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": test_message}],
            temperature=0,
            max_tokens=10,
            extra_body={"cache_salt": "salt1"},
        )
        cached_tokens_1_second = int(
            response2.usage.prompt_tokens_details.cached_tokens
        )
        print(
            f"Second request with salt1 - cached_tokens: {cached_tokens_1_second}, prompt_tokens: {prompt_tokens_1}"
        )

        # Verify cache hit for same salt
        assert (
            cached_tokens_1_second > cached_tokens_1_first
        ), "Should have cache hit with same cache_salt"
        assert (
            cached_tokens_1_second == prompt_tokens_1 - 1
        ), "Should cache all prompt tokens except the last one"

        # Third request with different cache_salt "salt2" - should not get cache hit
        response3 = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": test_message}],
            temperature=0,
            max_tokens=10,
            extra_body={"cache_salt": "salt2"},
        )
        cached_tokens_2_first = int(response3.usage.prompt_tokens_details.cached_tokens)
        print(f"First request with salt2 - cached_tokens: {cached_tokens_2_first}")

        # Verify no cache hit for different salt (should be similar to first request with salt1)
        assert (
            cached_tokens_2_first <= cached_tokens_1_first + self.min_cached
        ), "Different cache_salt should not share cache"

        # Fourth request with same cache_salt "salt2" - should now get cache hit
        response4 = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": test_message}],
            temperature=0,
            max_tokens=10,
            extra_body={"cache_salt": "salt2"},
        )
        cached_tokens_2_second = int(
            response4.usage.prompt_tokens_details.cached_tokens
        )
        print(f"Second request with salt2 - cached_tokens: {cached_tokens_2_second}")

        # Verify cache hit for salt2
        assert (
            cached_tokens_2_second == cached_tokens_2_first
        ), "Should have cache hit with same cache_salt for salt2"


if __name__ == "__main__":
    unittest.main()
