from typing import Any, Dict, Union, Optional, List
from uuid import uuid4
import datetime
import os

from langchain_core.callbacks.base import BaseCallbackHandler

from mindsdb.utilities import log
from mindsdb.interfaces.storage import db

logger = log.getLogger(__name__)
logger.setLevel('DEBUG')


class LangfuseCallbackHandler(BaseCallbackHandler):
    """Langchain callback handler that traces tool & chain executions using Langfuse."""

    def __init__(self, langfuse, trace_id: Optional[str] = None, observation_id: Optional[str] = None):
        self.langfuse = langfuse
        self.chain_uuid_to_span = {}
        self.action_uuid_to_span = {}
        # if these are not available, we generate some UUIDs
        self.trace_id = trace_id or uuid4().hex
        self.observation_id = observation_id or uuid4().hex

    def on_tool_start(
            self, serialized: Dict[str, Any], input_str: str, **kwargs: Any
    ) -> Any:
        """Run when tool starts running."""
        parent_run_uuid = kwargs.get('parent_run_id', uuid4()).hex
        action_span = self.action_uuid_to_span.get(parent_run_uuid)
        metadata = {
            'tool_name': serialized.get("name", "tool"),
            'started': datetime.datetime.now().isoformat()
        }
        action_span.update(metadata=metadata)

    def on_tool_end(self, output: str, **kwargs: Any) -> Any:
        """Run when tool ends running."""
        parent_run_uuid = kwargs.get('parent_run_id', uuid4()).hex
        action_span = self.action_uuid_to_span.get(parent_run_uuid)
        action_span.update(
            output=output,  # tool output is action output (unless superseded by a global action output)
            metadata={'finished': datetime.datetime.now().isoformat()}
        )

    def on_tool_error(
            self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any
    ) -> Any:
        """Run when tool errors."""
        parent_run_uuid = kwargs.get('parent_run_id', uuid4()).hex
        action_span = self.action_uuid_to_span.get(parent_run_uuid)
        try:
            error_str = str(error)
        except Exception:
            error_str = "Couldn't get error string."
        action_span.update(metadata={'error_description': error_str})

    def on_chain_start(
            self, serialized: Dict[str, Any], inputs: Dict[str, Any], **kwargs: Any
    ) -> Any:
        """Run when chain starts running."""
        run_uuid = kwargs.get('run_id', uuid4()).hex
        chain_span = self.langfuse.span(
            name=f'{serialized.get("name", "chain")}-{run_uuid}',
            trace_id=self.trace_id,
            parent_observation_id=self.observation_id,
            input=str(inputs)
        )
        self.chain_uuid_to_span[run_uuid] = chain_span

    def on_chain_end(self, outputs: Dict[str, Any], **kwargs: Any) -> Any:
        """Run when chain ends running."""
        chain_uuid = kwargs.get('run_id', uuid4()).hex
        if chain_uuid not in self.chain_uuid_to_span:
            return
        chain_span = self.chain_uuid_to_span.pop(chain_uuid)
        chain_span.update(output=str(outputs))
        chain_span.end()

    def on_chain_error(self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any) -> Any:
        """Run when chain errors."""
        # Do nothing for now.
        pass

    def on_agent_action(self, action, **kwargs: Any) -> Any:
        """Run on agent action."""
        # Do nothing for now.
        run_uuid = kwargs.get('run_id', uuid4()).hex
        action_span = self.langfuse.span(
            name=f'{getattr(action, "type", "action")}-{getattr(action, "tool", "")}-{run_uuid}',
            trace_id=self.trace_id,
            parent_observation_id=self.observation_id,
            input=str(action)
        )
        self.action_uuid_to_span[run_uuid] = action_span

    def on_agent_finish(self, finish, **kwargs: Any) -> Any:
        """Run on agent end."""
        # Do nothing for now.
        run_uuid = kwargs.get('run_id', uuid4()).hex
        if run_uuid not in self.action_uuid_to_span:
            return
        action_span = self.action_uuid_to_span.pop(run_uuid)
        if finish is not None:
            action_span.update(output=finish)  # supersedes tool output
        action_span.end()

    def auth_check(self):
        if self.langfuse is not None:
            return self.langfuse.auth_check()
        return False


def get_metadata(model_using: Dict) -> Dict:
    """ Generates initial metadata mapping from information provided in a model's `using` clause.
    Includes providers and model name.
    """
    metadata_keys = ['provider', 'model_name', 'embedding_model_provider']  # keeps keys relevant for tracing
    trace_metadata = {}
    for key in metadata_keys:
        if key in model_using:
            trace_metadata[key] = model_using.get(key)
    return trace_metadata


def get_skills(agent: db.Agents) -> List:
    """ Retrieve skills from agent `skills` attribute. Specific to agent endpoints. """
    return [s.type for s in agent.skills]


def get_tags(metadata: Dict) -> List:
    """ Retrieves tags from existing langfuse metadata (built using `get_metadata` and `get_skills`), and environment variables. """
    trace_tags = []
    if os.getenv('FLASK_ENV'):
        trace_tags.append(os.getenv('FLASK_ENV'))  # Fix: use something other than flask_env
    if 'provider' in metadata:
        trace_tags.append(metadata['provider'])
    return trace_tags


def get_tool_usage(trace) -> Dict:
    """ Retrieves tool usage information from a langfuse trace.
    Note: assumes trace marks an action with string `AgentAction` """
    tool_usage = {}
    steps = [s.name for s in trace.observations]
    for step in steps:
        if 'AgentAction' in step:
            tool_name = step.split('-')[1]
            if tool_name not in tool_usage:
                tool_usage[tool_name] = 0
            tool_usage[tool_name] += 1
    return tool_usage
