"""
Middleware for managing Unity instance selection per session.

This middleware intercepts all tool calls and injects the active Unity instance
into the request-scoped state, allowing tools to access it via ctx.get_state("unity_instance").
"""
from threading import RLock
import logging
import time

from fastmcp.server.middleware import Middleware, MiddlewareContext

from core.config import config
from services.registry import get_registered_tools
from transport.plugin_hub import PluginHub

logger = logging.getLogger("mcp-for-unity-server")
# Separate logger that propagates to root -> stderr so diagnostics show in console
_diag = logging.getLogger("transport.unity_instance_middleware")

# Store a global reference to the middleware instance so tools can interact
# with it to set or clear the active unity instance.
_unity_instance_middleware = None
_middleware_lock = RLock()


def get_unity_instance_middleware() -> 'UnityInstanceMiddleware':
    """Get the global Unity instance middleware."""
    global _unity_instance_middleware
    if _unity_instance_middleware is None:
        with _middleware_lock:
            if _unity_instance_middleware is None:
                # Auto-initialize if not set (lazy singleton) to handle import order or test cases
                _unity_instance_middleware = UnityInstanceMiddleware()

    return _unity_instance_middleware


def set_unity_instance_middleware(middleware: 'UnityInstanceMiddleware') -> None:
    """Replace the global middleware instance.

    This is a test seam: production code uses ``get_unity_instance_middleware()``
    which lazy-initialises the singleton.  Tests call this function to inject a
    mock or pre-configured middleware before exercising tool/resource code.
    """
    global _unity_instance_middleware
    _unity_instance_middleware = middleware


