import os
from pathlib import Path
from typing import Any

import tiktoken
from litellm import ModelResponse  # type: ignore[import]
from llamux import Router  # type: ignore[import]
from loguru import logger

from notte_core.common.config import LlmModel, PerceptionType, config
from notte_core.errors.llm import InvalidPromptTemplateError
from notte_core.llms.engine import LLMEngine
from notte_core.llms.prompt import PromptLibrary
from notte_core.llms.types import TResponseFormat

PROMPT_DIR = Path(__file__).parent.parent / "llms" / "prompts"
LLAMUX_CONFIG = Path(__file__).parent.parent / "llms" / "config" / "endpoints.csv"


def get_llamux_config(verbose: bool = False) -> str:
    if "LLAMUX_CONFIG_PATH" in os.environ:
        if verbose:
            logger.debug(f"Using custom LLAMUX config path: {os.environ['LLAMUX_CONFIG_PATH']}")
    else:
        if verbose:
            logger.debug(f"Using default LLAMUX config path: {LLAMUX_CONFIG}")
    return os.getenv("LLAMUX_CONFIG_PATH", str(LLAMUX_CONFIG))


class LLMService:
    """
    LLM service for Notte.
    """

    def __init__(self, base_model: str | None = None) -> None:
        self.lib: PromptLibrary = PromptLibrary(str(PROMPT_DIR))
        self.router: Router | None = None

        if config.use_llamux:
            llamux_config = get_llamux_config(config.verbose)
            path = Path(llamux_config)
            if not path.exists():
                raise FileNotFoundError(f"LLAMUX config file not found at {path}")
            self.router = Router.from_csv(llamux_config)
        self.base_model: str = base_model or LlmModel.default()
        self.tokenizer: tiktoken.Encoding = tiktoken.get_encoding("cl100k_base")
        self.verbose: bool = config.verbose
        self.nb_retries_structured_output: int = config.nb_retries_structured_output

    @staticmethod
    def from_config(perception_type: PerceptionType = config.perception_type) -> "LLMService":
        model = config.perception_model
        if model is None:
            model = config.reasoning_model
            if perception_type == "deep":
                logger.warning(f"No perception model set, using reasoning model: {config.reasoning_model}")
        return LLMService(base_model=model)

    def context_length(self) -> int:
        return LlmModel.get_provider(self.base_model).context_length

    def get_base_model(self, messages: list[dict[str, Any]]) -> tuple[str, str | None]:
        eid: str | None = None

        if self.router is not None:
            router = "llamux"
            provider, model, eid, _ = self.router.query(messages=messages)
            base_model = f"{provider}/{model}"
        else:
            router = "fixed"
            base_model = self.base_model

        token_len = self.estimate_tokens(text="\n".join([m["content"] for m in messages]))
        if self.verbose:
            logger.debug(f"llm router '{router}' selected '{base_model}' for approx {token_len} tokens")
        return base_model, eid

    def clip_tokens(self, document: str, max_tokens: int | None = None) -> str:
        max_tokens = max_tokens or (self.context_length() - 2000)
        tokens = self.tokenizer.encode(document)
        if len(tokens) > max_tokens:
            logger.debug(f"Cannot process document, exceeds max tokens: {len(tokens)} > {max_tokens}. Clipping...")
            return self.tokenizer.decode(tokens[:max_tokens])
        return document

    def estimate_tokens(
        self, text: str | None = None, prompt_id: str | None = None, variables: dict[str, Any] | None = None
    ) -> int:
        if text is None:
            if prompt_id is None or variables is None:
                raise InvalidPromptTemplateError(
                    prompt_id=prompt_id or "unknown",
                    message="for token estimation, prompt_id and variables must be provided if text is not provided",
                )
            messages = self.lib.materialize(prompt_id, variables)
            text = "\n".join([m["content"] for m in messages])
        return len(self.tokenizer.encode(text))

    async def structured_completion(
        self,
        prompt_id: str,
        response_format: type[TResponseFormat],
        variables: dict[str, Any] | None = None,
        use_strict_response_format: bool = True,
    ) -> TResponseFormat:
        messages = self.lib.materialize(prompt_id, variables)
        base_model, _ = self.get_base_model(messages)
        return await LLMEngine(
            nb_retries_structured_output=self.nb_retries_structured_output, verbose=self.verbose
        ).structured_completion(
            messages=messages,  # type: ignore[arg-type]
            response_format=response_format,
            model=base_model,
            use_strict_response_format=use_strict_response_format,
        )

    async def completion(
        self,
        prompt_id: str,
        variables: dict[str, Any] | None = None,
    ) -> ModelResponse:
        messages = self.lib.materialize(prompt_id, variables)
        base_model, eid = self.get_base_model(messages)
        response = await LLMEngine(verbose=self.verbose).completion(
            messages=messages,  # type: ignore[arg-type]
            model=base_model,
        )
        if eid is not None and self.router is not None:
            # log usage to LLAMUX router if eid is provided
            tokens: int = response.usage.total_tokens  # type: ignore[attr-defined]
            self.router.log(tokens=tokens, endpoint_id=eid)  # type: ignore[arg-type]
        return response
