import os
from typing import Any, Optional, Generic, Iterable, cast
from uuid import UUID, uuid4
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,
    FRISK_SERVICE_NAME,
    FRISK_SERVICE_NAME_DEFAULT,
    FRISK_SERVICE_NAMESPACE,
    FRISK_SERVICE_NAMESPACE_DEFAULT,
)
from .frisk_options import normalize_frisk_options
from ..errors.configuration_errors import (
    MissingAPIKeyError,
    MissingBaseURLError,
    MissingTelemetryEndpointError,
)
from ..logging import get_logger, 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 .process_tool_call_result import ProcessToolCallResult

from frisk_sdk.telemetry.get_tracer import get_tracer
from .token_manager_runtime import get_or_init_token_runtime
from .frisk_session import FriskSession
from .session_registry import SessionRegistry
from frisk_sdk.framework_adapter.framework_adapter import (
    FrameworkAdapter,
    ToolCallInfo,
    ToolCallT,
    ToolT,
)
from frisk_sdk.framework_adapter.base_framework_adapter import BaseFrameworkAdapter
from ..telemetry.tracer_handle import TracerHandle
from ..utils import is_secure_url

# 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),
        *,
        base_url: Optional[str] = None,
        telemetry_endpoint: Optional[str] = None,
        framework_adapter: Optional[FrameworkAdapter[ToolCallT, ToolT]] = None,
        logger: Optional[Logger] = None,
        logger_name: Optional[str] = None,
        options: object | None = None,
    ):
        logger_name = ".".join(filter(None, [FRISK_LOGGER_NAME, logger_name]))
        self.logger = logger or get_logger(logger_name)
        self.logger_name = self.logger.name
        auto_enable_debug_logging(self.logger.name)

        self.options = normalize_frisk_options(options)

        if not api_key:
            raise MissingAPIKeyError()

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

        if telemetry_endpoint is None:
            telemetry_endpoint = os.getenv(
                FRISK_TELEMETRY_ENDPOINT, FRISK_TELEMETRY_ENDPOINT_DEFAULT
            )
        if not telemetry_endpoint:
            raise MissingTelemetryEndpointError()

        self._framework_adapter = framework_adapter or BaseFrameworkAdapter()
        self._tracer = get_tracer()  # todo: Standardize underscored properties.

        # api_key is used to obtain short-lived access tokens.
        # NOTE: This uses a shared TokenManagerRuntime so multiple Frisk() instances
        # do NOT spawn multiple background refresh threads.
        self._token_runtime = get_or_init_token_runtime(
            api_key=api_key,
            base_url=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.tracer_handle = TracerHandle.get_instance(
            current_token,
            endpoint=telemetry_endpoint,
            insecure=not is_secure_url(telemetry_endpoint),
            service_name=os.getenv(FRISK_SERVICE_NAME, FRISK_SERVICE_NAME_DEFAULT),
            service_namespace=os.getenv(
                FRISK_SERVICE_NAMESPACE, FRISK_SERVICE_NAMESPACE_DEFAULT
            ),
        ).acquire()

        self._handle = FriskCoreHandle(current_token, base_url, telemetry_endpoint)

        # Subscribe this instance to future token updates
        self._token_callbacks = [
            lambda token: self.update_access_token(token),
            lambda token: self.tracer_handle.update_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 create_session(self) -> UUID:
        """Create a new Frisk session and return its UUID.

        Returns:
            UUID: The session identifier to be used in context
        """
        session_id = uuid4()
        session = FriskSession(
            self,
            session_id=session_id,
            tracer=self._tracer,
            framework_adapter=self._framework_adapter,
            options=self.options,
        )

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

        return session_id

    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 get_tool_call_info(self, tool_call: ToolCallT) -> ToolCallInfo:
        """Extract normalized tool-call metadata using the configured adapter."""
        return self._framework_adapter.get_tool_call_info(tool_call)

    def evaluate_tool_call(
        self,
        tool_name: str,
        tool_args: Optional[dict[str, Any] | str],
        agent_state: Optional[dict[str, Any] | str],
        frisk_tool_call_id: Optional[str] = None,
        trace_context_carrier: Optional[CarrierT] = None,
    ) -> ProcessToolCallResult:
        """Evaluate a tool call against active policies and return a decision."""
        frisk_tool_call_id = frisk_tool_call_id or self._cuid_generator()

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

        serialized_tool_args = (
            tool_args
            if isinstance(tool_args, str)
            else self._framework_adapter.serialize_tool_args(tool_args).value
        )

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

        result = self._handle.process(
            tool_name,
            serialized_tool_args,
            serialized_agent_state,
            frisk_tool_call_id,
            self.options,
            trace_context_carrier,
        )

        return ProcessToolCallResult.from_rust_binding(result)

    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 TokenManagerRuntime 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()
        # Release our reference to the process-wide tracer handle.
        self.tracer_handle.release()

    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,
        )

    def wrap_tools(self, tools: Iterable[ToolT]) -> list[ToolT]:
        """Wrap framework tools through the configured adapter."""
        return cast(list[ToolT], self._framework_adapter.wrap_tools(tools))
