#!/usr/bin/env python3
"""
公交查询MCP服务
提供公交线路搜索和公交位置查询功能
"""

import asyncio
import json
import re
from typing import Any, Optional
from urllib.parse import quote, unquote, urlparse, parse_qs

import aiohttp
from bs4 import BeautifulSoup
from mcp.server import FastMCP
from mcp.types import Tool


def extract_bus_position(html_content: str) -> dict:
    """
    从公交站牌HTML中提取车牌号和对应站点

    Args:
        html_content: 公交站牌的HTML文本

    Returns:
        字典 {车牌号: 站点名称}
    """
    # 解析HTML
    soup = BeautifulSoup(html_content, 'html.parser')
    # 找到站点表格
    table = soup.find('table', id='tablevalue')
    if not table:
        return {"error": "未找到站点表格"}

    # 提取所有行
    rows = table.find_all('tr')
    bus_info = {}  # 存储最终结果：{车牌号: 站点名}
    last_station = ""  # 记录上一个有效站点名称

    # 遍历每一行，定位车辆位置
    for i, row in enumerate(rows):
        # 获取当前行的所有列
        tds = row.find_all(['th', 'td'])
        if len(tds) < 3:
            continue

        # 识别"站点行"（有站点名称的行）
        station_name_td = tds[2]  # 第三列是站点名称列
        station_name = station_name_td.get_text(strip=True)
        if station_name and station_name != last_station:
            last_station = station_name  # 更新最近的站点名称

        # 识别"带车牌号的行"（车辆位置行）
        # 车牌号在第一列的span标签内
        left_td = tds[0]  # 第一列是车牌号/起终点标记列
        # 获取第一列的所有文本（包括嵌套标签中的文本）
        left_text = left_td.get_text(strip=True)
        # 匹配车牌号（鲁F+数字+可选字母，D可选的格式）
        plate_match = re.search(r'鲁F\d+[A-Z]?D?', left_text)
        if plate_match:
            plate_number = plate_match.group()
            # 车牌号行的站点名称需要从上一行获取
            # 查找上一行的站点名称
            if i > 0:
                prev_row = rows[i - 1]
                prev_tds = prev_row.find_all(['th', 'td'])
                if len(prev_tds) >= 3:
                    prev_station = prev_tds[2].get_text(strip=True)
                    if prev_station:
                        bus_info[plate_number] = prev_station
                    else:
                        bus_info[plate_number] = last_station
                else:
                    bus_info[plate_number] = last_station
            else:
                bus_info[plate_number] = last_station

    return bus_info


def extract_bus_line_info(html_content: str) -> dict:
    """
    从公交详情页HTML中提取核心线路信息

    Args:
        html_content: 完整的公交详情页HTML文本

    Returns:
        字典，包含线路名称、起止点、首末班、票价、运行季节等信息
    """
    # 初始化结果字典
    line_info = {
        "线路名称": "",
        "起点站": "",
        "终点站": "",
        "运行季节": "",
        "首班时间": "",
        "末班时间": "",
        "票价": ""
    }

    # 解析HTML
    soup = BeautifulSoup(html_content, 'html.parser')

    # 提取线路名称（页面标题/导航栏标题）
    title_tag = soup.find('h1', class_='mui-title')
    if title_tag:
        line_info["线路名称"] = title_tag.get_text(strip=True)

    # 提取起止点（核心：起/终按钮旁边的大字号文本）
    media_body = soup.find('div', class_='mui-media-body')
    if media_body:
        # 找到所有包含"起/终"按钮的p标签
        p_tags = media_body.find_all('p', class_='mui-ellipsis')
        for p in p_tags:
            # 提取按钮文字和站点名称
            btn = p.find('button', class_='mui-btn-primary')
            if btn:
                btn_text = btn.get_text(strip=True)
                # 提取按钮旁边的span文本（站点名）
                span = p.find('span', style=lambda s: s and 'font-size:20px' in s)
                if span:
                    station_name = span.get_text(strip=True)
                    if btn_text == "起":
                        line_info["起点站"] = station_name
                    elif btn_text == "终":
                        line_info["终点站"] = station_name

    # 提取运行季节、首末班时间
    time_div = soup.find('div', style=lambda s: s and 'padding:0px 20px; font-size: 12px;' in s)
    if time_div:
        p_tags = time_div.find_all('p')
        for p in p_tags:
            p_text = p.get_text(strip=True)
            if "夏季" in p_text or "冬季" in p_text:
                line_info["运行季节"] = p_text
            elif "首班" in p_text:
                line_info["首班时间"] = p_text
            elif "末班" in p_text:
                line_info["末班时间"] = p_text

    # 提取票价（页面中"票价X.00元"的文本）
    price_div = soup.find('div', style=lambda s: s and 'position: absolute; bottom:0.8rem; right:0.8rem;' in s)
    if price_div:
        line_info["票价"] = price_div.get_text(strip=True)

    return line_info


