import logging
import os

from achatbot.common.types import (
    DEFAULT_SYSTEM_PROMPT, MODELS_DIR,
    PersonalAIProxyArgs)
from achatbot.types.llm.transformers import TransformersLMArgs
from achatbot.common import interface
from achatbot.common.factory import EngineClass, EngineFactory

from dotenv import load_dotenv
load_dotenv(override=True)


class LLMEnvInit():

    @staticmethod
    def getEngine(tag, **kwargs) -> interface.IAsr | EngineClass:
        if "llm_llamacpp" in tag:
            from . import llamacpp
        elif "llm_personalai_proxy" in tag:
            from . import personalai
        elif "llm_transformers_manual_vision_qwen" in tag:
            from .transformers import manual_vision_qwen
        elif "llm_transformers_manual_vision_llama" in tag:
            from .transformers import manual_vision_llama
        elif "llm_transformers_manual_vision_molmo" in tag:
            from .transformers import manual_vision_molmo
        elif "llm_transformers_manual" in tag:
            from .transformers import manual
        elif "llm_transformers_pipeline" in tag:
            from .transformers import pipeline

        engine = EngineFactory.get_engine_by_tag(EngineClass, tag, **kwargs)
        return engine

    @staticmethod
    def initLLMEngine(
            tag: str | None = None,
            kwargs: dict | None = None) -> interface.ILlm | EngineClass:
        # llm
        tag = tag if tag else os.getenv('LLM_TAG', "llm_llamacpp")
        if kwargs is None:
            kwargs = LLMEnvInit.map_config_func[tag]()
        engine = LLMEnvInit.getEngine(tag, **kwargs)
        logging.info(f"initLLMEngine: {tag}, {engine}")
        return engine

    @staticmethod
    def get_llm_llamacpp_args() -> dict:
        kwargs = {}
        kwargs["model_name"] = os.getenv('LLM_MODEL_NAME', 'qwen2')
        kwargs["model_path"] = os.getenv('LLM_MODEL_PATH', os.path.join(
            MODELS_DIR, "qwen2-1_5b-instruct-q8_0.gguf"))
        kwargs["model_type"] = os.getenv('LLM_MODEL_TYPE', "chat")
        kwargs["n_threads"] = os.cpu_count()
        kwargs["verbose"] = False
        kwargs["n_gpu_layers"] = int(os.getenv('N_GPU_LAYERS', "0"))
        kwargs["flash_attn"] = bool(os.getenv('FLASH_ATTN', ""))
        kwargs["llm_stream"] = True
        # if logger.getEffectiveLevel() != logging.DEBUG:
        #    kwargs["verbose"] = False
        kwargs['llm_stop'] = [
            "<|end|>", "<|im_end|>", "<|endoftext|>", "<{/end}>",
            "</s>", "/s>", "</s", "<s>",
            "<|user|>", "<|assistant|>", "<|system|>",
        ]
        kwargs["llm_chat_system"] = os.getenv('LLM_CHAT_SYSTEM', DEFAULT_SYSTEM_PROMPT)
        kwargs["llm_tool_choice"] = os.getenv('LLM_TOOL_CHOICE', None)
        kwargs["chat_format"] = os.getenv('LLM_CHAT_FORMAT', None)
        kwargs["tokenizer_path"] = os.getenv('LLM_TOKENIZER_PATH', None)
        kwargs["clip_model_path"] = os.getenv('LLM_CLIP_MODEL_PATH', None)
        return kwargs

    @staticmethod
    def get_llm_personal_ai_proxy_args() -> dict:
        kwargs = PersonalAIProxyArgs(
            api_url=os.getenv('API_URL', "http://localhost:8787/"),
            chat_bot=os.getenv('CHAT_BOT', "openai"),
            model_type=os.getenv('CHAT_TYPE', "chat_only"),
            openai_api_base_url=os.getenv('OPENAI_API_BASE_URL', "https://api.groq.com/openai/v1/"),
            model_name=os.getenv('LLM_MODEL_NAME', "llama3-70b-8192"),
            llm_chat_system=os.getenv('LLM_CHAT_SYSTEM', DEFAULT_SYSTEM_PROMPT),
            llm_stream=False,
            llm_max_tokens=int(os.getenv('LLM_MAX_TOKENS', "1024")),
            func_search_name=os.getenv('FUNC_SEARCH_NAME', "search_api"),
            func_weather_name=os.getenv('FUNC_WEATHER_NAME', "openweathermap_api"),
        ).__dict__
        return kwargs

    @staticmethod
    def _get_llm_transformers_args() -> dict:
        kwargs = TransformersLMArgs(
            lm_model_name_or_path=os.getenv(
                'LLM_MODEL_NAME_OR_PATH',
                os.path.join(MODELS_DIR, "Qwen/Qwen2-0.5B-Instruct")),
            lm_attn_impl=os.getenv('LLM_ATTN_IMPL', None),
            lm_device=os.getenv('LLM_DEVICE', "cpu"),
            lm_torch_dtype=os.getenv('LLM_TORCH_DTYPE', 'auto'),
            lm_gen_max_new_tokens=int(os.getenv('LLM_GEN_MAX_NEW_TOKENS', '1024')),
            lm_gen_min_new_tokens=int(os.getenv('LLM_GEN_MIN_NEW_TOKENS', '0')),
            lm_gen_do_sample=bool(os.getenv('LLM_GEN_DO_SAMPLE', '')),
            lm_gen_temperature=float(os.getenv('LLM_GEN_TEMPERATURE', '0.0')),
            lm_stream=bool(os.getenv('LLM_STREAM', '1')),
            init_chat_prompt=os.getenv('LLM_INIT_CHAT_PROMPT', ''),
            chat_history_size=int(os.getenv('LLM_CHAT_HISTORY_SIZE', '10')),  # cache 10 round
            model_type=os.getenv('LLM_MODEL_TYPE', "chat_completion")
        ).__dict__
        return kwargs

    @staticmethod
    def get_llm_transformers_args() -> dict:
        hf_parser_file = os.getenv('HF_TRANSFOMERS_PARSER_FILE', None)
        if hf_parser_file:
            from .transformers import TransformersLLMInit
            return TransformersLLMInit.get_transformers_llm_args(hf_parser_file)
        return LLMEnvInit._get_llm_transformers_args()

    # TAG : config
    map_config_func = {
        'llm_llamacpp': get_llm_llamacpp_args,
        'llm_personalai_proxy': get_llm_personal_ai_proxy_args,
        'llm_transformers_manual': get_llm_transformers_args,
        'llm_transformers_pipeline': get_llm_transformers_args,
        'llm_transformers_manual_vision_qwen': get_llm_transformers_args,
        'llm_transformers_manual_vision_llama': get_llm_transformers_args,
        'llm_transformers_manual_vision_molmo': get_llm_transformers_args,
    }
