from typing import Dict, List, Optional, Tuple, Union

from lm_eval.api.registry import register_model
from lm_eval.models.openai_completions import LocalCompletionsAPI
from lm_eval.models.utils import handle_stop_sequences


@register_model("sglang-generate")
class SGLANGGENERATEAPI(LocalCompletionsAPI):
    def __init__(
        self,
        base_url=None,
        tokenizer_backend="huggingface",
        **kwargs,
    ):
        super().__init__(
            base_url=base_url, tokenizer_backend=tokenizer_backend, **kwargs
        )

    def _create_payload(
        self,
        messages: Union[List[List[int]], List[dict], List[str], str],
        generate=False,
        gen_kwargs: Optional[dict] = None,
        seed: int = 1234,
        eos=None,
        **kwargs,
    ) -> dict:
        is_string = (
            True
            if (isinstance(messages, str) or isinstance(messages[0], str))
            else False
        )
        if generate:
            gen_kwargs.pop("do_sample", False)
            if "max_tokens" in gen_kwargs:
                max_tokens = gen_kwargs.pop("max_tokens")
            else:
                max_tokens = gen_kwargs.pop("max_gen_toks", self._max_gen_toks)
            temperature = gen_kwargs.pop("temperature", 0)
            stop = handle_stop_sequences(gen_kwargs.pop("until", None), eos)
            request = {
                "sampling_params": {
                    "max_new_tokens": max_tokens,
                    "temperature": temperature,
                    "stop": stop,
                    **gen_kwargs,
                },
            }
            request.update({"text": messages}) if is_string else request.update(
                {"input_ids": messages}
            )
            return request
        else:
            assert not is_string, "Logprobs are only supported for tokenized inputs"
            request = {
                "input_ids": messages,
                "sampling_params": {"max_new_tokens": 1, "temperature": 0},
                "logprob_start_len": 0,
                "top_logprobs_num": 1,
                "return_logprob": True,
            }
            return request

    @staticmethod
    def parse_logprobs(
        outputs: Union[Dict, List[Dict]],
        tokens: List[List[int]] = None,
        ctxlens: List[int] = None,
        **kwargs,
    ) -> List[Tuple[float, bool]]:
        res = []
        if not isinstance(outputs, list):
            outputs = [outputs]
        for choice, ctxlen in zip(outputs, ctxlens):
            choice = choice["meta_info"]
            assert ctxlen > 0, "Context length must be greater than 0"
            logprobs = sum(x[0] for x in choice["input_token_logprobs"][ctxlen:])
            is_greedy = all(
                x[1] != y[0][1]
                for x, y in zip(
                    choice["input_token_logprobs"][ctxlen:],
                    choice["input_top_logprobs"][ctxlen:],
                )
            )
            res.append((logprobs, is_greedy))
        return res

    @staticmethod
    def parse_generations(outputs: Union[Dict, List[Dict]], **kwargs) -> List[str]:
        res = []
        if not isinstance(outputs, list):
            outputs = [outputs]
        for out in outputs:
            res.append(out["text"])
        return res

    @property
    def api_key(self):
        return ""