from __future__ import annotations

from typing import Optional, Dict, Any, List, TYPE_CHECKING
from uuid import UUID

from langchain_core.agents import AgentAction, AgentFinish
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

from frisk_sdk.logging import get_logger

if TYPE_CHECKING:
    from frisk_sdk.core.frisk import Frisk
    from .frisk_langchain_session import FriskLangchainSession


class FriskCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler that manages Frisk session tracing lifecycle."""

    def __init__(self, frisk: Frisk, session: FriskLangchainSession):
        """Initialize FriskCallbackHandler.

        Args:
            frisk: Frisk instance
            session_id: UUID or string representation of the session ID
        """
        self._frisk = frisk

        self._frisk_session = session

        self._logger = get_logger(f"{frisk.logger_name}.callback_handler")
        super().__init__()

    def on_chain_start(
        self,
        serialized: Optional[Dict[str, Any]],
        inputs: Dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Open Frisk root tracing span when a top-level chain run starts."""
        if parent_run_id is None:
            self._frisk_session.init_tracing(
                remote_session_id=str(run_id), inputs=inputs
            )

    def on_chain_end(
        self,
        outputs: dict[str, Any],
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> None:
        """Close Frisk root tracing span when the tracked chain run ends."""
        if self._frisk_session._root_run_id == str(run_id):
            self._frisk_session.end_tracing()

    def on_tool_start(
        self,
        serialized: dict[str, Any],
        input_str: str,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        metadata: Optional[dict[str, Any]] = None,
        inputs: Optional[dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        self._logger.debug(
            f"on_tool_start, run ID: {run_id}, parent run ID: {parent_run_id}"
        )

    def on_tool_end(
        self,
        output: Any,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._logger.debug(f"on_tool_end, run ID: {run_id}")

    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,
    ) -> None:
        self._logger.debug(
            f"on_llm_start, run ID: {run_id}, parent run ID: {parent_run_id}"
        )

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._logger.debug(f"on_llm_end, run ID: {run_id}")

    def on_agent_action(
        self,
        action: AgentAction,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._logger.debug(f"on_agent_action, type: {action.type}, tool: {action.tool}")

    def on_agent_finish(
        self,
        finish: AgentFinish,
        *,
        run_id: UUID,
        parent_run_id: Optional[UUID] = None,
        tags: Optional[list[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._logger.debug("on_agent_finish")
