from typing import Optional, TypedDict

from ...core import Frisk as FriskCore
from ...core.logging_options import LoggingOptions
from ...redact.redaction_options import RedactionOptionsInput
from .frisk_tool_hook import FriskStrandsToolHook
from .strands_framework_adapter import StrandsFrameworkAdapter


class FriskHookControls(TypedDict, total=False):
    """Optional controls for `FriskStrands.hook` / `tool_hook`.

    Supported keys:
      - tool_policies: Optional[bool]
          When True or omitted, tool invocations are subject to Frisk tool
          policy decisions. Setting this to False is currently not supported
          for this hook and will raise a ValueError.
    """

    tool_policies: Optional[bool]


class FriskStrands(FriskCore):
    """Frisk client preconfigured with the Strands 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=StrandsFrameworkAdapter(),
            logging=logging,
            redact=redact,
        )

    def tool_hook(self, controls: Optional[FriskHookControls] = None):
        """Create a Strands HookProvider that enforces Frisk policy decisions."""
        return self.hook(controls)

    def hook(self, controls: Optional[FriskHookControls] = None):
        """Create a Strands HookProvider that enforces Frisk policy decisions."""
        if (controls or {}).get("tool_policies") is False:
            raise ValueError(
                "Disabling tool policies is not yet implemented for this hook. If this functionality is not required, do not use this hook."
            )

        return FriskStrandsToolHook(self)