def parse_bus_url(url: str) -> dict:
    """
    从详情页链接中提取参数

    Args:
        url: 详情页链接，如 "./index.php?i=5&c=entry&linename=67路&search=search&do=line&m=ytbus"

    Returns:
        解析后的参数字典
    """
    # 解析URL
    parsed = urlparse(url)
    # 解析查询参数
    params = parse_qs(parsed.query)
    
    # 提取参数并解码
    result = {
        "i": params.get("i", ["5"])[0],
        "c": params.get("c", ["entry"])[0],
        "linename": unquote(params.get("linename", [""])[0]),
        "do": params.get("do", ["line"])[0],
        "m": params.get("m", ["ytbus"])[0]
    }
    
    return result


# 创建MCP服务器
mcp = FastMCP("bus-query")

# API基础URL
BASE_URL = "https://ytbus.edong500.com/app/index.php"


async def search_bus_routes(session: aiohttp.ClientSession, key: str) -> dict:
    """
    搜索可用公交线路

    Args:
        session: aiohttp会话
        key: 查询参数，如"67"表示67路公交

    Returns:
        解析后的线路信息
    """
    url = f"{BASE_URL}?i=5&c=entry&do=search_new&m=ytbus"
    data = {"key": key}

    try:
        async with session.post(url, data=data) as response:
            html = await response.text()
            
            # 解析HTML响应，提取公交线路信息
            soup = BeautifulSoup(html, 'html.parser')
            
            # 找到线路列表的ul（排除搜索框的ul，定位线路列表）
            bus_line_uls = soup.find_all('ul', class_='mui-table-view')
            
            if len(bus_line_uls) < 2:
                return {
                    "success": True,
                    "key": key,
                    "routes": [],
                    "message": "未找到公交线路列表",
                    "html_snippet": html[:500] if html else ""
                }
            
            bus_line_ul = bus_line_uls[1]
            routes = []
            
            # 遍历每个线路li
            for li in bus_line_ul.find_all('li', class_='mui-table-view-cell'):
                a_tag = li.find('a', class_='mui-navigate-right')
                if not a_tag:
                    continue
                
                # 提取线路名称
                line_name_div = a_tag.find('div', style='float: left; margin-left: 0.5rem;')
                line_name = line_name_div.get_text(strip=True) if line_name_div else ""
                
                # 提取详情页链接
                line_url = a_tag.get('href', "")
                # 解码URL
                decoded_url = unquote(line_url)
                
                if line_name and line_url:
                    routes.append({
                        "line_name": line_name,
                        "url": decoded_url
                    })
            
            return {
                "success": True,
                "key": key,
                "routes": routes,
                "html_snippet": html[:500] if html else ""
            }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "key": key
        }