class UnityInstanceMiddleware(Middleware):
    """
    Middleware that manages per-session Unity instance selection.

    Stores active instance per session_id and injects it into request state
    for all tool and resource calls.
    """

    # Key used in FastMCP's session-scoped state store for the active instance.
    _ACTIVE_INSTANCE_STATE_KEY = "mcpforunity.active_instance"

    def __init__(self):
        super().__init__()
        self._metadata_lock = RLock()
        self._unity_managed_tool_names: set[str] = set()
        self._tool_alias_to_unity_target: dict[str, str] = {}
        self._server_only_tool_names: set[str] = set()
        self._tool_visibility_signature: tuple[tuple[str, str], ...] = ()
        self._last_tool_visibility_refresh = 0.0
        self._tool_visibility_refresh_interval_seconds = 0.5
        self._has_logged_empty_registry_warning = False

    async def set_active_instance(self, ctx, instance_id: str) -> None:
        """Store the active instance for this MCP session.

        Persisted via FastMCP's session-scoped state store, which keys by
        ``ctx.session_id`` (the MCP-Session-Id header on HTTP, a per-subprocess
        UUID on stdio). Two MCP sessions cannot share state — see #1023 for the
        bug this replaces, which keyed on the peer-supplied ``client_id`` and
        collapsed multiple clients onto the same record.
        """
        await ctx.set_state(self._ACTIVE_INSTANCE_STATE_KEY, instance_id)

    async def get_active_instance(self, ctx) -> str | None:
        """Retrieve the active instance for this MCP session."""
        return await ctx.get_state(self._ACTIVE_INSTANCE_STATE_KEY)

    async def clear_active_instance(self, ctx) -> None:
        """Clear the stored instance for this MCP session.

        Overwrites with None rather than calling ``delete_state``: the read
        path already treats None as "no active instance", and this keeps the
        method usable from minimal context shims that don't implement
        ``delete_state``.
        """
        await ctx.set_state(self._ACTIVE_INSTANCE_STATE_KEY, None)

    async def _discover_instances(self, ctx) -> list:
        """
        Return running Unity instances across both HTTP (PluginHub) and stdio transports.

        Returns a list of objects with .id (Name@hash) and .hash attributes.
        """
        from types import SimpleNamespace
        transport = (config.transport_mode or "stdio").lower()
        results: list = []

        if PluginHub.is_configured():
            try:
                user_id = None
                get_state_fn = getattr(ctx, "get_state", None)
                if callable(get_state_fn) and config.http_remote_hosted:
                    user_id = await get_state_fn("user_id")
                sessions_data = await PluginHub.get_sessions(user_id=user_id)
                sessions = sessions_data.sessions or {}
                for session_info in sessions.values():
                    project = getattr(session_info, "project", None) or "Unknown"
                    hash_value = getattr(session_info, "hash", None)
                    if hash_value:
                        results.append(SimpleNamespace(
                            id=f"{project}@{hash_value}",
                            hash=hash_value,
                            name=project,
                        ))
            except Exception as exc:
                if isinstance(exc, (SystemExit, KeyboardInterrupt)):
                    raise
                logger.debug("PluginHub instance discovery failed (%s)", type(exc).__name__, exc_info=True)

        if not results and transport != "http":
            try:
                from transport.legacy.unity_connection import get_unity_connection_pool
                pool = get_unity_connection_pool()
                results = pool.discover_all_instances(force_refresh=True)
            except Exception as exc:
                if isinstance(exc, (SystemExit, KeyboardInterrupt)):
                    raise
                logger.debug("Stdio instance discovery failed (%s)", type(exc).__name__, exc_info=True)

        return results

    async def _resolve_instance_value(self, value: str, ctx) -> str:
        """
        Resolve a unity_instance string to a validated instance identifier.

        Accepts:
          - Bare port number like "6401" (stdio only) -> resolved Name@hash
          - "Name@hash" exact match
          - Hash prefix (unique prefix match against running instances)

        Raises ValueError with a user-friendly message on failure.
        """
        value = value.strip()
        if not value:
            raise ValueError("unity_instance value must not be empty.")

        transport = (config.transport_mode or "stdio").lower()

        # Port number (stdio only) — resolve to Name@hash via status file lookup
        if value.isdigit():
            if transport == "http":
                raise ValueError(
                    f"Port-based targeting ('{value}') is not supported in HTTP transport mode. "
                    "Use Name@hash or a hash prefix. Read mcpforunity://instances for available instances."
                )
            port_int = int(value)
            instances = await self._discover_instances(ctx)
            for inst in instances:
                if getattr(inst, "port", None) == port_int:
                    return inst.id
            available = ", ".join(
                f"{getattr(i, 'id', '?')} (port {getattr(i, 'port', '?')})"
                for i in instances
            ) or "none"
            raise ValueError(
                f"No Unity instance found on port {value}. Available: {available}."
            )

        instances = await self._discover_instances(ctx)
        ids = {
            getattr(inst, "id", None): inst
            for inst in instances
            if getattr(inst, "id", None)
        }

        # Exact Name@hash match
        if "@" in value:
            if value in ids:
                return value
            available = ", ".join(ids) or "none"
            raise ValueError(
                f"Instance '{value}' not found. Available: {available}. "
                "Read mcpforunity://instances for current sessions."
            )

        # Hash prefix match
        lookup = value.lower()
        matches = [
            inst for inst in instances
            if getattr(inst, "hash", "") and getattr(inst, "hash", "").lower().startswith(lookup)
        ]
        if len(matches) == 1:
            return matches[0].id
        if len(matches) > 1:
            ambiguous = ", ".join(getattr(m, "id", "?") for m in matches)
            raise ValueError(
                f"Hash prefix '{value}' is ambiguous ({ambiguous}). "
                "Provide the full Name@hash from mcpforunity://instances."
            )
        available = ", ".join(ids) or "none"
        raise ValueError(
            f"No running Unity instance matches '{value}'. Available: {available}. "
            "Read mcpforunity://instances for current sessions."
        )

    async def _maybe_autoselect_instance(self, ctx) -> str | None:
        """
        Auto-select the sole Unity instance when no active instance is set.

        Note: This method both *discovers* and *persists* the selection via
        `set_active_instance` as a side-effect, since callers expect the selection
        to stick for subsequent tool/resource calls in the same session.
        """
        try:
            transport = (config.transport_mode or "stdio").lower()
            # This implicit behavior works well for solo-users, but is dangerous for multi-user setups
            if transport == "http" and config.http_remote_hosted:
                return None
            if PluginHub.is_configured():
                try:
                    sessions_data = await PluginHub.get_sessions()
                    sessions = sessions_data.sessions or {}
                    ids: list[str] = []
                    for session_info in sessions.values():
                        project = getattr(
                            session_info, "project", None) or "Unknown"
                        hash_value = getattr(session_info, "hash", None)
                        if hash_value:
                            ids.append(f"{project}@{hash_value}")
                    if len(ids) == 1:
                        chosen = ids[0]
                        await self.set_active_instance(ctx, chosen)
                        logger.info(
                            "Auto-selected sole Unity instance via PluginHub: %s",
                            chosen,
                        )
                        return chosen
                    if len(ids) > 1:
                        logger.info(
                            "Multiple Unity instances found (%d). Pass unity_instance on any tool call "
                            "or call set_active_instance to choose one. Available: %s",
                            len(ids), ", ".join(ids),
                        )
                except (ConnectionError, ValueError, KeyError, TimeoutError, AttributeError) as exc:
                    logger.debug(
                        "PluginHub auto-select probe failed (%s); falling back to stdio",
                        type(exc).__name__,
                        exc_info=True,
                    )
                except Exception as exc:
                    if isinstance(exc, (SystemExit, KeyboardInterrupt)):
                        raise
                    logger.debug(
                        "PluginHub auto-select probe failed with unexpected error (%s); falling back to stdio",
                        type(exc).__name__,
                        exc_info=True,
                    )

            if transport != "http":
                try:
                    # Import here to avoid circular imports in legacy transport paths.
                    from transport.legacy.unity_connection import get_unity_connection_pool

                    pool = get_unity_connection_pool()
                    instances = pool.discover_all_instances(force_refresh=True)
                    ids = [getattr(inst, "id", None) for inst in instances]
                    ids = [inst_id for inst_id in ids if inst_id]
                    if len(ids) == 1:
                        chosen = ids[0]
                        await self.set_active_instance(ctx, chosen)
                        logger.info(
                            "Auto-selected sole Unity instance via stdio discovery: %s",
                            chosen,
                        )
                        return chosen
                    if len(ids) > 1:
                        logger.info(
                            "Multiple Unity instances found (%d). Pass unity_instance on any tool call "
                            "or call set_active_instance to choose one. Available: %s",
                            len(ids), ", ".join(ids),
                        )
                except (ConnectionError, ValueError, KeyError, TimeoutError, AttributeError) as exc:
                    logger.debug(
                        "Stdio auto-select probe failed (%s)",
                        type(exc).__name__,
                        exc_info=True,
                    )
                except Exception as exc:
                    if isinstance(exc, (SystemExit, KeyboardInterrupt)):
                        raise
                    logger.debug(
                        "Stdio auto-select probe failed with unexpected error (%s)",
                        type(exc).__name__,
                        exc_info=True,
                    )
        except Exception as exc:
            if isinstance(exc, (SystemExit, KeyboardInterrupt)):
                raise
            logger.debug(
                "Auto-select path encountered an unexpected error (%s)",
                type(exc).__name__,
                exc_info=True,
            )

        return None

    async def _resolve_user_id(self) -> str | None:
        """Extract user_id from the current HTTP request's API key."""
        if not config.http_remote_hosted:
            return None
        # Lazy import to avoid circular dependencies (same pattern as _maybe_autoselect_instance).
        from transport.unity_transport import _resolve_user_id_from_request
        return await _resolve_user_id_from_request()

    async def _inject_unity_instance(self, context: MiddlewareContext) -> None:
        """Inject active Unity instance and user_id into context if available."""
        ctx = context.fastmcp_context

        # Resolve user_id from the HTTP request's API key header
        user_id = await self._resolve_user_id()
        if config.http_remote_hosted and user_id is None:
            raise RuntimeError(
                "API key authentication required. Provide a valid X-API-Key header."
            )
        if user_id:
            await ctx.set_state("user_id", user_id)

        # Per-call routing: check if this tool call explicitly specifies unity_instance.
        # context.message.arguments is a mutable dict on CallToolRequestParams; resource
        # reads use ReadResourceRequestParams which has no .arguments, so this is a no-op for them.
        # We pop the key here so Pydantic's type_adapter.validate_python() never sees it.
        active_instance: str | None = None
        msg_args = getattr(getattr(context, "message", None), "arguments", None)
        if isinstance(msg_args, dict) and "unity_instance" in msg_args:
            raw = msg_args.pop("unity_instance")
            if raw is not None:
                raw_str = str(raw).strip()
                if raw_str:
                    # Raises ValueError with a user-friendly message on invalid input.
                    active_instance = await self._resolve_instance_value(raw_str, ctx)
                    logger.debug("Per-call unity_instance resolved to: %s", active_instance)

        if not active_instance:
            active_instance = await self.get_active_instance(ctx)
        if not active_instance:
            active_instance = await self._maybe_autoselect_instance(ctx)
        if active_instance:
            # If using HTTP transport (PluginHub configured), validate session
            # But for stdio transport (no PluginHub needed or maybe partially configured),
            # we should be careful not to clear instance just because PluginHub can't resolve it.
            # The 'active_instance' (Name@hash) might be valid for stdio even if PluginHub fails.

            session_id: str | None = None
            # Only validate via PluginHub if we are actually using HTTP transport.
            # For stdio transport, skip PluginHub entirely - we only need the instance ID.
            from transport.unity_transport import _is_http_transport
            if _is_http_transport() and PluginHub.is_configured():
                try:
                    # resolving session_id might fail if the plugin disconnected
                    # We only need session_id for HTTP transport routing.
                    # For stdio, we just need the instance ID.
                    # Pass user_id for remote-hosted mode session isolation
                    session_id = await PluginHub._resolve_session_id(active_instance, user_id=user_id)
                except (ConnectionError, ValueError, KeyError, TimeoutError) as exc:
                    # If resolution fails, it means the Unity instance is not reachable via HTTP/WS.
                    # If we are in stdio mode, this might still be fine if the user is just setting state?
                    # But usually if PluginHub is configured, we expect it to work.
                    # Let's LOG the error but NOT clear the instance immediately to avoid flickering,
                    # or at least debug why it's failing.
                    logger.debug(
                        "PluginHub session resolution failed for %s: %s; leaving active_instance unchanged",
                        active_instance,
                        exc,
                        exc_info=True,
                    )
                except Exception as exc:
                    # Re-raise unexpected system exceptions to avoid swallowing critical failures
                    if isinstance(exc, (SystemExit, KeyboardInterrupt)):
                        raise
                    logger.error(
                        "Unexpected error during PluginHub session resolution for %s: %s",
                        active_instance,
                        exc,
                        exc_info=True
                    )

            await ctx.set_state("unity_instance", active_instance)
            if session_id is not None:
                await ctx.set_state("unity_session_id", session_id)

    async def on_call_tool(self, context: MiddlewareContext, call_next):
        """Inject active Unity instance into tool context if available."""
        await self._inject_unity_instance(context)
        return await call_next(context)

    async def on_read_resource(self, context: MiddlewareContext, call_next):
        """Inject active Unity instance into resource context if available."""
        await self._inject_unity_instance(context)
        return await call_next(context)

    async def on_list_tools(self, context: MiddlewareContext, call_next):
        """Filter MCP tool listing to the Unity-enabled set when session data is available."""
        try:
            await self._inject_unity_instance(context)
        except Exception as exc:
            # Re-raise authentication errors so callers get a proper auth failure
            if isinstance(exc, RuntimeError) and "authentication" in str(exc).lower():
                raise
            _diag.warning(
                "on_list_tools: _inject_unity_instance failed (%s: %s), continuing without instance",
                type(exc).__name__, exc,
            )

        tools = await call_next(context)

        tool_names_from_fastmcp = sorted(getattr(t, "name", "?") for t in tools)
        _diag.debug(
            "on_list_tools: FastMCP returned %d tools: %s",
            len(tools), tool_names_from_fastmcp,
        )

        if not self._should_filter_tool_listing():
            _diag.debug("on_list_tools: skipping middleware filter (not HTTP or PluginHub not configured)")
            return tools

        self._refresh_tool_visibility_metadata_from_registry()
        enabled_tool_names = await self._resolve_enabled_tool_names_for_context(context)
        if enabled_tool_names is None:
            _diag.debug("on_list_tools: no Unity session data, returning %d tools from FastMCP as-is", len(tools))
            return tools

        filtered = []
        for tool in tools:
            tool_name = getattr(tool, "name", None)
            if self._is_tool_visible(tool_name, enabled_tool_names):
                filtered.append(tool)

        _diag.debug(
            "on_list_tools: filtered %d/%d tools visible (Unity register_tools). "
            "enabled_names=%s",
            len(filtered), len(tools), sorted(enabled_tool_names),
        )
        return filtered

    def _should_filter_tool_listing(self) -> bool:
        transport = (config.transport_mode or "stdio").lower()
        return transport == "http" and PluginHub.is_configured()

    async def _resolve_enabled_tool_names_for_context(
        self,
        context: MiddlewareContext,
    ) -> set[str] | None:
        ctx = context.fastmcp_context
        user_id = (await ctx.get_state("user_id")) if config.http_remote_hosted else None
        active_instance = await ctx.get_state("unity_instance")
        project_hashes = self._resolve_candidate_project_hashes(active_instance)
        try:
            sessions_data = await PluginHub.get_sessions(user_id=user_id)
            sessions = sessions_data.sessions if sessions_data else {}
        except Exception as exc:
            logger.debug(
                "Failed to fetch sessions for tool filtering (user_id=%s, %s)",
                user_id,
                type(exc).__name__,
                exc_info=True,
            )
            return None

        session_hashes = {
            getattr(session, "hash", None)
            for session in sessions.values()
            if getattr(session, "hash", None)
        }

        if project_hashes:
            active_hash = project_hashes[0]
            # Stale active_instance should not hide all Unity-managed tools.
            if active_hash not in session_hashes:
                return None
        else:
            if not sessions:
                return None

            if len(sessions) == 1:
                only_session = next(iter(sessions.values()))
                only_hash = getattr(only_session, "hash", None)
                if only_hash:
                    project_hashes = [only_hash]
            else:
                # Multiple sessions without explicit selection: use a union so we don't
                # hide tools that are valid in at least one visible Unity instance.
                project_hashes = [hash_value for hash_value in session_hashes if hash_value]

        if not project_hashes:
            return None

        enabled_tool_names: set[str] = set()
        resolved_any_project = False
        for project_hash in project_hashes:
            try:
                registered_tools = await PluginHub.get_tools_for_project(project_hash, user_id=user_id)
                # Only mark as resolved if tools are actually registered.
                # An empty list means register_tools hasn't been sent yet.
                if registered_tools:
                    resolved_any_project = True
            except Exception as exc:
                logger.debug(
                    "Failed to fetch tools for project hash %s (user_id=%s, %s)",
                    project_hash,
                    user_id,
                    type(exc).__name__,
                    exc_info=True,
                )
                continue

            for tool in registered_tools:
                tool_name = getattr(tool, "name", None)
                if isinstance(tool_name, str) and tool_name:
                    enabled_tool_names.add(tool_name)

        if not resolved_any_project:
            return None

        return enabled_tool_names

    def _refresh_tool_visibility_metadata_from_registry(self) -> None:
        now = time.monotonic()
        if now - self._last_tool_visibility_refresh < self._tool_visibility_refresh_interval_seconds:
            return

        with self._metadata_lock:
            now = time.monotonic()
            if now - self._last_tool_visibility_refresh < self._tool_visibility_refresh_interval_seconds:
                return

            try:
                registry_tools = get_registered_tools()
            except Exception:
                logger.warning(
                    "Failed to refresh tool visibility metadata from registry; keeping previous metadata.",
                    exc_info=True,
                )
                self._last_tool_visibility_refresh = now
                return

            if not registry_tools and not self._has_logged_empty_registry_warning:
                logger.warning(
                    "Tool registry is empty during tool-list filtering; treating tools as unknown/visible."
                )
                self._has_logged_empty_registry_warning = True
            elif registry_tools:
                self._has_logged_empty_registry_warning = False

            unity_managed_tool_names: set[str] = set()
            tool_alias_to_unity_target: dict[str, str] = {}
            server_only_tool_names: set[str] = set()
            signature_entries: list[tuple[str, str]] = []

            for tool_info in registry_tools:
                tool_name = tool_info.get("name")
                if not isinstance(tool_name, str) or not tool_name:
                    continue

                unity_target = tool_info.get("unity_target", tool_name)
                if unity_target is None:
                    server_only_tool_names.add(tool_name)
                    signature_entries.append((tool_name, "<server-only>"))
                    continue

                if not isinstance(unity_target, str) or not unity_target:
                    logger.debug(
                        "Skipping tool visibility metadata with invalid unity_target: %s",
                        tool_info,
                    )
                    continue

                if unity_target == tool_name:
                    unity_managed_tool_names.add(tool_name)
                    signature_entries.append((tool_name, unity_target))
                    continue

                tool_alias_to_unity_target[tool_name] = unity_target
                unity_managed_tool_names.add(unity_target)
                signature_entries.append((tool_name, unity_target))

            signature = tuple(sorted(signature_entries, key=lambda item: item[0]))
            if signature == self._tool_visibility_signature:
                self._last_tool_visibility_refresh = now
                return

            self._unity_managed_tool_names = unity_managed_tool_names
            self._tool_alias_to_unity_target = tool_alias_to_unity_target
            self._server_only_tool_names = server_only_tool_names
            self._tool_visibility_signature = signature
            self._last_tool_visibility_refresh = now

    @staticmethod
    def _resolve_candidate_project_hashes(active_instance: str | None) -> list[str]:
        if not active_instance:
            return []

        if "@" in active_instance:
            _, _, suffix = active_instance.rpartition("@")
            return [suffix] if suffix else []

        return [active_instance]

    def _is_tool_visible(self, tool_name: str | None, enabled_tool_names: set[str]) -> bool:
        if not isinstance(tool_name, str) or not tool_name:
            return True

        if tool_name in self._server_only_tool_names:
            return True

        if tool_name in enabled_tool_names:
            return True

        unity_target = self._tool_alias_to_unity_target.get(tool_name)
        if unity_target:
            return unity_target in enabled_tool_names

        # Keep unknown tools visible for forward compatibility.
        if tool_name not in self._unity_managed_tool_names:
            return True

        return False
