import asyncio
import time

from dotenv import load_dotenv

load_dotenv()

from typing import Dict, List, Optional, Union

import pytest

import litellm
from litellm import Router
from litellm.router import CustomRoutingStrategyBase

from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE


def _create_router():
    return Router(
        model_list=[
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/very-special-endpoint",
                    "api_base": FAKE_OPENAI_API_BASE,
                    "api_key": "fake-key",
                },
                "model_info": {"id": "very-special-endpoint"},
            },
            {
                "model_name": "azure-model",
                "litellm_params": {
                    "model": "openai/fast-endpoint",
                    "api_base": FAKE_OPENAI_API_BASE,
                    "api_key": "fake-key",
                },
                "model_info": {"id": "fast-endpoint"},
            },
        ],
        set_verbose=True,
        debug_level="DEBUG",
    )


class CustomRoutingStrategy(CustomRoutingStrategyBase):
    def __init__(self, router_instance: Router):
        self._router = router_instance

    async def async_get_available_deployment(
        self,
        model: str,
        messages: Optional[List[Dict[str, str]]] = None,
        input: Optional[Union[str, List]] = None,
        specific_deployment: Optional[bool] = False,
        request_kwargs: Optional[Dict] = None,
    ):
        print("In CUSTOM async get available deployment")
        model_list = self._router.model_list
        print("router model list=", model_list)
        for model in model_list:
            if isinstance(model, dict):
                if model["litellm_params"]["model"] == "openai/very-special-endpoint":
                    return model
        pass

    def get_available_deployment(
        self,
        model: str,
        messages: Optional[List[Dict[str, str]]] = None,
        input: Optional[Union[str, List]] = None,
        specific_deployment: Optional[bool] = False,
        request_kwargs: Optional[Dict] = None,
    ):
        pass


def test_reset_custom_routing_strategy():
    """
    Setting a custom routing strategy installs instance-level overrides for
    get_available_deployment / async_get_available_deployment. Re-initializing the
    routing strategy must clear them so the class implementations are used again.
    """
    router = _create_router()
    router.set_custom_routing_strategy(CustomRoutingStrategy(router))

    assert "get_available_deployment" in router.__dict__
    assert "async_get_available_deployment" in router.__dict__

    router._reset_custom_routing_strategy()

    assert "get_available_deployment" not in router.__dict__
    assert "async_get_available_deployment" not in router.__dict__
    assert (
        router.async_get_available_deployment.__func__
        is Router.async_get_available_deployment
    )

    # idempotent: resetting again when nothing is overridden must not raise
    router._reset_custom_routing_strategy()


@pytest.mark.asyncio
async def test_custom_routing():
    litellm.set_verbose = True

    router = _create_router()
    router.set_custom_routing_strategy(CustomRoutingStrategy(router))

    # make 4 requests
    for _ in range(4):
        try:
            response = await router.acompletion(
                model="azure-model", messages=[{"role": "user", "content": "hello"}]
            )
            print(response)
        except Exception as e:
            print("got exception", e)

    await asyncio.sleep(1)
    print("done sending initial requests to collect latency")

    deployments = {}
    # make 10 requests
    for _ in range(10):
        response = await router.acompletion(
            model="azure-model", messages=[{"role": "user", "content": "hello"}]
        )
        print(response)
        _picked_model_id = response._hidden_params["model_id"]
        if _picked_model_id not in deployments:
            deployments[_picked_model_id] = 1
        else:
            deployments[_picked_model_id] += 1
    print("deployments", deployments)