import json
import time
from typing import Any, AsyncGenerator, AsyncIterable

from mlflow.gateway.config import CohereConfig, RouteConfig
from mlflow.gateway.exceptions import AIGatewayException
from mlflow.gateway.providers.base import BaseProvider, ProviderAdapter
from mlflow.gateway.providers.utils import rename_payload_keys, send_request, send_stream_request
from mlflow.gateway.schemas import chat, completions, embeddings


class CohereAdapter(ProviderAdapter):
    @staticmethod
    def _scale_temperature(payload):
        # The range of Cohere's temperature is 0-5, but ours is 0-2, so we scale it.
        if temperature := payload.get("temperature"):
            payload["temperature"] = 2.5 * temperature
        return payload

    @classmethod
    def model_to_completions(cls, resp, config):
        # Response example (https://docs.cohere.com/reference/generate)
        # ```
        # {
        #   "id": "string",
        #   "generations": [
        #     {
        #       "id": "string",
        #       "text": "string"
        #     }
        #   ],
        #   "prompt": "string"
        # }
        # ```
        return completions.ResponsePayload(
            created=int(time.time()),
            object="text_completion",
            model=config.model.name,
            choices=[
                completions.Choice(
                    index=idx,
                    text=c["text"],
                    finish_reason=None,
                )
                for idx, c in enumerate(resp["generations"])
            ],
            usage=completions.CompletionsUsage(
                prompt_tokens=None,
                completion_tokens=None,
                total_tokens=None,
            ),
        )

    @classmethod
    def model_to_completions_streaming(cls, resp, config):
        # Response example (https://docs.cohere.com/reference/generate)
        #
        # Streaming chunks:
        # ```
        # {"index":0,"text":" Hi","is_finished":false,"event_type":"text-generation"}
        # ```
        # ```
        # {"index":1,"text":" Hi","is_finished":false,"event_type":"text-generation"}
        # ```
        # notes: "index" is only present if "num_generations" > 1
        #
        # Final chunk:
        # ```
        # {"is_finished":true,"event_type":"stream-end","finish_reason":"COMPLETE",
        #   "response":{"id":"b32a70c5-8c91-4f96-958f-d942801ed22f",
        #       "generations":[
        #           {
        #               "id":"5d5d0851-35ac-4c25-a9a9-2fbb391bd415",
        #               "index":0,
        #               "text":" Hi there! How can I assist you today? ",
        #               "finish_reason":"COMPLETE"
        #           },
        #           {
        #               "id":"0a24787f-504e-470e-a088-0bf801a2c72d",
        #               "index":1,
        #               "text":" Hi there, how can I assist you today? ",
        #               "finish_reason":"COMPLETE"
        #           }
        #       ],
        #       "prompt":"Hello"
        #   }}
        # ```
        response = resp.get("response")
        return completions.StreamResponsePayload(
            id=response["id"] if response else None,
            created=int(time.time()),
            model=config.model.name,
            choices=[
                completions.StreamChoice(
                    index=resp.get("index", 0),
                    finish_reason=resp.get("finish_reason"),
                    text=resp.get("text"),
                )
            ],
            usage=completions.CompletionsUsage(
                prompt_tokens=None,
                completion_tokens=None,
                total_tokens=None,
            ),
        )

    @classmethod
    def model_to_embeddings(cls, resp, config):
        # Response example (https://docs.cohere.com/reference/embed):
        # ```
        # {
        #   "id": "bc57846a-3e56-4327-8acc-588ca1a37b8a",
        #   "texts": [
        #     "hello world"
        #   ],
        #   "embeddings": [
        #     [
        #       3.25,
        #       0.7685547,
        #       2.65625,
        #       ...
        #       -0.30126953,
        #       -2.3554688,
        #       1.2597656
        #     ]
        #   ],
        #   "meta": [
        #     {
        #       "api_version": [
        #         {
        #           "version": "1"
        #         }
        #       ]
        #     }
        #   ]
        # }
        # ```
        return embeddings.ResponsePayload(
            data=[
                embeddings.EmbeddingObject(
                    embedding=output,
                    index=idx,
                )
                for idx, output in enumerate(resp["embeddings"])
            ],
            model=config.model.name,
            usage=embeddings.EmbeddingsUsage(
                prompt_tokens=None,
                total_tokens=None,
            ),
        )

    @classmethod
    def completions_to_model(cls, payload, config):
        key_mapping = {
            "stop": "stop_sequences",
            "n": "num_generations",
        }
        cls.check_keys_against_mapping(key_mapping, payload)
        payload = cls._scale_temperature(payload)
        return rename_payload_keys(payload, key_mapping)

    @classmethod
    def completions_streaming_to_model(cls, payload, config):
        return cls.completions_to_model(payload, config)

    @classmethod
    def embeddings_to_model(cls, payload, config):
        key_mapping = {"input": "texts"}
        for k1, k2 in key_mapping.items():
            if k2 in payload:
                raise AIGatewayException(
                    status_code=422, detail=f"Invalid parameter {k2}. Use {k1} instead."
                )
        return rename_payload_keys(payload, key_mapping)

    @classmethod
    def chat_to_model(cls, payload, config):
        if payload["n"] != 1:
            raise AIGatewayException(
                status_code=422,
                detail=f"Parameter n must be 1 for Cohere chat, got {payload['n']}.",
            )
        del payload["n"]

        if "stop" in payload:
            raise AIGatewayException(
                status_code=422,
                detail="Parameter stop is not supported for Cohere chat.",
            )
        payload = cls._scale_temperature(payload)

        messages = payload.pop("messages")
        last_message = messages.pop()  # pydantic enforces min_length=1
        if last_message["role"] != "user":
            raise AIGatewayException(
                status_code=422,
                detail=f"Last message must be from user, got {last_message['role']}.",
            )
        payload["message"] = last_message["content"]

        # Cohere uses `preamble_override` to set the system message
        # we concatenate all system messages from the user with a newline
        system_messages = [m for m in messages if m["role"] == "system"]
        if len(system_messages) > 0:
            payload["preamble_override"] = "\n".join(m["content"] for m in system_messages)

        # remaining messages are chat history
        # we want to include only user and assistant messages
        messages = [m for m in messages if m["role"] in ("user", "assistant")]
        if messages:
            payload["chat_history"] = [
                {
                    "role": "USER" if m["role"] == "user" else "CHATBOT",
                    "message": m["content"],
                }
                for m in messages
            ]
        return payload

    @classmethod
    def chat_streaming_to_model(cls, payload, config):
        return cls.chat_to_model(payload, config)

    @classmethod
    def model_to_chat(cls, resp, config):
        # Response example (https://docs.cohere.com/reference/chat)
        # ```
        # {
        #   "response_id": "string",
        #   "text": "string",
        #   "generation_id": "string",
        #   "token_count": {
        #     "prompt_tokens": 0,
        #     "response_tokens": 0,
        #     "total_tokens": 0,
        #     "billed_tokens": 0
        #   },
        #   "meta": {
        #     "api_version": {
        #       "version": "1"
        #     },
        #     "billed_units": {
        #       "input_tokens": 0,
        #       "output_tokens": 0
        #     }
        #   },
        #   "tool_inputs": null
        # }
        # ```
        return chat.ResponsePayload(
            id=resp["response_id"],
            object="chat.completion",
            created=int(time.time()),
            model=config.model.name,
            choices=[
                chat.Choice(
                    index=0,
                    message=chat.ResponseMessage(
                        role="assistant",
                        content=resp["text"],
                    ),
                    finish_reason=None,
                ),
            ],
            usage=chat.ChatUsage(
                prompt_tokens=resp["token_count"]["prompt_tokens"],
                completion_tokens=resp["token_count"]["response_tokens"],
                total_tokens=resp["token_count"]["total_tokens"],
            ),
        )

    @classmethod
    def model_to_chat_streaming(cls, resp, config):
        # Response example (https://docs.cohere.com/reference/chat)
        # Streaming chunks:
        # ```
        # {
        #   "is_finished":false,
        #   "event_type":"stream-start",
        #   "generation_id":"string"
        # }
        # {"is_finished":false,"event_type":"text-generation","text":"How"}
        # {"is_finished":false,"event_type":"text-generation","text":" are"}
        # {"is_finished":false,"event_type":"text-generation","text":" you"}
        # {
        #   "is_finished":true,
        #   "event_type":"stream-end",
        #   "response":{
        #     "response_id":"string",
        #     "text":"How are you",
        #     "generation_id":"string",
        #     "token_count":{
        #       "prompt_tokens":83,"response_tokens":63,"total_tokens":146,"billed_tokens":128
        #     },
        #     "tool_inputs":null
        #   },
        #   "finish_reason":"COMPLETE"
        # }
        # ```
        response = resp.get("response")
        return chat.StreamResponsePayload(
            # first chunk has "generation_id" but not "response_id"
            id=response["response_id"] if response else None,
            created=int(time.time()),
            model=config.model.name,
            choices=[
                chat.StreamChoice(
                    index=0,
                    finish_reason=resp.get("finish_reason"),
                    delta=chat.StreamDelta(
                        role=None,
                        content=resp.get("text"),
                    ),
                )
            ],
            usage=chat.ChatUsage(
                prompt_tokens=response["token_count"]["prompt_tokens"] if response else None,
                completion_tokens=response["token_count"]["response_tokens"] if response else None,
                total_tokens=response["token_count"]["total_tokens"] if response else None,
            ),
        )


