from typing import Optional, TypedDict, cast

from .frisk_tool_middleware import FriskToolMiddleware
from .frisk_langchain_session import FriskLangchainSession
from .langchain_framework_adapter import LangchainFrameworkAdapter
from ... import SessionRegistry
from ...core import Frisk as FriskCore
from ...core.logging_options import LoggingOptions
from ...redact.redaction_options import RedactionOptionsInput


class FriskGuardControls(TypedDict, total=False):
    tool_policies: Optional[bool]


class FriskLangchain(FriskCore):
    """Frisk client preconfigured with the LangChain framework adapter."""

    def __init__(
        self,
        api_key: str,
        *,
        api_base_url: Optional[str] = None,
        logging: Optional[LoggingOptions] = None,
        redact: Optional[RedactionOptionsInput] = None,
    ):
        super().__init__(
            api_key,
            api_base_url=api_base_url,
            adapter=LangchainFrameworkAdapter(),
            logging=logging,
            redact=redact,
        )

    def guard(self, controls: FriskGuardControls | None = None):
        if (controls or {}).get("tool_policies") is False:
            raise ValueError(
                "Disabling tool policies is not yet implemented for this middleware. If this functionality is not required, do not use this middleware."
            )
        return FriskToolMiddleware(
            self, logging=cast(LoggingOptions, {"logger": self.logger})
        )

    def session(self):
        """Create and return a new :class:`FriskLangchainSession`.

        The returned session object exposes its identifier and context via
        the ``id`` and ``context`` attributes (e.g. ``session.id``,
        ``session.context``).

        Returns:
            FriskLangchainSession: The newly created session object.
        """
        session = FriskLangchainSession(
            self,
            tracer=self._tracing_manager.get_tracer(),
            redact=self.redaction,
        )

        # Register session in the global registry
        registry = SessionRegistry()
        registry.register(session.id, session)

        return session
