"""A360 MCP server: tool exposure, schema compaction, and the call pipeline.

The server is intentionally lean. Unlike the sibling ServiceNow MCP it has:

* **No multi-instance routing.** A360 has one Control Room per server process.
* **No skills/resources layer.** Phase 4 scope is tool plumbing only.
* **No god "standard" package.** Packages are layered ``core`` →
  ``bot_read``/``log_read`` → ``diagnose`` → ``full`` from day one, because a
  bloated tool list taxes every request (the LLM re-reads it each turn).

Two responsibilities dominate this module:

1. **Schema compaction.** Pydantic JSON schemas carry a lot of filler
   (``title`` on every node, ``anyOf`` nullable wrappers, model docstrings,
   long default strings). The filler is noise to the LLM and bloats the tool
   advertisement. We strip it WITHOUT ever truncating: a truncated default or
   parameter description is worse than the original, because it looks valid
   but is wrong.
2. **The call_tool pipeline.** Order matters and is documented inline. The
   contract is: return a structured dict on every outcome, including errors.
   Raising across the MCP boundary makes the client show a generic error;
   a structured dict lets the LLM see the failure and decide what to do next.

Write classification and policy guards live in
``policies/write_guards.py`` (the single source of truth). This module
imports from there — the identity-pin tests assert both modules expose the
same function objects, so the classification cannot drift.
"""

from __future__ import annotations

import logging
import os
from typing import Any

import orjson
from mcp.server import NotificationOptions, Server
from mcp.server.stdio import stdio_server
from mcp.types import TextContent, Tool
from pydantic import BaseModel, ValidationError

from .auth.auth_manager import AuthManager
from .policies.write_guards import (
    ENVIRONMENT_CONFIRM_FIELD,
    is_mutating_tool,
    is_read_only_call,
    run_post_confirm_guards,
    run_pre_confirm_guards,
)
from .utils.config import ServerConfig
from .utils.registry import ToolEntry, discover_tools_lazy
from .utils.response_budget import enforce_response_budget

logger = logging.getLogger(__name__)

# --------------------------------------------------------------------------- #
# Constants
# --------------------------------------------------------------------------- #

# String defaults longer than this are dropped from the advertised schema. We
# never TRUNCATE a default: a half default is a corrupt default. If it does
# not fit, the LLM does not need to see it pre-filled.
_MAX_DEFAULT_STR = 60

# Field names that are self-explanatory to an LLM; their description is filler
# and gets stripped. Add sparingly — only when the name alone tells the whole
# story.
_SELF_EXPLANATORY_FIELDS = frozenset(
    {
        "dry_run",
        "limit",
        "offset",
        "page",
        "page_size",
        "fields",
        "sort",
        "filter",
        "verbose",
    }
)

# Package layering. ``None`` (or "*") means "all tools". Sets are explicit
# allow-sets so adding a new tool does not silently leak into a smaller
# package — the author must opt it in by editing this map. Packages are
# cumulative: each tier includes the tools of the tiers below it.
_CORE_TOOLS = frozenset({"get_execution", "list_executions", "get_bot_actions"})
_DIAGNOSE_TOOLS = frozenset({"diagnose_execution"})
_BROWSER_READ_TOOLS = frozenset(
    {"inspect_console", "inspect_network", "screenshot_debug", "get_computed_style"}
)
_BROWSER_WRITE_TOOLS = frozenset({"open_debug_window", "close_debug_window"})

_PACKAGE_DEFINITIONS: dict[str, set[str] | None] = {
    "core": set(_CORE_TOOLS),
    "bot_read": set(_CORE_TOOLS),
    "log_read": set(_CORE_TOOLS),
    "diagnose": set(_CORE_TOOLS) | set(_DIAGNOSE_TOOLS),
    # Browser package: read-only inspection only. Window lifecycle tools
    # (open/close) are NOT included here — opt into them with the ``browser``
    # package, which adds both halves.
    "browser_read": set(_CORE_TOOLS) | set(_DIAGNOSE_TOOLS) | set(_BROWSER_READ_TOOLS),
    "browser": set(_CORE_TOOLS)
    | set(_DIAGNOSE_TOOLS)
    | set(_BROWSER_READ_TOOLS)
    | set(_BROWSER_WRITE_TOOLS),
    "full": None,  # every registered tool
}

