from .lesb import lesbRequest
from pydantic import BaseModel
from mcp.shared.exceptions import McpError
from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource
from mcp.server.stdio import stdio_server
from mcp.server import Server
from typing import Sequence
from enum import Enum
import json
import re

from pydantic import BaseModel


import logging

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(name)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s",
    handlers=[logging.FileHandler("ssp.log", encoding="utf-8")],
)
logger = logging.getLogger(__name__)


class SSPTools(str, Enum):
    QUERY_PERAGE_AND_CBXX = ("queryPerAgeAndCbxx",)
    GET_FDYCTXNL = ("getFdyctxnl",)
    GET_AGED_ACCOUNT = ("getAgedAccount",)
    QUERY_FDBD_AND_SPZZL = "queryFdbdAndSpzzl"
    CALC_STAFF_EXAM = "calcStaffExam"
    CALC_SPGZ = "calcSpgz"


class QueryPerAgeAndCbxxResult(BaseModel):
    _lesb__errcode_: str
    errflag: str
    errtext: str
    perinfo: str
    # 解析开的
    xb: str
    age: str
    jfnx: str
    csrq: str


class GetFdyctxnlResult(BaseModel):
    _lesb__errcode_: str
    errflag: str
    errtext: str
    fdyctxnl: str


class GetAgedAccountResult(BaseModel):
    _lesb__errcode_: str
    errflag: str
    errtext: str
    jfnx: str
    zhye: str


class QueryFdbdAndSpzzlResult(BaseModel):
    _lesb__errcode_: str
    errflag: str
    errtext: str
    jfjssx: str
    jfjsxx: str
    zzl: str


class CalcStaffExamResult(BaseModel):
    _lesb__errcode_: str
    errflag: str
    errtext: str
    wljnys: str  # 未来缴费月数
    txnl: str  # 退休年龄
    wljfbj: str  # 未来缴费成本
    # 每个月可以领取
    jcylj: str  # 基础养老金
    zhylj: str  # 账户养老金
    gdylj: str  # 过度养老金
    zgzhye: str


class CalcSpgzResult(BaseModel):
    spgz: float


def analyzePerInfo(arg1: str) -> dict:
    # 正则表达式：兼容"缴费年限"前的不同前缀（0到多个汉字），非贪婪模式跳过中间字段
    pattern = r"性别：([^，,]+)[，,]\s*出生日期：(\d{4}年\d{2}月)[，,]\s*年龄：(\d+岁\d+个月).*?(?:\S+)?缴费年限：(\d+年\d+个月)"

    match = re.search(pattern, arg1)
    if not match:
        raise ValueError("字段匹配失败，请检查输入格式")

    gender = match.group(1).strip()
    csrq = match.group(2).strip()
    age = match.group(3).strip()
    jfnx = match.group(4).strip()

    return {
        "xb": gender,
        "age": age,
        "jfnx": jfnx,
        "csrq": csrq,  # 出生日期
    }


class SSPServer:
    def queryPerAgeAndCbxx(
        self, ip: str, sfzhm: str, rsxtid: str
    ) -> QueryPerAgeAndCbxxResult:
        response = lesbRequest(
            ip,
            "SiServiceCenter",
            "queryPerAgeAndCbxx",
            {"sfzhm": sfzhm, "rsxtid": rsxtid},
        )

        perinfoObj = analyzePerInfo(response["perinfo"])
        for key in perinfoObj:
            response[key] = perinfoObj[key]

        return QueryPerAgeAndCbxxResult(**response)

    def getFdyctxnl(self, ip: str, sfzhm: str) -> GetFdyctxnlResult:
        response = lesbRequest(ip, "SiServiceCenter", "getFdyctxnl", {"sfzhm": sfzhm})

        return GetFdyctxnlResult(**response)

    def getAgedAccount(self, ip: str, sfzhm: str, rsxtid: str) -> GetAgedAccountResult:
        response = lesbRequest(
            ip, "SiUrService", "getAgedAccount", {"sfzhm": sfzhm, "rsxtid": rsxtid}
        )

        return GetAgedAccountResult(**response)

    def queryFdbdAndSpzzl(
        self, ip: str, zg_jbjgid: str, rsxtid: str, nd: str
    ) -> QueryFdbdAndSpzzlResult:
        response = lesbRequest(
            ip,
            "SiServiceCenter",
            "queryFdbdAndSpzzl",
            {"jbjgid": zg_jbjgid, "rsxtid": rsxtid, "nd": nd},
        )

        return QueryFdbdAndSpzzlResult(**response)

    def calcStaffExam(
        self,
        ip: str,
        sfzhm: str,
        yctxbz: str,  # 延迟退休标志
        pjgzzzl: str,  # 在岗职工平均工资增长率
        wlgrzhjzll: str,  # 未来个人账户记账利率
        jmjfje: str,  # 居民账户余额
        jfjs: str,
    ) -> CalcStaffExamResult:
        response = lesbRequest(
            ip,
            "SiServiceCenter",
            "calcStaffExam",
            {
                "shbzhm": sfzhm,
                "sfyt": yctxbz,
                "ypjgzzzl": pjgzzzl,
                "jfjs": jfjs,
                "jzll": wlgrzhjzll,
                "jmzhye": jmjfje,
            },
        )

        return CalcStaffExamResult(**response)

    def calcSpgz(self, ip: str, zg_jbjgid: str, rsxtid: str, nd: str) -> CalcSpgzResult:

        r = self.queryFdbdAndSpzzl(ip, zg_jbjgid, rsxtid, nd)

        spgz = float(r.jfjsxx) / 0.6

        return CalcSpgzResult(
            **{
                "spgz": round(spgz, 2),
            }
        )


