import os
from typing import Any, Optional, Generic, Iterable, cast
from cuid2 import cuid_wrapper

from opentelemetry import propagate
from opentelemetry.propagators.textmap import CarrierT

from frisk_sdk.config.environment_variables import (
    FRISK_API_KEY,
    FRISK_BASE_URL,
    FRISK_BASE_URL_DEFAULT,
    FRISK_TELEMETRY_ENDPOINT,
    FRISK_TELEMETRY_ENDPOINT_DEFAULT,
    FRISK_SHUTDOWN_TIMEOUT_S,
    FRISK_SHUTDOWN_TIMEOUT_S_DEFAULT,
)
from .llm_reasoning import remove_llm_reasoning_arg
from .logging_options import LoggingOptions
from ..redact.redaction_options import normalize_redaction_options
from .wrap_tool_options import WrapToolOptions
from ..errors.configuration_errors import (
    MissingAPIKeyError,
    MissingBaseURLError,
    MissingOtlpEndpointError,
)
from ..logging import get_logger, auto_enable_debug_logging
from ..logging.frisk_logger_name import FRISK_LOGGER_NAME
from ..process import install_frisk_shutdown
from .._core import FriskHandle as FriskCoreHandle  # type: ignore
from .tool_call_decision import ToolCallDecision

from .access_token_provider_runtime import get_or_init_access_token_provider_runtime
from .frisk_session import FriskSession
from .session_registry import SessionRegistry
from frisk_sdk.framework_adapter.framework_adapter import (
    FrameworkAdapter,
    NormalizedToolCall,
    ToolCallT,
    ToolT,
)
from frisk_sdk.framework_adapter.base_framework_adapter import BaseFrameworkAdapter
from ..telemetry.constants import TRACER_NAME, TRACER_VERSION
from ..telemetry.tracing_manager import TracingManager

# todo: config. https://linear.app/friskai/issue/POL-55/add-config-library-to-sdk
TOKEN_REFRESH_INTERVAL_MINUTES = 5.0


