"""
ERP MCP Server for XiaoZhi AI.

Reuses the existing /api/ai/chat pipeline (same path as Telegram text messages).
Logs in once on startup with ERP credentials to obtain a JWT, refreshes on 401.

Usage:
    1. Add to .env file (same directory as mcp_pipe.py):
       ERP_URL=http://your-erp-server:8000
       ERP_USERNAME=your-erp-username
       ERP_PASSWORD=your-erp-password
    2. python mcp_pipe.py erp_mcp.py

    MCP_ENDPOINT is already in .env — no need to set it again.
"""

import json
import logging
import os
import urllib.error
import urllib.request
from typing import Any

# Fix: mcp>=1.27 passes bare `str` to pydantic create_model which requires (type, default) tuples.
try:
    import mcp.server.fastmcp.utilities.func_metadata as _fm
    from pydantic import create_model as _create_model

    def _patched_create_wrapped_model(func_name: str, annotation: Any) -> Any:
        return _create_model(f"{func_name}Output", **{"result": (annotation, ...)})

    _fm._create_wrapped_model = _patched_create_wrapped_model
except Exception:
    pass

from mcp.server.fastmcp import FastMCP

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("erp_mcp")

mcp = FastMCP("ERP Assistant")

ERP_URL = os.getenv("ERP_URL", "http://localhost:8000").rstrip("/")
ERP_USERNAME = os.getenv("ERP_USERNAME", "")
ERP_PASSWORD = os.getenv("ERP_PASSWORD", "")

_jwt_token: str = ""

# Per-session conversation context (keyed by session_id passed in by caller)
_sessions: dict[str, list[dict]] = {}
_SESSION_CONTEXT_LIMIT = 16
_REPLY_LIMIT = 900  # stay under XiaoZhi's ~1024 byte response cap


def _http(method: str, path: str, body: Any = None, token: str = "") -> dict:
    """Minimal HTTP helper — returns parsed JSON or {"error": ...}."""
    url = f"{ERP_URL}{path}"
    data = json.dumps(body).encode() if body is not None else None
    headers: dict[str, str] = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(url, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=20) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        body_text = e.read().decode() if e.fp else ""
        return {"error": f"HTTP {e.code}", "detail": body_text[:200], "status_code": e.code}
    except Exception as e:
        return {"error": str(e)}


def _login() -> str:
    """Login and return JWT access_token, or empty string on failure."""
    if not ERP_USERNAME or not ERP_PASSWORD:
        logger.error("ERP_USERNAME / ERP_PASSWORD not set")
        return ""
    result = _http("POST", "/api/auth/login", {"username": ERP_USERNAME, "password": ERP_PASSWORD})
    token = result.get("access_token", "")
    if token:
        logger.info("ERP login successful")
    else:
        logger.error("ERP login failed: %s", result.get("error") or result)
    return token


def _ensure_token() -> str:
    """Return cached JWT, logging in first if not yet obtained."""
    global _jwt_token
    if not _jwt_token:
        _jwt_token = _login()
    return _jwt_token


def _ai_chat(message: str, context: list[dict]) -> str:
    """
    Call POST /api/ai/chat — the same endpoint the Telegram bot uses.
    Retries once after re-login on 401.
    Returns the reply string.
    """
    global _jwt_token

    payload = {"message": message, "context": context}

    for attempt in range(2):
        token = _ensure_token()
        if not token:
            return "ERP 登录失败，请检查 ERP_USERNAME / ERP_PASSWORD 配置。"

        result = _http("POST", "/api/ai/chat", payload, token=token)

        if result.get("error") and result.get("status_code") == 401:
            logger.warning("JWT expired, re-logging in")
            _jwt_token = ""
            continue

        if result.get("error"):
            logger.error("ai/chat error: %s", result)
            return f"ERP 接口错误: {result.get('error')}"

        reply = result.get("reply") or "（无回复）"
        logger.info("ai/chat ok, reply_len=%d", len(reply))
        return reply

    return "ERP 登录失败，无法获取有效 token。"


def _get_context(session_id: str) -> list[dict]:
    return _sessions.get(session_id, [])


def _save_context(session_id: str, user_msg: str, assistant_reply: str) -> None:
    ctx = _sessions.setdefault(session_id, [])
    ctx.extend([
        {"role": "user", "content": user_msg},
        {"role": "assistant", "content": assistant_reply},
    ])
    if len(ctx) > _SESSION_CONTEXT_LIMIT:
        _sessions[session_id] = ctx[-_SESSION_CONTEXT_LIMIT:]


def _truncate(text: str) -> str:
    encoded = text.encode()
    if len(encoded) <= _REPLY_LIMIT:
        return text
    return encoded[:_REPLY_LIMIT].decode(errors="ignore") + "…"


@mcp.tool()
def chat_with_erp(message: str, session_id: str = "default") -> str:
    """
    向ERP系统发送消息，查询业务数据或执行操作。
    当用户询问发票、客户、收入、支出、订阅、GST税务、票据等ERP相关信息时使用此工具。
    也可以通过自然语言创建发票草稿或新建客户资料。

    message: 用户的完整问题或指令，例如"查询本月未付款发票"、"上个月收入多少"。
    session_id: 会话标识，用于保持多轮对话上下文，同一次对话请传相同的值，默认 default。
    """
    ctx = _get_context(session_id)
    reply = _ai_chat(message, ctx)
    _save_context(session_id, message, reply)
    return _truncate(reply)


@mcp.tool()
def reset_erp_session(session_id: str = "default") -> str:
    """
    清除指定会话的对话历史，开启新一轮对话。
    当用户说"重新开始"、"清除记录"、"新对话"时使用此工具。
    session_id: 要重置的会话标识，默认 default。
    """
    _sessions.pop(session_id, None)
    logger.info("Session %s reset", session_id)
    return "已清除对话记录，可以开始新的对话。"


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