import logging
import time
from abc import abstractmethod
from typing import List, Optional

from mirix.errors import LLMError
from mirix.observability.context import get_trace_context
from mirix.observability.langfuse_client import get_langfuse_client
from mirix.schemas.llm_config import LLMConfig
from mirix.schemas.message import Message
from mirix.schemas.openai.chat_completion_response import ChatCompletionResponse
from mirix.services.cloud_file_mapping_manager import CloudFileMappingManager
from mirix.services.file_manager import FileManager


class LLMClientBase:
    """
    Abstract base class for LLM clients, formatting the request objects,
    handling the downstream request and parsing into chat completions response format
    """

    def __init__(
        self,
        llm_config: LLMConfig,
        use_tool_naming: bool = True,
    ):
        self.llm_config = llm_config
        self.use_tool_naming = use_tool_naming
        self.file_manager = FileManager()
        self.cloud_file_mapping_manager = CloudFileMappingManager()
        self.logger = logging.getLogger(f"Mirix.LLMClientBase")

    def send_llm_request(
        self,
        messages: List[Message],
        tools: Optional[List[dict]] = None,
        stream: bool = False,
        force_tool_call: Optional[str] = None,
        get_input_data_for_debugging: bool = False,
        existing_file_uris: Optional[List[str]] = None,
    ) -> ChatCompletionResponse:
        """
        Issues a request to the downstream model endpoint and parses response.
        """
        request_data = self.build_request_data(
            messages,
            self.llm_config,
            tools,
            force_tool_call,
            existing_file_uris=existing_file_uris,
        )

        if get_input_data_for_debugging:
            return request_data

        # ========== LANGFUSE INSTRUMENTATION ==========
        langfuse = get_langfuse_client()
        generation = None
        trace_context = get_trace_context() if langfuse else {}

        if langfuse:
            try:
                # Prepare input for tracing
                messages_for_trace = [
                    {
                        "role": (
                            m.role.value if hasattr(m.role, "value") else str(m.role)
                        ),
                        "content": str(m.text) if hasattr(m, "text") else str(m),
                    }
                    for m in messages
                ]

                model_parameters = {}
                if (
                    hasattr(self.llm_config, "temperature")
                    and self.llm_config.temperature is not None
                ):
                    model_parameters["temperature"] = self.llm_config.temperature

                if trace_context.get("trace_id"):
                    self.logger.info(
                        f"LLM generation linked to trace_id={trace_context.get('trace_id')}"
                    )

                generation = langfuse.generation(
                    name="llm_completion",
                    model=self.llm_config.model,
                    model_parameters=model_parameters,
                    input=messages_for_trace,
                    metadata={
                        "provider": self.llm_config.model_endpoint_type,
                        "tools_count": len(tools) if tools else 0,
                    },
                    trace_id=trace_context.get("trace_id"),
                    parent_observation_id=trace_context.get("observation_id"),
                    user_id=trace_context.get("user_id"),
                    session_id=trace_context.get("session_id"),
                )
            except Exception as e:
                self.logger.warning(f"Failed to create LangFuse generation: {e}")

        try:
            t1 = time.time()
            response_data = self.request(request_data)
            t2 = time.time()
            self.logger.debug("LLM request time: %.2f seconds", t2 - t1)
        except Exception as e:
            # Update LangFuse generation with error
            if generation:
                try:
                    generation.update(
                        status_message=str(e),
                        level="ERROR",
                        metadata={"error_type": type(e).__name__},
                    )
                    generation.end()
                except Exception:
                    pass
            raise self.handle_llm_error(e)

        chat_completion_data = self.convert_response_to_chat_completion(
            response_data, messages
        )

        # Update LangFuse generation with success
        if generation:
            try:
                output_message = None
                if (
                    hasattr(chat_completion_data, "choices")
                    and len(chat_completion_data.choices) > 0
                ):
                    choice = chat_completion_data.choices[0]
                    if hasattr(choice, "message"):
                        msg = choice.message
                        output_message = {
                            "role": msg.role if hasattr(msg, "role") else "assistant",
                            "content": (
                                str(msg.content) if hasattr(msg, "content") else ""
                            ),
                        }

                usage_dict = None
                if (
                    hasattr(chat_completion_data, "usage")
                    and chat_completion_data.usage
                ):
                    usage_dict = {
                        "input": getattr(
                            chat_completion_data.usage, "prompt_tokens", 0
                        ),
                        "output": getattr(
                            chat_completion_data.usage, "completion_tokens", 0
                        ),
                        "total": getattr(chat_completion_data.usage, "total_tokens", 0),
                    }

                generation.update(
                    output=output_message,
                    usage=usage_dict,
                )
                generation.end()
            except Exception as e:
                self.logger.warning(f"Failed to update LangFuse generation: {e}")

        return chat_completion_data

    @abstractmethod
    def build_request_data(
        self,
        messages: List[Message],
        llm_config: LLMConfig,
        tools: Optional[List[dict]] = None,
        force_tool_call: Optional[str] = None,
        existing_file_uris: Optional[List[str]] = None,
    ) -> dict:
        """
        Constructs a request object in the expected data format for this client.
        """
        raise NotImplementedError

    @abstractmethod
    def request(self, request_data: dict) -> dict:
        """
        Performs underlying request to llm and returns raw response.
        """
        raise NotImplementedError

    @abstractmethod
    def convert_response_to_chat_completion(
        self,
        response_data: dict,
        input_messages: List[Message],
    ) -> ChatCompletionResponse:
        """
        Converts custom response format from llm client into an OpenAI
        ChatCompletionsResponse object.
        """
        raise NotImplementedError

    @abstractmethod
    def handle_llm_error(self, e: Exception) -> Exception:
        """
        Maps provider-specific errors to common LLMError types.
        Each LLM provider should implement this to translate their specific errors.

        Args:
            e: The original provider-specific exception

        Returns:
            An LLMError subclass that represents the error in a provider-agnostic way
        """
        return LLMError(f"Unhandled LLM error: {str(e)}")
