# Copyright (C) 2026 GoodData Corporation

"""Workspace root resolution utilities for MCP server."""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import unquote, urlparse

if TYPE_CHECKING:
    from typing import Any

    from mcp.server.fastmcp import Context as MCPContext

    Context = MCPContext[Any, Any, Any]

logger = logging.getLogger(__name__)

# Workspace root resolution (allows overrides for different repositories)
ENV_WORKSPACE_ROOT = "AIDA_WORKSPACE_ROOT"

# Global state for workspace root override and request/session cache
_workspace_root_override: Path | None = None
WORKSPACE_CACHE: dict[str, Path] = {}


def _fallback_workspace_root() -> Path:
    """Fallback workspace root when no override/roots are available."""
    return Path.cwd().resolve()


def _determine_workspace_override() -> Path | None:
    """Determine workspace root override from environment variable."""
    override = os.environ.get(ENV_WORKSPACE_ROOT)
    if not override:
        logger.info(
            "No %s override provided; will rely on client-advertised workspace roots.",
            ENV_WORKSPACE_ROOT,
        )
        return None

    candidate = Path(override).expanduser()
    if not candidate.is_absolute():
        candidate = (Path.cwd() / candidate).resolve()

    if candidate.exists():
        logger.info("Using workspace root override from %s: %s", ENV_WORKSPACE_ROOT, candidate)
        return candidate

    logger.warning(
        "Workspace root override ignored; path does not exist (%s): %s",
        ENV_WORKSPACE_ROOT,
        candidate,
    )
    return None


def path_from_uri(uri: str) -> Path:
    """Convert a file:// URI to a Path object."""
    parsed = urlparse(uri)
    if parsed.scheme != "file":
        raise ValueError(f"Unsupported URI scheme for workspace root: {uri}")

    netloc = parsed.netloc or ""
    path = unquote(parsed.path or "")
    netloc_lower = netloc.lower()

    if os.name == "nt":
        # Windows file URIs can encode drive letters in path and UNC shares in netloc.
        candidate = Path(f"//{netloc}{path}") if netloc and netloc_lower != "localhost" else Path(path)
    else:
        # For local file URIs, host is empty or "localhost".
        if netloc and netloc_lower != "localhost":
            raise ValueError(f"Unsupported non-local file URI host for workspace root: {uri}")
        candidate = Path(path)

    return candidate.resolve()


async def get_workspace_root(ctx: Context | None) -> Path:
    """
    Determine the effective workspace root for this request.

    Priority:
    1. Environment variable override (AIDA_WORKSPACE_ROOT)
    2. Client-advertised workspace roots (cached per transport/client session)
    3. Current working directory
    """
    global _workspace_root_override  # noqa: PLW0603

    # Initialize override on first call
    if _workspace_root_override is None:
        _workspace_root_override = _determine_workspace_override()

    if _workspace_root_override:
        os.environ.setdefault(ENV_WORKSPACE_ROOT, str(_workspace_root_override))
        return _workspace_root_override

    if ctx:
        cache_key = _workspace_cache_key(ctx)
        if cache_key and cache_key in WORKSPACE_CACHE:
            root = WORKSPACE_CACHE[cache_key]
            os.environ.setdefault(ENV_WORKSPACE_ROOT, str(root))
            return root

        # Try to list workspace roots from the client
        # This may fail if:
        # - Context is not available (outside of request)
        # - Client doesn't support roots capability
        # - Session is not properly initialized
        try:
            session = ctx.session
            roots_result = await session.list_roots()
        except (ValueError, AttributeError) as exc:
            # ValueError: "Context is not available outside of a request"
            # AttributeError: session or list_roots not available
            logger.debug("Cannot access session for workspace roots: %s", exc)
        except Exception as exc:  # noqa: BLE001
            # Other errors (e.g., client doesn't support roots capability)
            logger.warning("Unable to list client workspace roots: %s", exc)
        else:
            for root in roots_result.roots:
                try:
                    root_path = path_from_uri(str(root.uri))
                except ValueError:
                    continue
                if not root_path.exists():
                    continue
                if cache_key:
                    WORKSPACE_CACHE[cache_key] = root_path
                logger.info(
                    "Detected workspace root from client (cache_key=%s, name=%s): %s",
                    cache_key,
                    getattr(root, "name", None),
                    root_path,
                )
                os.environ.setdefault(ENV_WORKSPACE_ROOT, str(root_path))
                return root_path
            logger.debug(
                "Client did not advertise usable workspace roots; falling back to default: %s",
                _fallback_workspace_root(),
            )

    fallback = _fallback_workspace_root()
    os.environ.setdefault(ENV_WORKSPACE_ROOT, str(fallback))
    return fallback


def _workspace_cache_key(ctx: Context) -> str | None:
    """Build a stable cache key for workspace roots from transport/session hints."""
    request_context = getattr(ctx, "request_context", None)
    request = getattr(request_context, "request", None) if request_context is not None else None
    headers = getattr(request, "headers", None)
    if headers is not None:
        try:
            mcp_session_id = headers.get("mcp-session-id") or headers.get("Mcp-Session-Id")
        except Exception:  # noqa: BLE001
            mcp_session_id = None
        if isinstance(mcp_session_id, str) and mcp_session_id.strip():
            return f"mcp:{mcp_session_id.strip()}"

    try:
        client_id = ctx.client_id
    except (ValueError, AttributeError):
        client_id = None
    if isinstance(client_id, str) and client_id.strip():
        return f"client:{client_id.strip()}"

    return None
