"""FastMCP server exposing secure code search via Scalekit auth."""
from __future__ import annotations

import logging
import os
from html import escape
from typing import Any, Protocol

from fastmcp import FastMCP
from fastmcp.server.auth import TokenVerifier
from fastmcp.server.auth.providers.scalekit import ScalekitProvider

from Retriever import Retriever

class _RetrieverProtocol(Protocol):
    def retrieve(self, query: str, *, top_k: int = 10,
                 repo: str | None = None, branch: str | None = None) -> list: ...


logger = logging.getLogger(__name__)


def _env_truthy(value: str | None) -> bool:
    return str(value or "").strip().lower() in {"1", "true", "yes", "on"}


def _resolve_require_auth(require_auth: bool | None) -> bool:
    if require_auth is not None:
        return require_auth
    if _env_truthy(os.environ.get("GITRAG_MCP_DISABLE_AUTH")):
        return False
    configured = os.environ.get("GITRAG_MCP_REQUIRE_AUTH")
    if configured is not None:
        return _env_truthy(configured)
    return True


def build_scalekit_provider(*, token_verifier: TokenVerifier | None = None,
                            base_url: str | None = None, ) -> ScalekitProvider:
    """Construct Scalekit auth provider from environment variables."""
    environment_url = os.environ.get("SCALEKIT_ENVIRONMENT_URL")
    client_id = os.environ.get("SCALEKIT_CLIENT_ID")
    resource_id = os.environ.get("SCALEKIT_RESOURCE_ID")

    missing = [
        name
        for name, value in {
            "SCALEKIT_ENVIRONMENT_URL": environment_url,
            "SCALEKIT_CLIENT_ID": client_id,
            "SCALEKIT_RESOURCE_ID": resource_id,
        }.items()
        if not value
    ]
    if missing:
        raise RuntimeError(f"Missing required Scalekit environment variables: {', '.join(missing)}")

    resolved_base_url = base_url or os.environ.get("MCP_BASE_URL") or "http://127.0.0.1:8000/mcp"

    return ScalekitProvider(
        environment_url=environment_url,
        client_id=client_id,
        resource_id=resource_id,
        base_url=resolved_base_url,
        token_verifier=token_verifier,
    )


def create_mcp_server(*, retriever: Retriever, token_verifier: TokenVerifier | None = None,
                      base_url: str | None = None, require_auth: bool | None = None, ) -> FastMCP:
    """Create an authenticated MCP server with `search_code` tool."""
    resolved_require_auth = _resolve_require_auth(require_auth)
    auth_provider = None
    if resolved_require_auth:
        auth_provider = build_scalekit_provider(token_verifier=token_verifier, base_url=base_url)
    else:
        logger.info("Authentication disabled for MCP server")
    mcp = FastMCP(name="GitRag MCP Server", auth=auth_provider)

    @mcp.tool(name="search_code")
    def search_code(
            query: str,
            top_k: int = 5,
            repo: str | None = None,
            branch: str | None = None,
    ) -> dict[str, Any]:
        """Search indexed code and return top snippets."""
        logger.info("search_code start query=%r top_k=%s repo=%r branch=%r", query, top_k, repo, branch)
        chunks = retriever.retrieve(query, top_k=top_k, repo=repo, branch=branch)
        output: list[dict[str, Any]] = []
        markdown_blocks: list[str] = []
        for chunk in chunks:
            payload = vars(chunk).copy()
            if isinstance(payload.get("embeddings"), (bytes, bytearray, memoryview)):
                payload["embeddings"] = None
            escaped_path = escape(str(payload.get("path", "")), quote=True)
            escaped_repo = escape(str(payload.get("repo", "")), quote=True)
            escaped_branch = escape(str(payload.get("branch", "")), quote=True)
            escaped_chunk = escape(str(payload.get("chunk", "")))
            payload["formatted"] = (
                f'<file path="{escaped_path}" repo="{escaped_repo}" '
                f'branch="{escaped_branch}">\n{escaped_chunk}\n</file>'
            )
            markdown_blocks.append(
                f"### {payload.get('path')}\n\n```{payload.get('language', '')}\n{payload.get('chunk', '')}\n```"
            )
            output.append(payload)
        response = {"results": output, "markdown": "\n\n".join(markdown_blocks)}
        logger.info("search_code end query=%r results=%s", query, len(output))
        return response

    return mcp