#!/usr/bin/env python3
"""MNote LightRAG MCP facade.

This is intentionally thin: LightRAG owns retrieval, MNote owns citation/open
mapping through its source registry.
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import httpx
from mcp.server.fastmcp import FastMCP
from mcp.types import ToolAnnotations


DEFAULT_ENV_FILE = "/mnt/Data1T/Mnote_data/lightrag/LightRAG/.env"
DEFAULT_MNOTE_WEB_URL = "http://127.0.0.1:3000"
DEFAULT_ROOT_URI = "file:///mnt/Data1T/Mnote_data/users/mnote-e2e/workspaces/my-space"
DEFAULT_WORKSPACE_ID = "local-ws:mnote-e2e:my-space"

mcp = FastMCP("MNote-LightRAG-Server")
READ_ONLY_TOOL = ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=False)


def _read_env_value(key: str, env_file: str) -> str:
    path = Path(env_file)
    if not path.exists():
        return ""
    for line in path.read_text(encoding="utf-8").splitlines():
        if not line.startswith(f"{key}="):
            continue
        value = line.split("=", 1)[1].strip()
        if (value.startswith('"') and value.endswith('"')) or (
            value.startswith("'") and value.endswith("'")
        ):
            value = value[1:-1]
        return value
    return ""


def _lightrag_base_url() -> str:
    env_file = os.environ.get("MNOTE_LIGHTRAG_ENV_FILE", DEFAULT_ENV_FILE)
    host = os.environ.get("LIGHTRAG_HOST") or _read_env_value("HOST", env_file) or "127.0.0.1"
    port = os.environ.get("LIGHTRAG_PORT") or _read_env_value("PORT", env_file) or "9621"
    if host in {"0.0.0.0", "::"}:
        host = "127.0.0.1"
    return f"http://{host}:{port}".rstrip("/")


def _lightrag_api_key() -> str:
    env_file = os.environ.get("MNOTE_LIGHTRAG_ENV_FILE", DEFAULT_ENV_FILE)
    return os.environ.get("LIGHTRAG_API_KEY") or _read_env_value("LIGHTRAG_API_KEY", env_file)


def _mnote_web_url() -> str:
    return os.environ.get("MNOTE_WEB_URL", DEFAULT_MNOTE_WEB_URL).rstrip("/")


def _mnote_headers() -> dict[str, str]:
    return {
        "content-type": "application/json",
        "x-mnote-actor-id": os.environ.get("MNOTE_ACTOR_ID", "mnote-e2e"),
        "x-mnote-actor-type": os.environ.get("MNOTE_ACTOR_TYPE", "user"),
    }


async def _request_lightrag(path: str, *, method: str = "GET", json_body: Any = None) -> Any:
    headers = {"accept": "application/json"}
    api_key = _lightrag_api_key()
    if api_key:
        # 当前 LightRAG /query 接受 X-API-Key；Bearer 在本机版本会返回 Invalid token。
        headers["X-API-Key"] = api_key
    async with httpx.AsyncClient(timeout=180) as client:
        response = await client.request(
            method,
            f"{_lightrag_base_url()}{path}",
            headers=headers,
            json=json_body,
        )
    try:
        payload = response.json()
    except Exception:
        payload = {"text": response.text}
    if response.status_code >= 400:
        return {"status": "error", "response": None, "error": payload, "httpStatus": response.status_code}
    return {"status": "success", "response": payload, "error": None, "httpStatus": response.status_code}


async def _request_mnote(path: str, *, method: str = "GET", json_body: Any = None, params: dict[str, Any] | None = None) -> Any:
    headers = _mnote_headers()
    if method.upper() == "GET":
        headers = {key: value for key, value in headers.items() if key != "content-type"}
    async with httpx.AsyncClient(timeout=180) as client:
        response = await client.request(
            method,
            f"{_mnote_web_url()}{path}",
            headers=headers,
            json=json_body,
            params=params,
        )
    try:
        payload = response.json()
    except Exception:
        payload = {"text": response.text}
    if response.status_code >= 400:
        return {"status": "error", "response": None, "error": payload, "httpStatus": response.status_code}
    return {"status": "success", "response": payload, "error": None, "httpStatus": response.status_code}


@mcp.tool(
    name="connect",
    description="No-op connection probe for agents that expect MCP servers to expose a connect tool.",
    annotations=READ_ONLY_TOOL,
)
async def connect() -> Any:
    return {
        "status": "success",
        "response": {
            "server": "mnote_lightrag_bridge",
            "connected": True,
        },
        "error": None,
        "httpStatus": 200,
    }


@mcp.tool(
    name="verify_server_health",
    description="Check whether the configured local LightRAG server is healthy.",
    annotations=READ_ONLY_TOOL,
)
async def verify_server_health() -> Any:
    return await _request_lightrag("/health")


@mcp.tool(
    name="check_indexing_status",
    description="Check the LightRAG document processing pipeline status.",
    annotations=READ_ONLY_TOOL,
)
async def check_indexing_status() -> Any:
    return await _request_lightrag("/documents/pipeline_status")


@mcp.tool(
    name="list_all_docs",
    description="List documents currently known to LightRAG.",
    annotations=READ_ONLY_TOOL,
)
async def list_all_docs() -> Any:
    return await _request_lightrag("/documents")


@mcp.tool(
    name="query_knowledge_graph",
    description="Ask MNote knowledge_rag.query. This calls MNote's LightRAG facade so references/citations are already mapped to MNote source registry and clickable locators.",
    annotations=READ_ONLY_TOOL,
)
async def query_knowledge_graph(
    prompt: str,
    search_mode: str = "mix",
    limit: int = 60,
    include_chunk_content: bool = True,
    include_document_structure_index: bool = False,
    source_paths: list[str] | None = None,
    workspace_id: str = DEFAULT_WORKSPACE_ID,
    root_uri: str = DEFAULT_ROOT_URI,
) -> Any:
    return await mnote_knowledge_rag_query(
        query=prompt,
        mode=search_mode,
        top_k=limit,
        chunk_top_k=limit,
        include_chunk_content=include_chunk_content,
        include_document_structure_index=include_document_structure_index,
        source_paths=source_paths,
        workspace_id=workspace_id,
        root_uri=root_uri,
    )


@mcp.tool(
    name="mnote_knowledge_rag_status",
    description="Call MNote mnote.knowledge_rag.status for provider health, source registry and sync state.",
    annotations=READ_ONLY_TOOL,
)
async def mnote_knowledge_rag_status(
    workspace_id: str = DEFAULT_WORKSPACE_ID,
    root_uri: str = DEFAULT_ROOT_URI,
) -> Any:
    return await _request_mnote(
        "/api/knowledge-rag/status",
        method="GET",
        params={"workspaceId": workspace_id, "rootUri": root_uri},
    )


@mcp.tool(
    name="mnote_knowledge_rag_query",
    description="Call MNote mnote.knowledge_rag.query. Use this for knowledge-library answers that need references/citations/clickable MNote locators. For book-like files pass source_paths and include_document_structure_index=true.",
    annotations=READ_ONLY_TOOL,
)
async def mnote_knowledge_rag_query(
    query: str,
    mode: str = "mix",
    top_k: int = 40,
    chunk_top_k: int = 20,
    include_chunk_content: bool = True,
    include_document_structure_index: bool = False,
    source_paths: list[str] | None = None,
    workspace_id: str = DEFAULT_WORKSPACE_ID,
    root_uri: str = DEFAULT_ROOT_URI,
) -> Any:
    mode_aliases = {
        "keyword": "naive",
        "semantic": "hybrid",
    }
    resolved_mode = mode_aliases.get(mode, mode)
    body = {
        "workspaceId": workspace_id,
        "rootUri": root_uri,
        "query": query,
        "mode": resolved_mode,
        "topK": top_k,
        "chunkTopK": chunk_top_k,
        "includeChunkContent": include_chunk_content,
        "includeDocumentStructureIndex": include_document_structure_index,
    }
    if source_paths:
        body["sourcePaths"] = source_paths
    return await _request_mnote("/api/knowledge-rag/query", method="POST", json_body=body)


@mcp.tool(
    name="mnote_knowledge_rag_section_context",
    description="Call MNote mnote.knowledge_rag.section_context. Use documentStructureIndex section ranges to fetch bounded LightRAG sidecar blocks/chunks for second-pass reading of book-like/long documents.",
    annotations=READ_ONLY_TOOL,
)
async def mnote_knowledge_rag_section_context(
    source_path: str = "",
    source_id: str = "",
    light_rag_doc_id: str = "",
    file_path: str = "",
    section_id: str = "",
    start_block_ordinal: int | None = None,
    end_block_ordinal: int | None = None,
    start_paragraph_ordinal: int | None = None,
    end_paragraph_ordinal: int | None = None,
    context_before: int = 1,
    context_after: int = 1,
    max_blocks: int = 24,
    max_chars: int = 12000,
    workspace_id: str = DEFAULT_WORKSPACE_ID,
    root_uri: str = DEFAULT_ROOT_URI,
) -> Any:
    body: dict[str, Any] = {
        "workspaceId": workspace_id,
        "rootUri": root_uri,
        "contextBefore": context_before,
        "contextAfter": context_after,
        "maxBlocks": max_blocks,
        "maxChars": max_chars,
    }
    optional_values = {
        "sourcePath": source_path,
        "sourceId": source_id,
        "lightRagDocId": light_rag_doc_id,
        "filePath": file_path,
        "sectionId": section_id,
        "startBlockOrdinal": start_block_ordinal,
        "endBlockOrdinal": end_block_ordinal,
        "startParagraphOrdinal": start_paragraph_ordinal,
        "endParagraphOrdinal": end_paragraph_ordinal,
    }
    for key, value in optional_values.items():
        if value is not None and value != "":
            body[key] = value
    return await _request_mnote("/api/knowledge-rag/section-context", method="POST", json_body=body)


@mcp.tool(
    name="open_mnote_reference",
    description="Map a LightRAG reference/file_path/chunk_id to an MNote clickable citationUrl using MNote source registry.",
    annotations=READ_ONLY_TOOL,
)
async def open_mnote_reference(
    file_path: str,
    chunk_id: str = "",
    reference_id: str = "",
    workspace_id: str = DEFAULT_WORKSPACE_ID,
    root_uri: str = DEFAULT_ROOT_URI,
    include_registry: bool = False,
) -> Any:
    body: dict[str, Any] = {
        "workspaceId": workspace_id,
        "rootUri": root_uri,
        "filePath": file_path,
    }
    if chunk_id:
        body["chunkId"] = chunk_id
    if reference_id:
        body["referenceId"] = reference_id
    result = await _request_mnote("/api/knowledge-rag/open-reference", method="POST", json_body=body)
    payload = result.get("response")
    if result.get("status") != "success":
        return result
    if not include_registry and isinstance(payload, dict):
        payload = {key: value for key, value in payload.items() if key != "registry"}
    return {"status": "success", "response": payload, "error": None, "httpStatus": result.get("httpStatus")}


@mcp.tool(
    name="mnote_knowledge_rag_open_reference",
    description="Call MNote mnote.knowledge_rag.open_reference to convert a returned reference/file_path/chunk_id into a clickable MNote locator.",
    annotations=READ_ONLY_TOOL,
)
async def mnote_knowledge_rag_open_reference(
    file_path: str,
    chunk_id: str = "",
    reference_id: str = "",
    workspace_id: str = DEFAULT_WORKSPACE_ID,
    root_uri: str = DEFAULT_ROOT_URI,
    include_registry: bool = False,
) -> Any:
    return await open_mnote_reference(
        file_path=file_path,
        chunk_id=chunk_id,
        reference_id=reference_id,
        workspace_id=workspace_id,
        root_uri=root_uri,
        include_registry=include_registry,
    )


if __name__ == "__main__":
    mcp.run()