from typing import Any, List, Optional, Iterator, AsyncIterator
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatResult, ChatGeneration, ChatGenerationChunk
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from pydantic import Field

from trim.utils import Trim
from langchain_openai import ChatOpenAI


class TrimChatModel(ChatOpenAI):
    trim_client: Trim = Field(default_factory=Trim)
    llm: ChatOpenAI = Field(default_factory=ChatOpenAI)

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)

    @property
    def _llm_type(self) -> str:
        return "trim-chat-model"

    def _generate(self,
                  messages: List[BaseMessage],
                  stop: Optional[List[str]] = None,
                  run_manager: Optional[CallbackManagerForLLMRun] = None,
                  **kwargs: Any) -> ChatResult:
        routing_response = self.trim_client.route_llm_request(
            prompt=str(messages),
            intended_model_id=None
        )

        response = self.llm.invoke(
            model=routing_response.chosen_model_id,
            input=messages,
            **kwargs
        )

        n_input_tokens = response.usage_metadata['input_tokens']
        n_output_tokens = response.usage_metadata['output_tokens']

        cache_response = self.trim_client.cache_llm_response(
            request_id=routing_response.request_id,
            response=response.content if response.content else "None",
            number_of_input_tokens=n_input_tokens,
            number_of_output_tokens=n_output_tokens
        )

        gen = ChatGeneration(message=response)
        return ChatResult(generations=[gen])

    def _stream(self,
                messages: List[BaseMessage],
                stop: Optional[List[str]] = None,
                run_manager: Optional[CallbackManagerForLLMRun] = None,
                **kwargs: Any) -> Iterator[ChatGenerationChunk]:
        raise NotImplementedError()

    async def _agenerate(self, *args, **kwargs) -> ChatResult:
        return self._generate(*args, **kwargs)

    async def _astream(self, *args, **kwargs) -> AsyncIterator[ChatGenerationChunk]:
        raise NotImplementedError()