ENV_TOOL_PACKAGE = "A360_TOOL_PACKAGE"
DEFAULT_TOOL_PACKAGE = "core"

# Aliases let us rename a tool without breaking existing prompts. Empty for
# now; Phase 6 will populate as tools land.
_DEPRECATED_TOOL_ALIASES: dict[str, str] = {}

# Confirm-gate field name. Tools that mutate set ``approve=True`` in their
# arguments to acknowledge the destructive action; absence or False blocks.
CONFIRM_FIELD = "approve"

# Meta fields consumed by the guard/confirm layer; stripped before Pydantic
# validation because they are not part of any tool's params model.
_STRIP_FIELDS: frozenset[str] = frozenset({CONFIRM_FIELD, ENVIRONMENT_CONFIRM_FIELD})

# --------------------------------------------------------------------------- #
# Schema compaction
# --------------------------------------------------------------------------- #


def _is_self_explanatory_field(name: str, schema: dict[str, Any]) -> bool:
    """True when a field's name carries enough meaning that its description is filler."""
    if name in _SELF_EXPLANATORY_FIELDS:
        return True
    # ``include_*`` / ``*_only`` boolean toggles: the name is the contract.
    if schema.get("type") == "boolean":
        if name.startswith("include_") or name.endswith("_only"):
            return True
    # Bounded integer knobs: ``max_*`` / ``min_*`` carry their semantics in name + bounds.
    if schema.get("type") == "integer" and (name.startswith("max_") or name.startswith("min_")):
        return True
    return False


def _drop_trivial_default(value: Any) -> bool:
    """True when a default adds no information worth advertising."""
    if value is None or value is False or value == 0 or value == "" or value == [] or value == {}:
        return True
    if isinstance(value, str) and len(value) > _MAX_DEFAULT_STR:
        return True
    return False


def _strip_field_filler(name: str, schema: dict[str, Any]) -> dict[str, Any]:
    """Drop filler from a single property schema; recurse into nested objects."""
    out = dict(schema)
    if _is_self_explanatory_field(name, schema):
        out.pop("description", None)
    if "default" in out and _drop_trivial_default(out["default"]):
        out.pop("default")
    # Flatten anyOf nullable: {"anyOf":[{"type":X}, {"type":"null"}]} -> {"type":X}
    out = _flatten_anyof_nullable(out)
    # Recurse into nested object properties.
    props = out.get("properties")
    if isinstance(props, dict):
        out["properties"] = {
            k: _strip_field_filler(k, v) if isinstance(v, dict) else v for k, v in props.items()
        }
    # Recurse into array items.
    items = out.get("items")
    if isinstance(items, dict):
        out["items"] = _compact_schema(items)
    return out


def _flatten_anyof_nullable(schema: dict[str, Any]) -> dict[str, Any]:
    """Collapse ``anyOf:[T, {"type":"null"}]`` (or allOf) into just ``T``."""
    for key in ("anyOf", "oneOf"):
        node = schema.get(key)
        if not isinstance(node, list) or len(node) != 2:
            continue
        non_null = [n for n in node if not (isinstance(n, dict) and n.get("type") == "null")]
        if len(non_null) == 1 and isinstance(non_null[0], dict):
            merged = {k: v for k, v in schema.items() if k != key}
            merged.update(non_null[0])
            return merged
    return schema


def _compact_schema(schema: dict[str, Any], *, _top_level: bool = False) -> dict[str, Any]:
    """Strip filler from a JSON schema without ever truncating.

    Operations: drop ``title``; drop empty top-level ``required``; drop
    ``additionalProperties=False`` (it is the Pydantic default and adds noise);
    flatten nullable anyOf; recurse into properties and items.
    """
    if not isinstance(schema, dict):
        return schema

    out = _flatten_anyof_nullable(schema)
    out.pop("title", None)

    if _top_level:
        # The model docstring surfaces as top-level description; the tool's own
        # description (passed to register_tool) is what the LLM should see, so
        # drop the redundant model docstring here.
        out.pop("description", None)
        required = out.get("required")
        if isinstance(required, list) and len(required) == 0:
            out.pop("required")
        # additionalProperties=False is Pydantic's default; advertising it is noise.
        if out.get("additionalProperties") is False:
            out.pop("additionalProperties")

    if "default" in out and _drop_trivial_default(out["default"]):
        out.pop("default")

    props = out.get("properties")
    if isinstance(props, dict):
        out["properties"] = {
            k: _strip_field_filler(k, v) if isinstance(v, dict) else v for k, v in props.items()
        }

    items = out.get("items")
    if isinstance(items, dict):
        out["items"] = _compact_schema(items)

    defs = out.get("$defs") or out.get("definitions")
    if isinstance(defs, dict):
        # Strip descriptions from sub-model definitions but keep their structure;
        # leaf field descriptions are preserved via the recursion above.
        compact_defs = {}
        for def_name, def_schema in defs.items():
            if isinstance(def_schema, dict):
                d = dict(def_schema)
                d.pop("description", None)
                d.pop("title", None)
                compact_defs[def_name] = d
            else:
                compact_defs[def_name] = def_schema
        if "$defs" in out:
            out["$defs"] = compact_defs
        if "definitions" in out:
            out["definitions"] = compact_defs

    return out