class CohereProvider(BaseProvider):
    NAME = "Cohere"
    CONFIG_TYPE = CohereConfig

    def __init__(self, config: RouteConfig) -> None:
        super().__init__(config)
        if config.model.config is None or not isinstance(config.model.config, CohereConfig):
            raise TypeError(f"Unexpected config type {config.model.config}")
        self.cohere_config: CohereConfig = config.model.config

    @property
    def headers(self) -> dict[str, str]:
        return {"Authorization": f"Bearer {self.cohere_config.cohere_api_key}"}

    @property
    def base_url(self) -> str:
        return "https://api.cohere.ai/v1"

    @property
    def adapter_class(self) -> type[ProviderAdapter]:
        return CohereAdapter

    def get_endpoint_url(self, route_type: str) -> str:
        if route_type == "llm/v1/chat":
            return f"{self.base_url}/chat"
        elif route_type == "llm/v1/completions":
            return f"{self.base_url}/generate"
        elif route_type == "llm/v1/embeddings":
            return f"{self.base_url}/embed"
        else:
            raise ValueError(f"Invalid route type {route_type}")

    async def _request(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        return await send_request(
            headers=self.headers,
            base_url=self.base_url,
            path=path,
            payload=payload,
        )

    def _stream_request(self, path: str, payload: dict[str, Any]) -> AsyncGenerator[bytes, None]:
        return send_stream_request(
            headers=self.headers,
            base_url=self.base_url,
            path=path,
            payload=payload,
        )

    async def chat_stream(
        self, payload: chat.RequestPayload
    ) -> AsyncIterable[chat.StreamResponsePayload]:
        from fastapi.encoders import jsonable_encoder

        payload = jsonable_encoder(payload, exclude_none=True)
        self.check_for_model_field(payload)
        stream = self._stream_request(
            "chat",
            {
                "model": self.config.model.name,
                **CohereAdapter.chat_streaming_to_model(payload, self.config),
            },
        )
        async for chunk in stream:
            if not chunk:
                continue

            resp = json.loads(chunk)
            if resp["event_type"] == "stream-start":
                continue
            yield CohereAdapter.model_to_chat_streaming(resp, self.config)

    async def chat(self, payload: chat.RequestPayload) -> chat.ResponsePayload:
        from fastapi.encoders import jsonable_encoder

        payload = jsonable_encoder(payload, exclude_none=True)
        self.check_for_model_field(payload)
        resp = await self._request(
            "chat",
            {
                "model": self.config.model.name,
                **CohereAdapter.chat_to_model(payload, self.config),
            },
        )
        return CohereAdapter.model_to_chat(resp, self.config)

    async def completions_stream(
        self, payload: completions.RequestPayload
    ) -> AsyncIterable[completions.StreamResponsePayload]:
        from fastapi.encoders import jsonable_encoder

        payload = jsonable_encoder(payload, exclude_none=True)
        self.check_for_model_field(payload)
        stream = self._stream_request(
            "generate",
            {
                "model": self.config.model.name,
                **CohereAdapter.completions_streaming_to_model(payload, self.config),
            },
        )
        async for chunk in stream:
            if not chunk:
                continue

            resp = json.loads(chunk)
            yield CohereAdapter.model_to_completions_streaming(resp, self.config)

    async def completions(self, payload: completions.RequestPayload) -> completions.ResponsePayload:
        from fastapi.encoders import jsonable_encoder

        payload = jsonable_encoder(payload, exclude_none=True)
        self.check_for_model_field(payload)
        resp = await self._request(
            "generate",
            {
                "model": self.config.model.name,
                **CohereAdapter.completions_to_model(payload, self.config),
            },
        )
        return CohereAdapter.model_to_completions(resp, self.config)

    async def embeddings(self, payload: embeddings.RequestPayload) -> embeddings.ResponsePayload:
        from fastapi.encoders import jsonable_encoder

        payload = jsonable_encoder(payload, exclude_none=True)
        self.check_for_model_field(payload)
        resp = await self._request(
            "embed",
            {
                "model": self.config.model.name,
                **CohereAdapter.embeddings_to_model(payload, self.config),
            },
        )
        return CohereAdapter.model_to_embeddings(resp, self.config)