class Frisk(Generic[ToolCallT, ToolT]):
    """Main Frisk SDK client for policy evaluation and tool instrumentation."""

    def __init__(
        self,
        api_key: Optional[str] = os.getenv(FRISK_API_KEY),
        *,
        api_base_url: Optional[str] = None,
        otlp_endpoint: Optional[str] = None,
        adapter: Optional[FrameworkAdapter[ToolCallT, ToolT]] = None,
        logging: Optional[LoggingOptions] = None,
        redact: object | None = None,
    ):
        logger_name = ".".join(
            filter(None, [FRISK_LOGGER_NAME, (logging or {}).get("logger_name")])
        )
        self.logger = (logging or {}).get("logger") or get_logger(logger_name)
        self.logger_name = self.logger.name
        auto_enable_debug_logging(self.logger.name)

        self.logger.debug("Initializing Frisk SDK...")

        self.redaction = normalize_redaction_options(redact)

        if not api_key:
            raise MissingAPIKeyError()

        if api_base_url is None:
            api_base_url = os.getenv(FRISK_BASE_URL, FRISK_BASE_URL_DEFAULT)
        if not api_base_url:
            raise MissingBaseURLError()

        self.logger.debug(f"Connecting to Frisk API at {api_base_url}")

        if otlp_endpoint is None:
            otlp_endpoint = os.getenv(
                FRISK_TELEMETRY_ENDPOINT, FRISK_TELEMETRY_ENDPOINT_DEFAULT
            )
        if not otlp_endpoint:
            raise MissingOtlpEndpointError()

        self.logger.debug(f"Sending telemetry to: {otlp_endpoint}")

        self._adapter = adapter or BaseFrameworkAdapter()

        # api_key is used to obtain short-lived access tokens.
        # NOTE: This uses a shared AccessTokenProviderRuntime so multiple Frisk() instances
        # do NOT spawn multiple background refresh threads.
        self._token_runtime = get_or_init_access_token_provider_runtime(
            api_key=api_key,
            base_url=api_base_url,
            refresh_interval_minutes=int(TOKEN_REFRESH_INTERVAL_MINUTES),
            parent_logger_name=self.logger_name,
        )
        self._cuid_generator = cuid_wrapper()

        # Initialize tracing with a current token (will fetch on demand)
        current_token = self._token_runtime.get_access_token()
        self._tracing_manager = TracingManager(
            otlp_endpoint=otlp_endpoint,
            auth_token=current_token,
            tracer_name=TRACER_NAME,
            tracer_version=TRACER_VERSION,
        )

        self._handle = FriskCoreHandle(current_token, api_base_url, otlp_endpoint)

        # Subscribe this instance to future token updates
        self._token_callbacks = [
            lambda token: self.update_access_token(token),
            lambda token: self._tracing_manager.update_auth_token(token),
        ]
        for cb in self._token_callbacks:
            self._token_runtime.subscribe(cb)

        shutdown_timeout_s = float(
            os.getenv(FRISK_SHUTDOWN_TIMEOUT_S, FRISK_SHUTDOWN_TIMEOUT_S_DEFAULT)
        )
        self.install_shutdown_handler(shutdown_timeout_s)

    def session(self) -> FriskSession:
        """Create and return a new FriskSession.

        Returns:
            FriskSession: The created session object. Use ``session.id`` as the
                session identifier in other APIs or when passing session context.
        """
        session = FriskSession(
            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

    def update_access_token(self, access_token: str):
        """Update the auth token used by the underlying Rust handle."""
        self._handle.update_auth_token(access_token)

    def normalize_tool_call(self, tool_call: ToolCallT) -> NormalizedToolCall:
        """Extract normalized tool-call metadata using the configured adapter."""
        if isinstance(tool_call, NormalizedToolCall):
            return tool_call

        return self._adapter.normalize_tool_call(tool_call)

    def decide_tool_call(
        self,
        *,
        tool_call: NormalizedToolCall,
        agent_state: Optional[dict[str, Any] | str],
        trace_context_carrier: Optional[CarrierT] = None,
    ) -> ToolCallDecision:
        """Evaluate a tool call against active policies and return a outcome.
        :param *:
        """

        # Use the active OTel context unless caller overrides
        if trace_context_carrier is None:
            trace_context_carrier = cast(CarrierT, {})
        propagate.inject(trace_context_carrier)

        base_tool_args = remove_llm_reasoning_arg(tool_call.args)
        serialized_tool_args = (
            base_tool_args
            if isinstance(base_tool_args, str)
            else self._adapter.serialize_tool_args(base_tool_args).value
        )

        serialized_agent_state = (
            agent_state
            if isinstance(agent_state, str)
            else self._adapter.serialize_agent_state(agent_state).value
        )

        result = self._handle.process(
            tool_call.name,
            serialized_tool_args,
            serialized_agent_state,
            tool_call.id,
            self.redaction,
            trace_context_carrier,
        )

        return ToolCallDecision.from_rust_binding(result)

    def wrap_tools(
        self,
        tools: Iterable[ToolT],
        options: WrapToolOptions | None = None,
    ) -> list[ToolT]:
        """Wrap framework tools through the configured adapter."""
        return cast(list[ToolT], self._adapter.wrap_tools(tools, options))

    def wrap_tool(self, tool: ToolT, options: WrapToolOptions | None = None) -> ToolT:
        return self._adapter.wrap_tool(tool, options)

    def shutdown(self, signal: Optional[float] = None) -> None:
        """Shut down SDK resources and unsubscribe token/tracing callbacks."""
        self.logger.info("Frisk is shutting down...")

        # Unsubscribe this Frisk instance from token updates.
        # (The shared AccessTokenProviderRuntime remains running for the process lifetime.)
        try:
            for cb in getattr(self, "_token_callbacks", []):
                self._token_runtime.unsubscribe(cb)
        except Exception:
            pass
        self._handle.shutdown()
        self._tracing_manager.shutdown()

    def install_shutdown_handler(self, shutdown_timeout_s: float):
        """Install process exit/signal hooks to flush and close SDK resources."""
        install_frisk_shutdown(
            lambda signal: self.shutdown(signal),  # shutdown(signal)
            timeout_s=shutdown_timeout_s,  # per instance
            install_signals=True,
            logger=self.logger,
        )
