"""Utilities to create langchain LLM and Chat Model instances."""
from abc import ABC, abstractmethod
from typing import Any
from langchain_core.embeddings import Embeddings
from langchain_core.language_models import BaseChatModel, BaseLLM
from langchain_community.chat_models.fake import FakeListChatModel
from langchain_community.llms.fake import FakeListLLM
from langchain_community.llms.loading import load_llm_from_config
from langchain_community.embeddings.fake import FakeEmbeddings
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from openai import OpenAI
from brevia.settings import get_settings
from brevia.utilities.types import load_type


class FakeBreviaLLM(FakeListLLM):
    """Fake LLM for testing purposes."""

    def get_token_ids(self, text: str) -> list[int]:
        """Fake for testing purposes."""
        return [10] * 10


LOREM_IPSUM = """Lorem ipsum dolor sit amet, consectetur adipisici elit,
sed eiusmod tempor incidunt ut labore et dolore magna aliqua."""


def load_llm(config: dict) -> BaseLLM:
    """Load langchain LLM, use Fake LLM in test mode"""
    if test_models_in_use():
        return FakeBreviaLLM(responses=[LOREM_IPSUM] * 10)

    return load_llm_from_config(config=config)


class FakeBreviaChatModel(FakeListChatModel):
    """Fake Chat Model for testing purposes."""

    def get_token_ids(self, text: str) -> list[int]:
        """Fake for testing purposes."""
        return [10] * 10


def load_chatmodel(config: dict) -> BaseChatModel:
    """Load Chat Model from Config Dict."""
    if test_models_in_use():
        return FakeBreviaChatModel(responses=[LOREM_IPSUM] * 10)

    chatmodel_aliases = {
        'openai-chat': ChatOpenAI,
        'fake-list-chat-model': FakeListChatModel,
    }
    model_type = config.pop('_type')
    if model_type in chatmodel_aliases:
        llm_cls = chatmodel_aliases[model_type]
    else:
        llm_cls = load_type(
            model_type,
            BaseChatModel,
            'langchain_community.chat_models',
        )

    return llm_cls(**config)


class BaseAudio(ABC):
    """Base class for Audio transcription"""
    @abstractmethod
    def transcribe(self, file: Any, **kwargs) -> dict:
        """transcribe Audio"""


class FakeAudio(BaseAudio):
    """Test class for Audio transcription"""

    def transcribe(self, file: Any, **kwargs) -> dict:
        """transcribe Audio"""
        return {'text': LOREM_IPSUM}


class AudioOpenAI(BaseAudio):
    """OpenAI Audio transcription"""

    def transcribe(self, file: Any, **kwargs) -> dict:
        """transcribe Audio"""
        client = OpenAI()
        audio_file = open(file, "rb")
        return client.audio.transcriptions.create(
            model="whisper-1",
            file=audio_file,
            **kwargs
        )


def load_audiotranscriber() -> BaseAudio:
    """Load Audio transcriber (only openAI supported for now)."""
    if test_models_in_use():
        return FakeAudio()

    return AudioOpenAI()


def load_embeddings(custom_conf: dict | None = None) -> Embeddings:
    """ Load Embeddings engine """
    settings = get_settings()
    if test_models_in_use():
        return FakeEmbeddings(size=1536)

    config = settings.embeddings.copy() if not custom_conf else custom_conf
    embed_aliases = {
        'openai-embeddings': OpenAIEmbeddings,
        'fake-embeddings': FakeEmbeddings,
    }
    model_type = config.pop('_type')
    if model_type in embed_aliases:
        emb_cls = embed_aliases[model_type]
    else:
        emb_cls = load_type(
            model_type,
            Embeddings,
            'langchain_community.embeddings',
        )

    return emb_cls(**config)


def test_models_in_use() -> bool:
    """Check if test models are in use (via `USE_TEST_MODELS` env var)"""
    return get_settings().use_test_models