# Per-(model, detail) schema cache. Schemas are computed once per process.
_TOOL_SCHEMA_CACHE: dict[tuple[type[BaseModel], bool], dict[str, Any]] = {}


def get_tool_schema(params_model: type[BaseModel], *, detail: bool = False) -> dict[str, Any]:
    """Return the compacted input schema for a tool's params model."""
    cache_key = (params_model, detail)
    cached = _TOOL_SCHEMA_CACHE.get(cache_key)
    if cached is not None:
        return cached

    raw = params_model.model_json_schema()
    compacted = _compact_schema(raw, _top_level=True)
    if not detail:
        # Non-detail mode drops $defs/definitions entirely: the LLM does not
        # need them to fill in flat kwargs, and they can be large.
        compacted.pop("$defs", None)
        compacted.pop("definitions", None)
    _TOOL_SCHEMA_CACHE[cache_key] = compacted
    return compacted


# --------------------------------------------------------------------------- #
# Serialization + budget
# --------------------------------------------------------------------------- #


def _compact_json(obj: Any) -> str:
    """Serialize with orjson (compact, UTF-8, no ASCII escaping)."""
    return orjson.dumps(obj).decode("utf-8")


def _compact_with_budget(obj: Any, *, tool_name: str) -> str:
    """Apply the response budget to ``obj`` and return compact JSON text."""
    bounded, _ = enforce_response_budget(obj, tool_name=tool_name)
    return _compact_json(bounded)


def serialize_tool_output(result: Any, *, tool_name: str) -> str:
    """Serialize a tool result to compact JSON text within the response budget.

    Strings are passed through unchanged (a tool may already return compact
    JSON). Pydantic models are dumped in JSON mode first. Everything else goes
    through ``_compact_with_budget``.
    """
    if isinstance(result, str):
        # Fast path: assume the tool already produced compact JSON. If it is
        # over budget, the tool chose to; we do not second-guess it.
        return result
    if isinstance(result, BaseModel):
        result = result.model_dump(mode="json")
    return _compact_with_budget(result, tool_name=tool_name)


# --------------------------------------------------------------------------- #
# Error result builder
# --------------------------------------------------------------------------- #


def _error_result(
    *,
    tool: str,
    error: str,
    error_type: str = "tool_error",
    action_required: str | None = None,
) -> dict[str, Any]:
    """Structured dict returned for every failure. Never raised across the wire."""
    out: dict[str, Any] = {
        "success": False,
        "error": error,
        "tool": tool,
        "error_type": error_type,
    }
    if action_required:
        out["action_required"] = action_required
    return out


# --------------------------------------------------------------------------- #
# Server
# --------------------------------------------------------------------------- #