async def serve(ip: str | None = None) -> None:

    if not ip:
        raise ValueError("Missing required argument: ip")

    logger.info(f"Using ip: {ip}")

    server = Server("mcp-ssp")
    ssp_server = SSPServer()

    @server.list_tools()
    async def list_tools() -> list[Tool]:
        """List available time tools."""
        return [
            Tool(
                name=SSPTools.QUERY_PERAGE_AND_CBXX.value,
                description="获取参保信息（性别、年龄、出生日期、参保年限）",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "sfzhm": {
                            "type": "string",
                            "description": f"身份证号码. (e.g., '37010119760201001X').",
                        },
                        "rsxtid": {
                            "type": "string",
                            "description": f"参保地区编号. (e.g., '3751').",
                        },
                    },
                    "required": ["sfzhm", "rsxtid"],
                },
            ),
            # Tool(
            #     name=SSPTools.GET_FDYCTXNL.value,
            #     description="获取法定延迟退休年龄",
            #     inputSchema={
            #         "type": "object",
            #         "properties": {
            #             "sfzhm": {
            #                 "type": "string",
            #                 "description": f"身份证号码. (e.g., '37010119760201001X').",
            #             }
            #         },
            #         "required": ["sfzhm"],
            #     },
            # ),
            Tool(
                name=SSPTools.GET_AGED_ACCOUNT.value,
                description="获取个人账户储存额",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "sfzhm": {
                            "type": "string",
                            "description": f"身份证号码. (e.g., '37010119760201001X').",
                        },
                        "rsxtid": {
                            "type": "string",
                            "description": f"参保地区编号. (e.g., '3751').",
                        },
                    },
                    "required": ["sfzhm", "rsxtid"],
                },
            ),
            # Tool(
            #     name=SSPTools.QUERY_FDBD_AND_SPZZL.value,
            #     description="查询某年度的缴费基数上下限",
            #     inputSchema={
            #         "type": "object",
            #         "properties": {
            #             "zg_jbjgid": {
            #                 "type": "string",
            #                 "description": f"职工经办机构编号. (e.g., '37019L01').",
            #             },
            #             "rsxtid": {
            #                 "type": "string",
            #                 "description": f"参保地区编号. (e.g., '3751').",
            #             },
            #             "nd": {
            #                 "type": "string",
            #                 "description": f"年度. (e.g., '2025').",
            #             },
            #         },
            #         "required": ["zg_jbjgid", "rsxtid", "nd"],
            #     },
            # ),
            Tool(
                name=SSPTools.CALC_STAFF_EXAM.value,
                description="测算养老保障水平",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "sfzhm": {
                            "type": "string",
                            "description": f"身份证号码. (e.g., '37010119760201001X').",
                        },
                        "yctxbz": {
                            "type": "string",
                            "description": f"延迟退休标志. (e.g., '1'、'0').",
                        },
                        "pjgzzzl": {
                            "type": "string",
                            "description": f"平均工资增长率. (e.g., '3').",
                        },
                        "wlgrzhjzll": {
                            "type": "string",
                            "description": f"未来个人账户记账利率. (e.g., '2').",
                        },
                        "jmjfje": {
                            "type": "string",
                            "description": f"个人账户储存额. (e.g., '10000').",
                        },
                        "jfjs": {
                            "type": "string",
                            "description": f"缴费基数. (e.g., '7506.67').",
                        },
                    },
                    "required": [
                        "sfzhm",
                        "yctxbz",
                        "pjgzzzl",
                        "wlgrzhjzll",
                        "jmjfje",
                        "jfjs",
                    ],
                },
            ),
            Tool(
                name=SSPTools.CALC_SPGZ.value,
                description="计算社平工资",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "zg_jbjgid": {
                            "type": "string",
                            "description": f"职工经办机构编号. (e.g., '37019L01').",
                        },
                        "rsxtid": {
                            "type": "string",
                            "description": f"参保地区编号. (e.g., '3751').",
                        },
                        "nd": {
                            "type": "string",
                            "description": f"年度. (e.g., '2025').",
                        },
                    },
                    "required": ["zg_jbjgid", "rsxtid", "nd"],
                },
            ),
        ]

    @server.call_tool()
    async def call_tool(
        name: str, arguments: dict
    ) -> Sequence[TextContent | ImageContent | EmbeddedResource]:
        """Handle tool calls for time queries."""
        try:
            match name:
                case SSPTools.QUERY_PERAGE_AND_CBXX.value:
                    sfzhm = arguments.get("sfzhm")
                    rsxtid = arguments.get("rsxtid")

                    if not sfzhm:
                        raise ValueError("Missing required argument: sfzhm")
                    if not rsxtid:
                        raise ValueError("Missing required argument: rsxtid")

                    result = ssp_server.queryPerAgeAndCbxx(ip, sfzhm, rsxtid)
                case SSPTools.GET_FDYCTXNL.value:
                    sfzhm = arguments.get("sfzhm")

                    if not sfzhm:
                        raise ValueError("Missing required argument: sfzhm")

                    result = ssp_server.getFdyctxnl(ip, sfzhm)
                case SSPTools.GET_AGED_ACCOUNT.value:
                    sfzhm = arguments.get("sfzhm")
                    rsxtid = arguments.get("rsxtid")

                    if not sfzhm:
                        raise ValueError("Missing required argument: sfzhm")
                    if not rsxtid:
                        raise ValueError("Missing required argument: rsxtid")

                    result = ssp_server.getAgedAccount(ip, sfzhm, rsxtid)
                case SSPTools.QUERY_FDBD_AND_SPZZL.value:
                    zg_jbjgid = arguments.get("zg_jbjgid")
                    nd = arguments.get("nd")
                    rsxtid = arguments.get("rsxtid")

                    if not rsxtid:
                        raise ValueError("Missing required argument: rsxtid")
                    if not nd:
                        raise ValueError("Missing required argument: nd")

                    result = ssp_server.queryFdbdAndSpzzl(ip, zg_jbjgid, rsxtid, nd)
                case SSPTools.CALC_STAFF_EXAM.value:
                    sfzhm = arguments.get("sfzhm")
                    yctxbz = arguments.get("yctxbz")
                    pjgzzzl = arguments.get("pjgzzzl")
                    wlgrzhjzll = arguments.get("wlgrzhjzll")
                    jmjfje = arguments.get("jmjfje")
                    jfjs = arguments.get("jfjs")

                    if not sfzhm:
                        raise ValueError("Missing required argument: sfzhm")
                    if not yctxbz:
                        raise ValueError("Missing required argument: yctxbz")
                    if not pjgzzzl:
                        raise ValueError("Missing required argument: pjgzzzl")
                    if not wlgrzhjzll:
                        raise ValueError("Missing required argument: wlgrzhjzll")
                    if not jmjfje:
                        raise ValueError("Missing required argument: jmjfje")
                    if not jfjs:
                        raise ValueError("Missing required argument: jfjs")

                    result = ssp_server.calcStaffExam(
                        ip, sfzhm, yctxbz, pjgzzzl, wlgrzhjzll, jmjfje, jfjs
                    )
                case SSPTools.CALC_SPGZ.value:
                    zg_jbjgid = arguments.get("zg_jbjgid")
                    nd = arguments.get("nd")
                    rsxtid = arguments.get("rsxtid")

                    result = ssp_server.calcSpgz(ip, zg_jbjgid, rsxtid, nd)

                case _:
                    raise ValueError(f"Unknown tool: {name}")

            return [
                TextContent(type="text", text=json.dumps(result.model_dump(), indent=2))
            ]
        except Exception as e:
            raise ValueError(f"Error processing mcp-server-time query: {str(e)}")

    options = server.create_initialization_options()
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, options)
    return None
