1#!/usr/bin/env python3
2# -*- coding: utf-8 -*-
3 
4from dataclasses import dataclass
5from datetime import datetime, timedelta
6from typing import Optional
7import re
8 
9 
10INDEX_REALTIME = "INDEX_REALTIME"
11KLINE_ANALYSIS = "KLINE_ANALYSIS"
12KLINE_CHART = "KLINE_CHART"  # 新增：K线绘图
13INTRADAY_ANALYSIS = "INTRADAY_ANALYSIS"
14VOLUME_ANALYSIS = "VOLUME_ANALYSIS"  # 新增：分时量能分析
15LIMIT_STATS = "LIMIT_STATS"
16MONEY_FLOW = "MONEY_FLOW"
17FUNDAMENTAL = "FUNDAMENTAL"
18STOCK_OVERVIEW = "STOCK_OVERVIEW"
19MARGIN_LHB = "MARGIN_LHB"
20SECTOR_ANALYSIS = "SECTOR_ANALYSIS"
21DERIVATIVES = "DERIVATIVES"
22FUND_BOND = "FUND_BOND"
23HK_US_MARKET = "HK_US_MARKET"
24NEWS = "NEWS"
25RESEARCH_REPORT = "RESEARCH_REPORT"
26STOCK_PICK = "STOCK_PICK"
27HELP = "HELP"  # 使用说明
28PORTFOLIO = "PORTFOLIO"  # 持仓管理
29 
30 
31# Common A-share stock aliases for quick name-to-symbol routing.
32# Keep this list lightweight and focused on frequently queried names.
33STOCK_NAME_MAP = {
34    "贵州茅台": "600519",
35    "茅台": "600519",
36    "宁德时代": "300750",
37    "比亚迪": "002594",
38    "五粮液": "000858",
39    "招商银行": "600036",
40    "中国平安": "601318",
41    "隆基绿能": "601012",
42    "药明康德": "603259",
43    "美的集团": "000333",
44    "格力电器": "000651",
45    "长城汽车": "601633",
46}
47 
48 
49@dataclass
50class IntentObj:
51    intent: str
52    query: str
53    symbol: Optional[str] = None
54    date: Optional[str] = None
55    period: Optional[str] = None
56    top_n: Optional[int] = None
57 
58 
59def _extract_symbol(query: str) -> Optional[str]:
60    m = re.search(r"\b(?:sh|sz)?(\d{6})\b", query.lower())
61    if m:
62        return m.group(1)
63 
64    for name in sorted(STOCK_NAME_MAP, key=len, reverse=True):
65        if name in query:
66            return STOCK_NAME_MAP[name]
67 
68    m = re.search(r"\b(hk\d{4,5}|[A-Z]{1,5})\b", query)
69    if m:
70        return m.group(1)
71 
72    return None
73 
74 
75def _extract_date(query: str) -> Optional[str]:
76    m = re.search(r"(\d{4})[-/]?(\d{2})[-/]?(\d{2})", query)
77    if m:
78        return f"{m.group(1)}-{m.group(2)}-{m.group(3)}"
79 
80    if "今天" in query or "今日" in query:
81        return datetime.now().strftime("%Y-%m-%d")
82    if "昨天" in query or "昨日" in query:
83        return (datetime.now() - timedelta(days=1)).strftime("%Y-%m-%d")
84 
85    return None
86 
87 
88def _extract_period(query: str) -> Optional[str]:
89    q = query.lower()
90    if "1m" in q or "1分钟" in query:
91        return "1"
92    if "5m" in q or "5分钟" in query:
93        return "5"
94    if "15m" in q or "15分钟" in query:
95        return "15"
96    if "30m" in q or "30分钟" in query:
97        return "30"
98    if "60m" in q or "60分钟" in query:
99        return "60"
100    if "周线" in query or "week" in q:
101        return "weekly"
102    if "月线" in query or "month" in q:
103        return "monthly"
104    if "日线" in query or "day" in q or "daily" in q:
105        return "daily"
106 
107    return None
108 
109 
110def _extract_top_n(query: str) -> Optional[int]:
111    m = re.search(r"top\s*(\d+)", query.lower())
112    if m:
113        return int(m.group(1))
114 
115    m = re.search(r"前\s*(\d+)\s*(名|条|个)?", query)
116    if m:
117        return int(m.group(1))
118    
119    # 支持"近N日"、"最近N天"等
120    m = re.search(r"近\s*(\d+)\s*(日|天|周|月)", query)
121    if m:
122        return int(m.group(1))
123    m = re.search(r"最近\s*(\d+)\s*(日|天|周|月)", query)
124    if m:
125        return int(m.group(1))
126 
127    return None
128 
129 
130def _classify_intent(query: str) -> str:
131    q = query.lower()
132 
133    if any(k in query for k in ["涨停", "跌停", "涨跌停"]):
134        return LIMIT_STATS
135    if any(k in query for k in ["推荐股票", "选股", "股票推荐", "有什么股票推荐"]):
136        return STOCK_PICK
137    if any(k in query for k in ["分时", "盘口", "逐笔"]):
138        return INTRADAY_ANALYSIS
139    # 分时量能分析
140    if any(k in query for k in ["量能", "放量", "缩量", "主力动向", "抢筹", "出货", "封单", "分时量能"]):
141        return VOLUME_ANALYSIS
142    if any(k in query for k in ["k线", "K线", "日线", "周线", "月线"]) or "kline" in q:
143        return KLINE_ANALYSIS
144    # 绘图功能：K线图、走势、行情图
145    if any(k in query for k in ["走势图", "趋势图", "行情图", "K线图", "k线图", "绘制", "画图", "图"]):
146        return KLINE_CHART
147    if any(k in query for k in ["怎么样", "分析", "看下", "评估", "综合"]):
148        return STOCK_OVERVIEW
149    if any(k in query for k in ["资金流", "主力资金", "北向资金", "南向资金", "东向资金", "行业资金", "板块资金"]):
150        return MONEY_FLOW
151    if any(k in query for k in ["基本面", "财报", "财务", "市盈率", "市净率", "估值", "roe", "ROE", "毛利率", "净利率", "资产负债率"]):
152        return FUNDAMENTAL
153    if any(k in query for k in ["融资融券", "龙虎榜", "两融", "融资余额", "融券余额"]):
154        return MARGIN_LHB
155    if any(k in query for k in ["研报", "研究报告", "机构评级"]):
156        return RESEARCH_REPORT
157    if "财经新闻" in query:
158        return NEWS
159 
160    if any(k in query for k in ["港股", "美股", "纳斯达克", "道琼斯", "标普", "恒生", "恒指"]) or any(k in q for k in ["nasdaq", "dow", "sp500", "s&p", "hang seng", "hk", "us"]):
161        return HK_US_MARKET
162    if any(k in query for k in ["期货", "期权", "衍生品", "主力合约", "if", "ih", "ic", "im"]):
163        return DERIVATIVES
164    if any(k in query for k in ["基金", "净值", "可转债", "转债", "债券", "etf", "ETF"]):
165        return FUND_BOND
166    if any(k in query for k in ["板块", "行业", "概念", "题材", "轮动", "涨幅榜", "跌幅榜"]):
167        return SECTOR_ANALYSIS
168 
169    if any(k in query for k in ["大盘", "指数", "上证", "深证", "创业板", "实时"]):
170        return INDEX_REALTIME
171 
172    # 使用说明
173    if any(k in query for k in ["介绍股市", "股市怎么用", "使用说明", "帮助", "help", "说明", "玩法", "有哪些功能"]):
174        return HELP
175 
176    # 持仓管理
177    if any(k in query for k in ["持仓", "仓位", "我的股票"]):
178        return PORTFOLIO
179 
180    return INDEX_REALTIME
181 
182 
183def parse_query(query: str) -> IntentObj:
184    query = (query or "").strip()
185    return IntentObj(
186        intent=_classify_intent(query),
187        query=query,
188        symbol=_extract_symbol(query),
189        date=_extract_date(query),
190        period=_extract_period(query),
191        top_n=_extract_top_n(query),
192    )
193