import logging
import uuid
from dataclasses import fields
from typing import Any, Dict, List, Optional, Sequence, Union
from uuid import UUID

# Use LangChain features
from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.documents import Document
from langchain_core.messages import BaseMessage, ToolMessage
from langchain_core.outputs import LLMResult

from ...expiring_key_value_store import ExpiringKeyValueStore
from ...logger import (
    GenerationConfig,
    Logger,
    RetrievalConfig,
    SpanConfig,
    ToolCallConfig,
    ToolCallError,
)
from ..models import Container, Metadata, SpanContainer, TraceContainer
from .utils import (
    parse_langchain_llm_error,
    parse_langchain_llm_result,
    parse_langchain_messages,
    parse_langchain_model_and_provider,
    parse_langchain_model_parameters,
    parse_langchain_provider,
)

logging = logging.getLogger("MaximSDK")


# 20 minutes
DEFAULT_TIMEOUT = 60 * 20


class MaximLangchainTracer(BaseCallbackHandler):
    """
    A callback handler that logs langchain outputs to Maxim logger

    Args:
        logger: Logger: Maxim Logger instance to log outputs
    """

    def __init__(
        self, logger: Logger, metadata: Optional[Dict[str, Any]] = None
    ) -> None:
        """Initializes the Langchain Tracer
        Args:
            logger: Logger: Maxim Logger instance to log outputs
        """
        super().__init__()
        self.run_inline = True
        self.logger = logger
        self.containers: Dict[str, Container] = {}
        self.metadata_store = ExpiringKeyValueStore()
        self.metadata = None
        if metadata is not None:
            self.__validate_maxim_metadata(metadata)
            self.metadata = metadata

    def __validate_maxim_metadata(self, metadata: Optional[Dict[str, Any]]):
        """
        Validates the metadata
        """
        if metadata is None:
            return
        id_keys = ["session_id", "trace_id", "span_id"]
        present_keys = [key for key in id_keys if key in metadata]
        if len(present_keys) > 1:
            raise ValueError(
                f"Multiple keys found in metadata: {present_keys}. You can pass only one of these."
            )
        valid_keys = [field.name for field in fields(Metadata)]
        invalid_keys = [key for key in metadata if key not in valid_keys]
        if len(invalid_keys) > 0:
            raise ValueError(
                f"Invalid keys found in metadata: {invalid_keys}. Valid keys are {valid_keys}"
            )

    def __get_container_from_metadata(self, run_id: UUID) -> Container:
        """
        Gets the container from the metadata
        """
        maxim_metadata = self.__get_metadata()
        container: Optional[Container] = None
        if maxim_metadata is not None:
            span_id: Optional[str] = maxim_metadata.span_id
            if span_id is not None:
                span_name: Optional[str] = maxim_metadata.span_name
                container = SpanContainer(
                    span_id=span_id,
                    logger=self.logger,
                    span_name=span_name,
                    mark_created=True,
                )
            trace_id = maxim_metadata.trace_id
            if trace_id is not None:
                trace_name = maxim_metadata.trace_name
                container = TraceContainer(
                    trace_id=trace_id,
                    logger=self.logger,
                    trace_name=trace_name,
                    mark_created=True,
                )

        if container is None:
            session_id: Optional[str] = (
                maxim_metadata.session_id if maxim_metadata else None
            )
            container = TraceContainer(
                trace_id=str(uuid.uuid4()),
                logger=self.logger,
                trace_name="Trace",
                parent=session_id,
            )
            container.create()

        return container

    def __get_container(
        self, run_id: UUID, parent_run_id: Optional[UUID] = None
    ) -> Optional[Container]:
        container: Optional[Container] = None
        if parent_run_id is None:
            # This is the first activity in this run
            container = self.__get_container_from_metadata(run_id)
        else:
            container = self.containers.get(str(parent_run_id))
        return container

    def __get_metadata(
        self, metadata: Optional[Dict[str, Any]] = None
    ) -> Optional[Metadata]:
        if metadata is not None and metadata.get("maxim", None) is not None:
            return Metadata(metadata.get("maxim", None))
        return Metadata(self.metadata)

    def on_llm_start(
        self,
        serialized: dict[str, Any],
        prompts: list[str],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        metadata: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Runs when LLM starts"""
        logging.debug("[MaximSDK: Langchain] on_llm_start called")
        model, model_parameters = parse_langchain_model_parameters(**kwargs)
        provider = parse_langchain_provider(serialized)
        model, provider = parse_langchain_model_and_provider(model, provider)
        messages = parse_langchain_messages(prompts)
        generation_id = str(run_id)
        maxim_metadata = self.__get_metadata(metadata)
        # Prepare generation config
        generation_config = GenerationConfig(
            id=generation_id,
            name=maxim_metadata.generation_name if maxim_metadata else None,
            provider=provider,
            model=model,
            messages=messages,
            model_parameters=model_parameters,
            tags=maxim_metadata.generation_tags if maxim_metadata else None,
        )
        # Adding it back to the container
        container = self.__get_container(run_id, parent_run_id)
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for generation")
            return
        if not container.is_created():
            container.create()
        container.add_generation(generation_config)

    def on_chat_model_start(
        self,
        serialized: Dict[str, Any],
        messages: List[List[BaseMessage]],
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Runs when a chat model call starts"""
        logging.debug("[Langchain] on_chat_model_start called")
        run_id = kwargs.get("run_id", None)
        model, model_parameters = parse_langchain_model_parameters(**kwargs)
        provider = parse_langchain_provider(serialized)
        model, provider = parse_langchain_model_and_provider(model, provider)
        maxim_messages = parse_langchain_messages(messages)
        generation_id = str(run_id)
        maxim_metadata = self.__get_metadata(metadata)
        generation_name = None
        if kwargs.get("parent_run_id", None) is None:
            generation_name = maxim_metadata.generation_name if maxim_metadata else None
        # Checking if generation id already present for this run_id
        generation_config = GenerationConfig(
            id=generation_id,
            name=generation_name,
            provider=provider,
            model=model,
            messages=maxim_messages,
            model_parameters=model_parameters,
            tags=maxim_metadata.generation_tags if maxim_metadata else None,
        )
        container = self.__get_container(run_id, kwargs.get("parent_run_id", None))
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for  generation")
            return
        if not container.is_created():
            container.create()
        container.add_generation(generation_config)
        if container.parent() is None:
            self.containers[str(run_id)] = container

    def on_llm_new_token(self, token: str, **kwargs: Any) -> Any:
        """Run on new LLM token. Only available when streaming is enabled."""

    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> Any:
        """Run when LLM ends running."""
        logging.debug("[MaximSDK: Langchain] on_llm_end called")
        run_id = kwargs.get("run_id", None)
        parent_run_id = kwargs.get("parent_run_id", None)
        result = parse_langchain_llm_result(response)
        container = self.__get_container(run_id, parent_run_id)
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for generation")
            return
        self.logger.generation_result(str(run_id), result)
        # We close the container - only if it was created by this run_id
        if container.parent() is None:
            container.end()
            self.containers.pop(str(run_id))

    def on_llm_error(
        self, error: Union[Exception, BaseException, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        """Run when LLM errors."""
        logging.debug("[Langchain] on_llm_error called")
        run_id = kwargs.get("run_id", None)
        parent_run_id = kwargs.get("parent_run_id", None)
        container = self.__get_container(run_id, parent_run_id)
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for generation")
            return
        generation_error = parse_langchain_llm_error(error)
        self.logger.generation_error(str(run_id), generation_error)
        # We close the container - only if it was created by this run_id
        if container.parent() is None:
            container.end()
            self.containers.pop(str(run_id))

    def on_retriever_start(
        self,
        serialized: dict[str, Any],
        query: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        metadata: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """
        Run when Retriever starts running.
        """
        logging.debug("[Langchain] on_retriever_start called")
        retrieval_id = str(run_id)
        retrieval_config = RetrievalConfig(id=retrieval_id)
        container = self.__get_container(run_id, parent_run_id)
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for retrieval")
            return
        if not container.is_created():
            container.create()
        retrieval = container.add_retrieval(retrieval_config)
        retrieval.input(query)

    def on_retriever_end(
        self,
        documents: Sequence[Document],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        logging.debug("[Langchain] on_retriever_end called")
        container = self.__get_container(run_id, parent_run_id)
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for retrieval")
            return
        self.logger.retrieval_output(str(run_id), documents)
        if container.parent is None:
            container.end()
            self.containers.pop(str(run_id))

    # Chain callbacks

    def on_chain_start(
        self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> Any:
        try:
            run_id = kwargs.get("run_id", None)
            # Creating new span for the chain node
            name = kwargs.get("name", None)
            tags: Dict[str, str] = {}
            if (
                self.metadata is not None
                and self.metadata.get("chain_tags") is not None
            ):
                # its a dict
                for key, value in self.metadata["chain_tags"].items():
                    tags[key] = str(value)
            tags["run_id"] = run_id
            if kwargs.get("parent_run_id", None) is not None:
                tags["parent_run_id"] = str(kwargs.get("parent_run_id", None))
            if kwargs.get("tags", None) is not None:
                for tag in kwargs["tags"]:
                    key, value = tag.split(":", 1)
                    tags[key.strip()] = value.strip()
            if kwargs.get("metadata", None) is not None:
                for key, value in kwargs["metadata"].items():
                    tags[key.strip()] = str(value)

            container = self.__get_container(run_id, kwargs.get("parent_run_id", None))

            if container is None:
                logging.error("[MaximSDK] Couldn't find a container for chain")
                return

            if not container.is_created():
                container.create()

            span_config = SpanConfig(id=str(run_id), name=name, tags=tags)
            container.add_span(span_config)
            self.containers[str(run_id)] = SpanContainer(
                span_id=str(run_id), logger=self.logger, mark_created=True
            )
        except Exception as e:
            logging.error("[MaximSDK] Error while processing chain_start %s", e)

    def on_chain_end(self, outputs: Union[str, Dict[str, Any]], **kwargs: Any) -> Any:
        try:
            run_id = kwargs.get("run_id", None)
            parent_run_id = kwargs.get("parent_run_id", None)
            # We hide the hidden chains
            tags = {
                tag.split(":")[0]: tag.split(":")[1]
                for tag in kwargs.get("tags", [])
                if ":" in tag
            }
            container = self.__get_container(run_id, parent_run_id)
            if container is None:
                logging.error("[MaximSDK] Couldn't find a container for chain")
                return

            container.add_tags(tags)
            container.end()
            self.containers.pop(str(run_id))

        except Exception as e:
            logging.error("[MaximSDK] Failed to parse chain-end: %s", e)

    def on_chain_error(
        self, error: Union[Exception, BaseException, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        logging.info("[MaximSDK] on_chain_error called")

    def on_custom_event(
        self,
        name: str,
        data: Any,
        *,
        run_id: UUID,
        tags: Optional[list[str]] = None,
        metadata: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        container = self.__get_container(run_id, kwargs.get("parent_run_id", None))
        if container is None:
            logging.error("[MaximSDK] Couldn't find a container for event")
            return
        if tags is not None:
            final_tags = {
                tag.split(":")[0]: tag.split(":")[1] for tag in tags if ":" in tag
            }
            container.add_event(str(run_id), name, final_tags)

    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        logging.warning("[MaximSDK] On agent action is not supported")

    def on_agent_finish(
        self,
        finish: AgentFinish,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        **kwargs: Any,
    ) -> Any:
        logging.warning("[MaximSDK] On agent action is not supported")

    # Tool callbacks

    def on_tool_start(
        self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> Any:
        try:
            run_id = kwargs.get("run_id", None)
            parent_run_id = kwargs.get("parent_run_id", None)
            name = serialized.get("name", None)
            description = serialized.get("description", None)
            container = self.__get_container(run_id, parent_run_id)
            if container is None:
                logging.error("[MaximSDK] Couldn't find a container for this tool call")
                return
            if not container.is_created():
                container.create()
            container.add_tool_call(
                ToolCallConfig(
                    id=str(run_id), name=name, description=description, args=input_str
                )
            )
            if container.parent() is None:
                self.containers[str(run_id)] = container
        except Exception as e:
            logging.error("[MaximSDK] Failed to parse tool-start: %s", e)

    def on_tool_end(self, output: Any, **kwargs: Any) -> Any:
        try:
            run_id = kwargs.get("run_id", None)
            parent_run_id = kwargs.get("parent_run_id", None)
            container = self.__get_container(run_id, parent_run_id)
            if container is None:
                logging.error("[MaximSDK] Couldn't find a container for tool call")
                return
            if isinstance(output, ToolMessage):
                if output.status == "success":
                    self.logger.tool_call_result(str(run_id), output.content)
                elif output.status == "error":
                    if isinstance(output.content, str):
                        self.logger.tool_call_error(
                            str(run_id), ToolCallError(message=output.content)
                        )
                    else:
                        self.logger.tool_call_error(
                            str(run_id), ToolCallError(message=str(output.content))
                        )
            if container.parent() is None:
                container.end()
                self.containers.pop(str(run_id))
        except Exception as e:
            logging.error("[MaximSDK] Failed to parse tool-end: %s", e)

    def on_tool_error(
        self, error: Union[Exception, BaseException, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        try:
            run_id = kwargs.get("run_id", None)
            container = self.__get_container(run_id, kwargs.get("parent_run_id", None))
            if container is None:
                logging.error("[MaximSDK] Couldn't find a container for tool_call")
                return
            if error is not None and hasattr(error, "__str__"):
                self.logger.tool_call_error(
                    str(run_id), ToolCallError(message=str(error))
                )
            if container.parent() is None:
                container.end()
                self.containers.pop(str(run_id))
        except Exception as e:
            logging.error("[MaximSDK] Failed to parse tool-end: %s", e)

    def on_text(self, text: str, **kwargs: Any) -> Any:
        logging.warning("[MaximSDK] Text models are not supported by Maxim callback")
