#!/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


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")


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"
    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}


@mcp.tool(
    name="verify_server_health",
    description="Check whether the configured local LightRAG server is healthy.",
)
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.",
)
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.",
)
async def list_all_docs() -> Any:
    return await _request_lightrag("/documents")


@mcp.tool(
    name="query_knowledge_graph",
    description="Search the local LightRAG knowledge base. Use mix by default; references include provider file_path/chunk_id for MNote citation mapping.",
)
async def query_knowledge_graph(
    prompt: str,
    search_mode: str = "mix",
    limit: int = 60,
    context_only: bool = False,
    prompt_only: bool = False,
    include_references: bool = True,
    include_chunk_content: bool = True,
) -> Any:
    mode_aliases = {
        "keyword": "naive",
        "semantic": "hybrid",
    }
    mode = mode_aliases.get(search_mode, search_mode)
    body = {
        "query": prompt,
        "mode": mode,
        "top_k": limit,
        "chunk_top_k": limit,
        "only_need_context": context_only,
        "only_need_prompt": prompt_only,
        "include_references": include_references,
        "include_chunk_content": include_chunk_content,
        "response_type": "Multiple Paragraphs",
    }
    return await _request_lightrag("/query", 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.",
)
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
    async with httpx.AsyncClient(timeout=60) as client:
        response = await client.post(
            f"{_mnote_web_url()}/api/knowledge-rag/open-reference",
            headers=_mnote_headers(),
            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}
    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": response.status_code}


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