class A360MCPServer:
    """Lean A360 MCP server.

    Responsibilities:
    * Resolve the active package (env-driven) and the resulting enabled tool set.
    * Advertise compacted tool schemas via ``list_tools``.
    * Run the ``call_tool`` pipeline (alias → enabled → unknown → mutating →
      confirm → pydantic → execute → structured dict → serialize+budget).
    * Own the ``AuthManager`` (created lazily so construction never blocks).
    """

    def __init__(
        self,
        config: ServerConfig,
        *,
        package: str | None = None,
        auth_manager: AuthManager | None = None,
    ) -> None:
        self.config = config
        self.package = self._resolve_package(package)
        # No owner override: AuthManager's default is "a360-mcp:<pid>", and the
        # per-process tag is what keeps two servers on one account counted as
        # two session slots instead of collapsing into one.
        self.auth_manager = auth_manager or AuthManager(config)
        self.enabled_names: set[str] | None = self._determine_enabled_tools(self.package)
        self.mcp_server = Server("a360-mcp")
        self._register_handlers()

    # ------------------------------------------------------------------ #
    # Configuration
    # ------------------------------------------------------------------ #

    @staticmethod
    def _resolve_package(package: str | None) -> str:
        if package:
            return package
        env_pkg = os.environ.get(ENV_TOOL_PACKAGE, "").strip()
        if env_pkg and env_pkg in _PACKAGE_DEFINITIONS:
            return env_pkg
        return DEFAULT_TOOL_PACKAGE

    @staticmethod
    def _determine_enabled_tools(package: str) -> set[str] | None:
        """Resolve the package allow-set. ``None`` means "all tools"."""
        definition = _PACKAGE_DEFINITIONS.get(package)
        if definition is None:
            return None
        # For packages with an explicit empty set, fall back to "all tools
        # registered so far" so a freshly-seeded package still exposes the
        # tools that exist at runtime. The author is expected to pin the set
        # once the tool roster stabilises.
        if len(definition) == 0:
            return None
        return set(definition)

    # ------------------------------------------------------------------ #
    # Handler registration
    # ------------------------------------------------------------------ #

    def _register_handlers(self) -> None:
        server = self

        @self.mcp_server.list_tools()
        async def _list_tools() -> list[Tool]:
            return server._list_tools_impl()

        @self.mcp_server.call_tool(validate_input=False)
        async def _call_tool(name: str, arguments: dict[str, Any]) -> list[TextContent]:
            return server._call_tool_impl(name, arguments)

    # ------------------------------------------------------------------ #
    # list_tools
    # ------------------------------------------------------------------ #

    def _list_tools_impl(self) -> list[Tool]:
        registry = discover_tools_lazy(self.enabled_names)
        tools: list[Tool] = []
        for name, entry in sorted(registry.items()):
            tools.append(self._build_tool_advertisement(name, entry))
        return tools

    def _build_tool_advertisement(self, name: str, entry: ToolEntry) -> Tool:
        _func, params_model, _return_type, description, _serialization = entry
        schema = get_tool_schema(params_model)
        return Tool(name=name, description=description, inputSchema=schema)

    # ------------------------------------------------------------------ #
    # call_tool pipeline
    # ------------------------------------------------------------------ #

    def _call_tool_impl(self, name: str, arguments: dict[str, Any]) -> list[TextContent]:
        # 1. Alias routing.
        canonical = _DEPRECATED_TOOL_ALIASES.get(name, name)

        # 2. Enabled-set check (before "unknown" so we can give a helpful
        #    "available in package X" message if the tool exists but is gated).
        registry = discover_tools_lazy(self.enabled_names)
        if canonical not in registry:
            if self.enabled_names is None:
                hint = ""
            else:
                hint = f" (not in the active '{self.package}' package)"
            # Is it known at all (in some other package)?
            full_registry = discover_tools_lazy(None)
            if canonical in full_registry:
                result = _error_result(
                    tool=name,
                    error=f"Tool '{name}' is registered but not enabled{hint}.",
                    error_type="not_enabled",
                    action_required=(
                        f"Switch the active package (env {ENV_TOOL_PACKAGE}) or pick a tool "
                        f"from the advertised list."
                    ),
                )
            else:
                result = _error_result(
                    tool=name,
                    error=f"Unknown tool: {name}",
                    error_type="unknown_tool",
                    action_required="Call list_tools to see advertised tools.",
                )
            return [TextContent(type="text", text=serialize_tool_output(result, tool_name=name))]

        entry = registry[canonical]
        func, params_model, _return_type, _description, _serialization = entry

        # 3. Read-only / mutating classification.
        if is_mutating_tool(canonical):
            if not self.config.allow_writes:
                result = _error_result(
                    tool=name,
                    error=(
                        f"Tool '{name}' is mutating but this server is read-only "
                        f"(ServerConfig.allow_writes=False)."
                    ),
                    error_type="write_blocked",
                    action_required="Enable writes in config to use mutating tools.",
                )
                return [
                    TextContent(type="text", text=serialize_tool_output(result, tool_name=name))
                ]

        # 4. Pre-confirm guards (local-only, 0 network).
        guard_error = self._run_pre_confirm_guards(canonical, arguments)
        if guard_error is not None:
            result = _error_result(
                tool=name,
                error=guard_error,
                error_type="write_guard_blocked",
            )
            return [TextContent(type="text", text=serialize_tool_output(result, tool_name=name))]

        # 5. Confirm gate for mutating tools.
        if is_mutating_tool(canonical):
            approved = bool(arguments.get(CONFIRM_FIELD, False))
            if not approved:
                result = _error_result(
                    tool=name,
                    error=(
                        f"Tool '{name}' mutates Control Room state. Set "
                        f"'{CONFIRM_FIELD}=true' to confirm."
                    ),
                    error_type="confirmation_required",
                    action_required=f"Re-call with {CONFIRM_FIELD}=true once reviewed.",
                )
                return [
                    TextContent(type="text", text=serialize_tool_output(result, tool_name=name))
                ]

        # 6. Post-confirm guards (live reads allowed).
        guard_error = self._run_post_confirm_guards(canonical, arguments)
        if guard_error is not None:
            result = _error_result(
                tool=name,
                error=guard_error,
                error_type="write_guard_blocked",
            )
            return [TextContent(type="text", text=serialize_tool_output(result, tool_name=name))]

        # Strip guard/confirm meta fields before validation; they are not part
        # of the tool's params model.
        clean_args = {k: v for k, v in arguments.items() if k not in _STRIP_FIELDS}

        # 7. Pydantic validation.
        try:
            params = params_model(**clean_args)
        except ValidationError as exc:
            result = _error_result(
                tool=name,
                error=f"Invalid arguments: {exc.errors()}",
                error_type="validation_error",
                action_required="Fix the argument errors and retry.",
            )
            return [TextContent(type="text", text=serialize_tool_output(result, tool_name=name))]

        # 8. Execute. A360AuthError is propagated distinctly (Phase 1 contract:
        #    never silence auth failures). Everything else becomes a structured
        #    error dict so the LLM can react.
        try:
            raw_result = func(self.config, self.auth_manager, params)
        except Exception as exc:  # noqa: BLE001 - boundary converts to structured dict
            # A360AuthError is re-raised so callers/tests can distinguish auth
            # failure from a tool error; the MCP stdio loop will surface it.
            from .auth.auth_manager import A360AuthError

            if isinstance(exc, A360AuthError):
                raise
            result = _error_result(
                tool=name,
                error=f"{type(exc).__name__}: {exc}",
                error_type="tool_error",
            )
            return [TextContent(type="text", text=serialize_tool_output(result, tool_name=name))]

        # 9. Serialize + budget.
        text = serialize_tool_output(raw_result, tool_name=name)
        return [TextContent(type="text", text=text)]

    # ------------------------------------------------------------------ #
    # Guard hooks (delegated to policies/write_guards.py)
    # ------------------------------------------------------------------ #

    def _run_pre_confirm_guards(self, name: str, arguments: dict[str, Any]) -> str | None:
        """Local-only checks run before the user confirms. 0 network calls."""
        return run_pre_confirm_guards(
            name, arguments, config=self.config, auth_manager=self.auth_manager
        )

    def _run_post_confirm_guards(self, name: str, arguments: dict[str, Any]) -> str | None:
        """Live-read checks run after confirmation. May touch the network."""
        return run_post_confirm_guards(
            name, arguments, config=self.config, auth_manager=self.auth_manager
        )

    # ------------------------------------------------------------------ #
    # Lifecycle
    # ------------------------------------------------------------------ #

    async def run_stdio(self) -> None:
        """Run the server over stdio. Phase 8 owns the CLI wrapper."""
        init_options = self.mcp_server.create_initialization_options(
            notification_options=NotificationOptions(),
            experimental_capabilities={},
        )
        async with stdio_server() as (read_stream, write_stream):
            await self.mcp_server.run(
                read_stream,
                write_stream,
                init_options,
            )


__all__ = [
    "A360MCPServer",
    "CONFIRM_FIELD",
    "DEFAULT_TOOL_PACKAGE",
    "ENV_TOOL_PACKAGE",
    "get_tool_schema",
    "is_mutating_tool",
    "is_read_only_call",
    "serialize_tool_output",
]