async def get_bus_location(
    session: aiohttp.ClientSession,
    detail_url: str,
    direction: str = "上行"
) -> dict:
    """
    查询公交位置

    Args:
        session: aiohttp会话
        detail_url: 从搜索结果中获取的详情页链接，用于提取查询参数
        direction: 方向，"上行"或"下行"

    Returns:
        公交位置信息，包含车牌号和所在站点的对应关系
    """
    # 从detail_url中解析参数
    params = parse_bus_url(detail_url)
    linename_param = params["linename"]
    line_name = linename_param  # 用于返回结果

    # 构造查询URL
    url = f"{BASE_URL}?i=5&c=entry&linename={quote(linename_param)}&upordown={quote(direction)}&sname=&do=line&m=ytbus"

    try:
        async with session.post(url) as response:
            html = await response.text()

            # 使用BeautifulSoup解析HTML，提取车牌号和站点的对应关系
            bus_info = extract_bus_position(html)

            # 检查是否有错误
            if "error" in bus_info:
                return {
                    "success": False,
                    "error": bus_info["error"],
                    "line_name": line_name,
                    "direction": direction,
                    "html_snippet": html[:500] if html else ""
                }

            # 格式化返回结果
            bus_locations = [
                {
                    "license_plate": plate,
                    "station": station,
                    "direction": direction
                }
                for plate, station in bus_info.items()
            ]

            return {
                "success": True,
                "line_name": line_name,
                "direction": direction,
                "bus_count": len(bus_locations),
                "bus_locations": bus_locations,
                "html_snippet": html[:500] if html else ""
            }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "line_name": line_name,
            "direction": direction
        }


async def get_bus_line_info(
    session: aiohttp.ClientSession,
    detail_url: str
) -> dict:
    """
    查询公交线路的基本信息

    Args:
        session: aiohttp会话
        detail_url: 从搜索结果中获取的详情页链接

    Returns:
        公交线路基本信息，包含起点站、终点站、首末班时间、票价等
    """
    # 从detail_url中解析参数
    params = parse_bus_url(detail_url)
    linename_param = params["linename"]
    line_name = linename_param  # 用于返回结果

    # 构造查询URL（默认查询上行方向）
    url = f"{BASE_URL}?i=5&c=entry&linename={quote(linename_param)}&upordown=上行&sname=&do=line&m=ytbus"

    try:
        async with session.post(url) as response:
            html = await response.text()

            # 使用BeautifulSoup解析HTML，提取线路基本信息
            line_info = extract_bus_line_info(html)

            # 如果没有提取到线路名称，使用解析出来的参数
            if not line_info.get("线路名称"):
                line_info["线路名称"] = line_name

            return {
                "success": True,
                "line_name": line_name,
                "line_info": line_info,
                "html_snippet": html[:500] if html else ""
            }
    except Exception as e:
        return {
            "success": False,
            "error": str(e),
            "line_name": line_name
        }


@mcp.tool()
async def search_bus_routes_tool(key: str) -> str:
    """搜索可用的公交线路。根据关键词查找匹配的公交路线。

    Args:
        key: 搜索关键词，如'67'表示67路公交
    """
    async with aiohttp.ClientSession() as session:
        result = await search_bus_routes(session, key)
        return json.dumps(result, ensure_ascii=False, indent=2)


@mcp.tool()
async def get_bus_location_tool(detail_url: str, direction: str = "上行") -> str:
    """查询指定公交线路的实时位置信息，包括车牌号等。

    Args:
        detail_url: 从search_bus_routes_tool获取的详情页链接，用于精确查询
        direction: 公交方向，'上行'或'下行'，默认为'上行'
    """
    async with aiohttp.ClientSession() as session:
        result = await get_bus_location(session, detail_url, direction)
        return json.dumps(result, ensure_ascii=False, indent=2)


@mcp.tool()
async def get_bus_line_info_tool(detail_url: str) -> str:
    """查询公交线路的基本信息，包括起点站、终点站、首末班时间、票价等。

    Args:
        detail_url: 从search_bus_routes_tool获取的详情页链接，用于精确查询
    """
    async with aiohttp.ClientSession() as session:
        result = await get_bus_line_info(session, detail_url)
        return json.dumps(result, ensure_ascii=False, indent=2)


def main():
    """命令行入口点"""
    mcp.run()


if __name__ == "__main__":
    main()