Capture Alphapai Unified

{
  "api_key": "",
  "base_url": "https://open-api.rabyte.cn",
  "account": "",
  "password": ""
}

#!/usr/bin/env python3
"""
AlphaPai Research CLI
=====================
Alpha派投研 API 命令行工具。

用法：
  python alphapai_client.py config --set-key YOUR_KEY
  python alphapai_client.py qa --question "贵州茅台2024年经营情况？"
  python alphapai_client.py recall --query "茅台市值" --type comment,qa --start 2025-01-01

模块用法（供其他脚本导入）：
  from alphapai_client import AlphaPaiClient, load_config
  client = AlphaPaiClient(load_config())
  result = client.qa_text("问题")
"""

import argparse
import json
import os
import sys
import urllib.request
import urllib.error
from typing import Any, Dict, List, Optional

# ============================================================
# 配置管理
# ============================================================

DEFAULT_CONFIG_PATH = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "config.json"
)


def _parse_env_file(env_path: str) -> Dict[str, str]:
    """解析 .env 文件，返回 key-value 字典"""
    result = {}
    if not os.path.exists(env_path):
        return result
    try:
        with open(env_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" in line:
                    key, _, value = line.partition("=")
                    key = key.strip()
                    value = value.strip().strip('"').strip("'")
                    if key:
                        result[key] = value
    except Exception:
        pass
    return result


def _find_env_value(key: str, start_dir: str) -> Optional[str]:
    """从当前目录向上逐级查找 .env 文件中的 key
    
    查找顺序：
    1. start_dir/.env
    2. start_dir/../.env (上一级)
    3. start_dir/../../.env (上上级)
    """
    env_file = ".env"
    max_levels = 3  # 当前 + 上两级
    
    current_dir = start_dir
    for _ in range(max_levels):
        env_path = os.path.join(current_dir, env_file)
        env_dict = _parse_env_file(env_path)
        if key in env_dict and env_dict[key]:
            return env_dict[key]
        # 向上一级
        parent = os.path.dirname(current_dir)
        if parent == current_dir:  # 已到根目录
            break
        current_dir = parent
    
    return None


def load_config(config_path: str = DEFAULT_CONFIG_PATH) -> Optional[Dict]:
    """加载API配置，按优先级查找
    
    查找顺序：
    1. 环境变量 ALPHAPAI_API_KEY / ALPHAPAI_BASE_URL
    2. 同级目录 .env
    3. 上一级目录 .env
    4. 上上级目录 .env
    5. config.json
    
    返回 {"api_key": "...", "base_url": "..."} 或 None
    """
    base_url = "https://open-api.rabyte.cn"
    
    # 1. 优先从环境变量读取（推荐方式，最安全）
    env_key = os.environ.get("ALPHAPAI_API_KEY")
    env_url = os.environ.get("ALPHAPAI_BASE_URL")
    if env_url:
        base_url = env_url
    
    if env_key:
        return {"api_key": env_key, "base_url": base_url}
    
    # 2-4. 从 .env 文件逐级查找
    skill_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    env_key = _find_env_value("ALPHAPAI_API_KEY", skill_dir)
    env_url = _find_env_value("ALPHAPAI_BASE_URL", skill_dir)
    
    if env_url:
        base_url = env_url
    
    if env_key:
        return {"api_key": env_key, "base_url": base_url}
    
    # 5. fallback 到 config.json（兼容旧方式）
    if os.path.exists(config_path):
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
            if config.get("api_key"):
                config.setdefault("base_url", base_url)
                return config
        except Exception:
            pass
    
    return None


def save_config(
    api_key: str,
    base_url: str = "https://open-api.rabyte.cn",
    config_path: str = DEFAULT_CONFIG_PATH,
) -> Dict:
    """保存API配置到 config.json"""
    config = {"api_key": api_key, "base_url": base_url}
    os.makedirs(os.path.dirname(config_path), exist_ok=True)
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)
    return config


# ============================================================
# API Client（可单独导入使用）
# ============================================================


class AlphaPaiClient:
    """AlphaPai Open API 客户端（零第三方依赖，使用 urllib）"""

    def __init__(self, config: Dict):
        self.base_url = config["base_url"].rstrip("/")
        self.headers = {
            "app-agent": config["api_key"],
            "Content-Type": "application/json",
        }

    def _post(self, endpoint: str, payload: Dict, stream: bool = False) -> Any:
        """通用 POST 请求。

        stream=True 时返回原始 HTTPResponse（用于 SSE 流式解析），
        stream=False 时返回解析后的 JSON 字典。
        """
        url = f"{self.base_url}{endpoint}"
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        req = urllib.request.Request(url, data=body, headers=self.headers, method="POST")
        resp = urllib.request.urlopen(req, timeout=120)
        if stream:
            return resp
        data = resp.read().decode("utf-8", errors="replace")
        return json.loads(data)

    @staticmethod
    def parse_sse(response) -> Dict:
        """解析 SSE 流式响应，返回 {"answer": str, "references": list}。

        从 urllib HTTPResponse 逐块读取，拼接后按 SSE 事件（双换行分隔）解析。
        """
        buffer, answer, references = "", "", []
        while True:
            chunk = response.read(4096)
            if not chunk:
                break
            buffer += chunk.decode("utf-8", errors="replace")
            while "\n\n" in buffer:
                event, buffer = buffer.split("\n\n", 1)
                event = event.strip()
                if not event.startswith("data: "):
                    continue
                line = event[6:].strip()
                if not line:
                    continue
                try:
                    data = json.loads(line)
                    if isinstance(data, dict):
                        d = data.get("data", data)
                        if d.get("answer"):
                            answer += d["answer"]
                        if d.get("references"):
                            references.extend(d["references"])
                except json.JSONDecodeError:
                    continue
        return {"answer": answer, "references": references}

    def qa_text(
        self,
        question: str,
        context: Optional[List[str]] = None,
        mode: str = "Flash",
        is_stream: bool = False,
        is_auto_route: bool = False,
        is_web_search: bool = False,
        is_deep_reasoning: bool = False,
        start_time: Optional[str] = None,
        end_time: Optional[str] = None,
        question_id: Optional[str] = None,
    ) -> Dict:
        """投研知识问答接口。is_stream=True时内部解析SSE并返回聚合结果。"""
        payload = {
            "question": question,
            "mode": mode,
            "isStream": is_stream,
            "isAutoRoute": is_auto_route,
            "userInfo": {
                "userCode": "001010",
                "userName": "rabyte",
                "roleName": "基金经理",
                "roleCode": 5,
            },
        }
        if context:
            payload["context"] = context
        if is_web_search:
            payload["isWebSearch"] = True
        if is_deep_reasoning:
            payload["isDeepReasoning"] = True
        if start_time:
            payload["requestSelectStartTime"] = start_time
        if end_time:
            payload["requestSelectEndTime"] = end_time
        if question_id:
            payload["questionId"] = question_id

        if is_stream:
            return self.parse_sse(
                self._post("/alpha/open-api/v1/paipai/qa-text", payload, stream=True)
            )
        return self._post("/alpha/open-api/v1/paipai/qa-text", payload)

    def recall_data(
        self,
        query: str,
        is_cut_off: bool = True,
        recall_type: Optional[List[str]] = None,
        start_time: Optional[str] = None,
        end_time: Optional[str] = None,
    ) -> Dict:
        """获取问题召回数据"""
        payload = {"query": query, "isCutOff": is_cut_off}
        if recall_type:
            payload["recallType"] = recall_type
        else:
            payload["recallType"] = []
        if start_time:
            payload["startTime"] = start_time
        if end_time:
            payload["endTime"] = end_time
        return self._post("/alpha/open-api/v1/paipai/recall-data", payload)

    def stock_agent(
        self,
        question: str,
        agent_mode: int,
        stock: Optional[Dict] = None,
        template: int = 0,
        template_text: str = "",
        report_type: Optional[str] = None,
        stock_report_id: Optional[str] = None,
        stock_report_title: Optional[str] = None,
        stock_report_period: Optional[str] = None,
        template_concern: Optional[str] = None,
        request_select_start_time: Optional[str] = None,
        request_select_end_time: Optional[str] = None,
        input_industry: Optional[str] = None,
        report_date: Optional[str] = None,
        fund_type: Optional[str] = None,
        if_annual: Optional[int] = None,
        stock_list: Optional[List[Dict]] = None,
        picture_color: Optional[List[str]] = None,
        picture_style: Optional[str] = None,
        source: Optional[int] = None,
        language: Optional[str] = None,
        only_answer: Optional[bool] = None,
    ) -> Dict:
        """Alpha派投研Agent接口（公司一页纸、业绩点评、调研大纲等），SSE 流式请求。"""
        payload: Dict = {
            "question": question,
            "agentMode": agent_mode,
            "template": template,
            "templateText": template_text,
        }
        if stock:
            payload["stock"] = stock
        if report_type:
            payload["reportType"] = report_type
        if stock_report_id:
            payload["stockReportId"] = stock_report_id
        if stock_report_title:
            payload["stockReportTitle"] = stock_report_title
        if stock_report_period:
            payload["stockReportPeriod"] = stock_report_period
        if template_concern:
            payload["templateConcern"] = template_concern
        if request_select_start_time:
            payload["requestSelectStartTime"] = request_select_start_time
        if request_select_end_time:
            payload["requestSelectEndTime"] = request_select_end_time
        if input_industry:
            payload["inputIndustry"] = input_industry
        if report_date:
            payload["reportDate"] = report_date
        if fund_type:
            payload["fundType"] = fund_type
        if if_annual is not None:
            payload["ifAnnual"] = if_annual
        if stock_list:
            payload["stockList"] = stock_list
        if picture_color:
            payload["pictureColor"] = picture_color
        if picture_style:
            payload["pictureStyle"] = picture_style
        if source is not None:
            payload["source"] = source
        if language:
            payload["language"] = language
        if only_answer is not None:
            payload["onlyAnswer"] = only_answer
        return self.parse_sse(
            self._post("/alpha/open-api/v1/paipai/stock/agent", payload, stream=True)
        )

    def stock_report(self, code: str) -> Dict:
        """获取某个股票的公告列表（上市公司在交易所公开的公告）"""
        return self._post("/alpha/open-api/v1/paipai/stock/report", {"code": code})

    def search_image(
        self,
        query_text: str,
        files_range: Optional[List[str]] = None,
        topk: int = 50,
        recall_mode: str = "both",
        use_llm_rank: bool = False,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
    ) -> Dict:
        """搜图表：从研报/公告中搜索相关图片和表格"""
        payload: Dict = {
            "queryText": query_text,
            "topk": topk,
            "recallMode": recall_mode,
            "useLlmRank": use_llm_rank,
        }
        if files_range:
            payload["filesRange"] = files_range
        if start_date:
            payload["startDate"] = start_date
        if end_date:
            payload["endDate"] = end_date
        return self._post("/alpha/open-api/v1/paipai/search-image", payload)


# ============================================================
# 输出格式化
# ============================================================

_TYPE_LABEL = {
    "comment": "机构点评",
    "vps": "基金定期报告",
    "report": "内资研报",
    "foreign_report": "外资研报",
    "wechat_public_article": "公众号",
    "ann": "公司公告",
    "roadShow": "路演纪要",
    "roadShow_ir": "上市公司披露投关纪要",
    "roadShow_us": "美股纪要",
    "table": "公告/研报表格",
    "image": "研报图片",
    "edb": "EDB数据库",
}


def _type_label(type_val: str) -> str:
    return _TYPE_LABEL.get(type_val, type_val)


def _require_success(result: Dict):
    """检查响应码，非成功则打印错误并退出"""
    if result.get("code") != 200000:
        print(
            f"[错误] code={result.get('code')} message={result.get('message')}",
            file=sys.stderr,
        )
        sys.exit(1)


def format_qa(result: Dict) -> str:
    # parse_sse returns {"answer": ..., "references": [...]} directly (no code wrapper)
    if "answer" in result:
        answer = result.get("answer", "（无回答）")
        references = result.get("references", [])
    else:
        _require_success(result)
        data = result.get("data", {})
        answer = data.get("answer", "（无回答）")
        references = data.get("references", [])

    lines = ["[回答]", answer]
    if references:
        lines.append(f"\n[引用来源 {len(references)} 条]")
        for i, ref in enumerate(references, 1):
            date = f" ({ref.get('publishDate', '')})" if ref.get("publishDate") else ""
            lines.append(
                f"  {i}. [{_type_label(ref.get('type', ''))}] {ref.get('title', '')}{date}"
            )
    return "\n".join(lines)


def format_recall(result: Dict) -> str:
    _require_success(result)
    items = result.get("data", [])
    lines = [f"[召回数据 {len(items)} 条]\n"]
    for i, item in enumerate(items, 1):
        item_type = item.get("type", "")
        context_info = item.get("contextInfo", "")
        lines.append(f"#{i} [{_type_label(item_type)}]")
        lines.append(f"  {context_info}")
        if item_type == "qa":
            if item.get("contextText"):
                lines.append(f"  Q: {item['contextText']}")
            if item.get("answer"):
                text = item["answer"]
                lines.append(f"  A: {text[:300]}{'...' if len(text) > 300 else ''}")
        elif item.get("chunks"):
            text = item["chunks"][0]
            lines.append(f"  {text[:300]}{'...' if len(text) > 300 else ''}")
        lines.append("")
    return "\n".join(lines)


def format_agent(result: Dict) -> str:
    if "answer" in result:
        answer = result.get("answer", "（无回答）")
        references = result.get("references", [])
    else:
        _require_success(result)
        data = result.get("data", {})
        answer = data.get("answer", "（无回答）")
        references = data.get("references", [])
    lines = ["[Agent 回答]", answer]
    if references:
        lines.append(f"\n[引用来源 {len(references)} 条]")
        for i, ref in enumerate(references, 1):
            date = f" ({ref.get('publishDate', '')})" if ref.get("publishDate") else ""
            lines.append(
                f"  {i}. [{_type_label(ref.get('type', ''))}] {ref.get('title', '')}{date}"
            )
    return "\n".join(lines)


def format_report(result: Dict) -> str:
    _require_success(result)
    items = result.get("data", [])
    if not items:
        return "[公告列表] 未找到相关公告"
    lines = [f"[公告列表 {len(items)} 条]\n"]
    for i, item in enumerate(items, 1):
        lines.append(
            f"#{i} [{item.get('reportType', '')}] {item.get('stockReportTitle', '')}"
        )
        lines.append(
            f"   期间: {item.get('reportPeriod', '')}  ID: {item.get('stockReportId', '')}"
        )
        lines.append("")
    return "\n".join(lines)


def format_image(result: Dict) -> str:
    _require_success(result)
    items = result.get("data", [])
    if not items:
        return "[图表搜索] 未找到相关图表"
    lines = [f"[图表搜索 {len(items)} 条]\n"]
    for i, item in enumerate(items, 1):
        lines.append(f"#{i} {item.get('articleTitle', '')}")
        caption = ", ".join(item.get("captionList") or [])
        if caption:
            lines.append(f"   说明: {caption}")
        lines.append(
            f"   来源: {item.get('source', '')}  日期: {item.get('publishDate', '')}"
        )
        if item.get("imageUrl"):
            lines.append(f"   图片: {item['imageUrl']}")
        lines.append("")
    return "\n".join(lines)


# ============================================================
# CLI 子命令
# ============================================================


def cmd_config(args):
    config_path = DEFAULT_CONFIG_PATH
    if args.set_key or args.set_url:
        config = load_config(config_path) or {}
        if args.set_key:
            config["api_key"] = args.set_key
        if args.set_url:
            config["base_url"] = args.set_url
        config.setdefault("base_url", "https://open-api.rabyte.cn")
        save_config(config["api_key"], config["base_url"], config_path)
        key = config["api_key"]
        masked = key[:4] + "****" + key[-4:] if len(key) > 8 else "****"
        print(f"配置已保存: api_key={masked}  base_url={config['base_url']}")
        print(f"路径: {config_path}")
    elif args.show:
        config = load_config(config_path)
        if not config:
            print(
                "未找到配置，请先运行: python alphapai_client.py config --set-key YOUR_KEY"
            )
            sys.exit(1)
        key = config["api_key"]
        masked = key[:4] + "****" + key[-4:] if len(key) > 8 else "****"
        print(f"api_key : {masked}")
        print(f"base_url: {config['base_url']}")
        print(f"路径    : {config_path}")
    else:
        print(
            "用法: python alphapai_client.py config --show | --set-key KEY [--set-url URL]"
        )


def cmd_qa(args):
    config = load_config()
    if not config:
        print(
            "[错误] 未找到配置，请先运行: python alphapai_client.py config --set-key YOUR_KEY",
            file=sys.stderr,
        )
        sys.exit(1)

    client = AlphaPaiClient(config)
    result = client.qa_text(
        question=args.question,
        context=args.context or None,
        mode=args.mode,
        is_stream=True,
        is_web_search=args.web_search,
        is_deep_reasoning=args.deep_reasoning,
        start_time=args.start,
        end_time=args.end,
    )

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(format_qa(result))


def cmd_recall(args):
    config = load_config()
    if not config:
        print(
            "[错误] 未找到配置，请先运行: python alphapai_client.py config --set-key YOUR_KEY",
            file=sys.stderr,
        )
        sys.exit(1)

    client = AlphaPaiClient(config)
    recall_type = [t.strip() for t in args.type.split(",")] if args.type else []

    result = client.recall_data(
        query=args.query,
        is_cut_off=not args.no_cutoff,
        recall_type=recall_type,
        start_time=args.start,
        end_time=args.end,
    )

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(format_recall(result))


def _parse_stock(s: str) -> Dict:
    """解析 CODE:NAME 格式的股票参数"""
    parts = s.split(":", 1)
    return {"code": parts[0], "name": parts[1] if len(parts) == 2 else ""}


def cmd_agent(args):
    config = load_config()
    if not config:
        print(
            "[错误] 未找到配置，请先运行: python alphapai_client.py config --set-key YOUR_KEY",
            file=sys.stderr,
        )
        sys.exit(1)

    client = AlphaPaiClient(config)
    stock = _parse_stock(args.stock) if args.stock else None
    stock_list = [_parse_stock(s) for s in args.stock_list] if args.stock_list else None

    # modes 5/8/9 have template fixed at 1 per API spec
    template = 1 if args.mode in (5, 8, 9) else args.template

    result = client.stock_agent(
        question=args.question,
        agent_mode=args.mode,
        stock=stock,
        template=template,
        template_text=args.template_text or "",
        report_type=args.report_type,
        stock_report_id=args.report_id,
        stock_report_title=args.report_title,
        stock_report_period=args.report_period,
        template_concern=args.concern,
        request_select_start_time=args.start,
        request_select_end_time=args.end,
        input_industry=args.industry,
        report_date=args.report_date,
        fund_type=args.fund_type,
        if_annual=args.if_annual,
        stock_list=stock_list,
        picture_color=args.picture_color,
        picture_style=args.picture_style,
        source=args.source,
        language=args.language,
        only_answer=True if args.only_answer else None,
    )

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(format_agent(result))


def cmd_report(args):
    config = load_config()
    if not config:
        print(
            "[错误] 未找到配置，请先运行: python alphapai_client.py config --set-key YOUR_KEY",
            file=sys.stderr,
        )
        sys.exit(1)

    client = AlphaPaiClient(config)
    result = client.stock_report(code=args.code)

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(format_report(result))


def cmd_image(args):
    config = load_config()
    if not config:
        print(
            "[错误] 未找到配置，请先运行: python alphapai_client.py config --set-key YOUR_KEY",
            file=sys.stderr,
        )
        sys.exit(1)

    client = AlphaPaiClient(config)
    result = client.search_image(
        query_text=args.query,
        files_range=args.files_range,
        topk=args.topk,
        recall_mode=args.recall_mode,
        use_llm_rank=args.llm_rank,
        start_date=args.start,
        end_date=args.end,
    )

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        print(format_image(result))


# ============================================================
# CLI 入口
# ============================================================


def main():
    parser = argparse.ArgumentParser(
        prog="alphapai_client.py",
        description="AlphaPai Research CLI — Alpha派投研API命令行工具",
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # --- config ---
    p_cfg = sub.add_parser("config", help="查看或设置API配置")
    p_cfg.add_argument("--show", action="store_true", help="显示当前配置")
    p_cfg.add_argument("--set-key", metavar="KEY", help="设置api_key")
    p_cfg.add_argument(
        "--set-url",
        metavar="URL",
        help="设置base_url（默认 https://open-api.rabyte.cn）",
    )

    # --- qa ---
    p_qa = sub.add_parser("qa", help="投研知识问答")
    p_qa.add_argument("--question", "-q", required=True, help="问题内容")
    p_qa.add_argument(
        "--mode",
        default="Flash",
        choices=["Flash", "Think"],
        help="问答模式: Flash=简单搜索问答，一问一搜一答（默认）; Think=Wide Search，一问多搜一答",
    )
    p_qa.add_argument(
        "--context",
        "-c",
        nargs="+",
        metavar="MSG",
        help="多轮对话上下文，按顺序传入历史消息列表",
    )
    p_qa.add_argument("--web-search", action="store_true", help="开启联网搜索")
    p_qa.add_argument("--deep-reasoning", action="store_true", help="开启深度推理")
    p_qa.add_argument("--start", metavar="YYYY-MM-DD", help="数据筛选开始日期")
    p_qa.add_argument("--end", metavar="YYYY-MM-DD", help="数据筛选结束日期")
    p_qa.add_argument("--json", action="store_true", help="输出原始JSON（供程序解析）")

    # --- recall ---
    p_rc = sub.add_parser("recall", help="获取问题召回的底层数据")
    p_rc.add_argument("--query", "-q", required=True, help="查询问题")
    p_rc.add_argument(
        "--type",
        "-t",
        metavar="TYPES",
        help="数据类型，逗号分隔，如 comment,qa,report（不传则全类型）",
    )
    p_rc.add_argument(
        "--no-cutoff",
        action="store_true",
        help="返回截断前完整内容（默认截断，与送入大模型的数据一致）",
    )
    p_rc.add_argument("--start", metavar="YYYY-MM-DD", help="数据筛选开始日期")
    p_rc.add_argument("--end", metavar="YYYY-MM-DD", help="数据筛选结束日期")
    p_rc.add_argument("--json", action="store_true", help="输出原始JSON（供程序解析）")

    # --- agent ---
    p_ag = sub.add_parser(
        "agent", help="股票 Agent（公司一页纸、业绩点评、调研大纲等）"
    )
    p_ag.add_argument(
        "--mode",
        "-m",
        type=int,
        required=True,
        help="Agent模式: 1=业绩点评 2=公司一页纸 3=调研大纲 5=主题选股 "
        "7=投资逻辑 8=可比公司 9=观点Challenge 11=行业一页纸 "
        "12=个股选基 13=主题选基 15=画图",
    )
    p_ag.add_argument("--question", "-q", required=True, help="问题内容")
    p_ag.add_argument(
        "--stock",
        metavar="CODE:NAME",
        help="股票信息，格式 CODE:NAME，如 300014.SZ:亿纬锂能",
    )
    p_ag.add_argument(
        "--template",
        type=int,
        default=0,
        help="模版类型: 0=alpha派模板(默认) 1=用户模板",
    )
    p_ag.add_argument(
        "--template-text", metavar="TEXT", help="用户模版正文（template=1时需传）"
    )
    p_ag.add_argument(
        "--report-type", metavar="TYPE", help="报告类型（业绩点评必填），如: 季报"
    )
    p_ag.add_argument("--report-id", metavar="ID", help="公告ID（业绩点评必填）")
    p_ag.add_argument(
        "--report-title", metavar="TITLE", help="报告标题（业绩点评必填）"
    )
    p_ag.add_argument(
        "--report-period",
        metavar="PERIOD",
        help="报告期（业绩点评必填），如: 2025年一季报",
    )
    p_ag.add_argument(
        "--concern", metavar="TEXT", help="用户关注内容（业绩点评/观点Challenge选填）"
    )
    p_ag.add_argument(
        "--industry", metavar="NAME", help="行业信息（行业一页纸必填），如: 白酒"
    )
    p_ag.add_argument(
        "--report-date",
        metavar="DATE",
        help="报告期（个股选基/主题选基必填），如: 2025-09-30",
    )
    p_ag.add_argument(
        "--fund-type",
        metavar="TYPE",
        help="基金类型（个股选基/主题选基必填）: 全部|主动|指数|ETF",
    )
    p_ag.add_argument(
        "--if-annual", type=int, choices=[0, 1], help="是否年报: 0=否 1=是"
    )
    p_ag.add_argument(
        "--stock-list",
        nargs="+",
        metavar="CODE:NAME",
        help="股票列表（个股选基必填），如: 601231.SH:环旭电子 300308.SZ:中际旭创",
    )
    p_ag.add_argument(
        "--picture-color",
        nargs="+",
        metavar="HEX",
        help="图片颜色HEX值（画图必填），如: 2A66F6 A5A8AF",
    )
    p_ag.add_argument(
        "--picture-style",
        metavar="STYLE",
        help="图片风格（画图必填）: PPT风格|科普风格",
    )
    p_ag.add_argument(
        "--source",
        type=int,
        choices=[0, 1],
        help="画图样式: 0=仅图片(默认) 1=图文",
    )
    p_ag.add_argument(
        "--language", metavar="LANG", help="语言（美股公司一页纸可选）: 中文|英文"
    )
    p_ag.add_argument(
        "--only-answer",
        action="store_true",
        help="仅返回最终答案（mode 7 投资逻辑可选）",
    )
    p_ag.add_argument("--start", metavar="YYYY-MM-DD", help="数据开始日期")
    p_ag.add_argument("--end", metavar="YYYY-MM-DD", help="数据结束日期")
    p_ag.add_argument("--json", action="store_true", help="输出原始JSON（供程序解析）")

    # --- report ---
    p_rp = sub.add_parser("report", help="获取股票公告列表（供业绩点评查询公告ID）")
    p_rp.add_argument("--code", "-c", required=True, help="股票编码，如 603380.SH")
    p_rp.add_argument("--json", action="store_true", help="输出原始JSON（供程序解析）")

    # --- image ---
    p_img = sub.add_parser("image", help="搜图表（从研报/公告中搜索图片和表格）")
    p_img.add_argument("--query", "-q", required=True, help="搜索内容")
    p_img.add_argument(
        "--files-range",
        nargs="+",
        metavar="CODE",
        help="来源类型代码列表，可选值: 3=内资研报 8=外资研报 6=公告 9=三方研报",
    )
    p_img.add_argument("--topk", type=int, default=50, help="返回数量(1-100，默认50)")
    p_img.add_argument(
        "--recall-mode",
        default="both",
        choices=["both", "vector_only", "es_only"],
        help="召回模式（默认both）",
    )
    p_img.add_argument("--llm-rank", action="store_true", help="使用LLM重排序")
    p_img.add_argument("--start", metavar="YYYY-MM-DD", help="开始日期")
    p_img.add_argument("--end", metavar="YYYY-MM-DD", help="结束日期")
    p_img.add_argument("--json", action="store_true", help="输出原始JSON（供程序解析）")

    args = parser.parse_args()
    {
        "config": cmd_config,
        "qa": cmd_qa,
        "recall": cmd_recall,
        "agent": cmd_agent,
        "report": cmd_report,
        "image": cmd_image,
    }[args.command](args)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
AlphaPai Web Client
===================
Alpha派 Web版 API 客户端（账号密码认证）。

用于获取路演纪要、点评、热门话题等功能。
"""

import base64
import json
import os
import time
import urllib.request
import urllib.parse
from pathlib import Path
from typing import Dict, List, Optional, Any
from html.parser import HTMLParser

BASE_URL = "https://alphapai-web.rabyte.cn/external/alpha/api"


def _strip_html(html: str) -> str:
    """去除 HTML 标签"""
    class _Strip(HTMLParser):
        def __init__(self):
            super().__init__()
            self.parts = []
        def handle_data(self, data):
            self.parts.append(data)
    s = _Strip()
    s.feed(html or "")
    return "".join(s.parts)


def _decode_jwt_exp(token: str) -> float:
    """从 JWT token 中提取 exp 时间戳。

    解析 JWT 第二段（payload）的 base64 编码，提取 exp 字段。
    返回 0 表示无法解析（不阻塞后续流程，由 API 401 兜底）。
    """
    try:
        payload = token.split(".")[1]
        # 补齐 base64 padding
        payload += "=" * (4 - len(payload) % 4)
        data = json.loads(base64.b64decode(payload))
        return float(data.get("exp", 0))
    except Exception:
        return 0


class AlphaPaiWebClient:
    """Alpha派 Web版 API 客户端"""

    def __init__(self, account: str, password: str, token_cache_path: str = None):
        self.account = account
        self.password = password
        self.token_cache_path = token_cache_path or os.path.join(
            os.path.dirname(__file__), "..", "auth_state.json"
        )
        self.token = None
        self._load_token()

    def _load_token(self):
        """从缓存加载 token，兼容两种格式，检查 JWT 过期时间。

        格式1（简化）: {"token": "...", "account": "..."}
        格式2（origins）: {"origins": [{"localStorage": [{"name": "USER_AUTH_TOKEN", "value": "..."}]}]}
        """
        if not os.path.exists(self.token_cache_path):
            return
        try:
            with open(self.token_cache_path, "r") as f:
                data = json.load(f)

            # 优先读 token 字段
            token = data.get("token")

            # fallback: origins 格式
            if not token:
                origins = data.get("origins", [])
                if origins:
                    for item in origins[0].get("localStorage", []):
                        if item.get("name") == "USER_AUTH_TOKEN":
                            token = item.get("value")
                            break

            if not token:
                return

            # 检查 JWT 过期时间，提前 60 秒判定过期
            exp = _decode_jwt_exp(token)
            if exp > 0 and time.time() > exp - 60:
                import sys
                print("[auth] 缓存 token 已过期，需重新登录", file=sys.stderr)
                return

            self.token = token
        except Exception:
            pass

    def _save_token(self):
        """保存 token 到缓存，同时写入两种格式保持兼容。"""
        os.makedirs(os.path.dirname(self.token_cache_path), exist_ok=True)
        auth = {
            "origin": "https://alphapai-web.rabyte.cn",
            "token": self.token,
            "account": self.account,
            "origins": [{
                "origin": "https://alphapai-web.rabyte.cn",
                "localStorage": [
                    {"name": "USER_AUTH_TOKEN", "value": self.token},
                ],
            }],
        }
        with open(self.token_cache_path, "w") as f:
            json.dump(auth, f, indent=2, ensure_ascii=False)
    
    def _request(self, method: str, endpoint: str, data: dict = None, auth: bool = True) -> dict:
        """发送请求"""
        url = f"{BASE_URL}/{endpoint}"
        headers = {"Content-Type": "application/json"}
        
        if auth and self.token:
            headers["Authorization"] = self.token
            headers["x-from"] = "web"
        
        body = None
        if data:
            body = json.dumps(data, ensure_ascii=False).encode("utf-8")
        
        req = urllib.request.Request(url, data=body, headers=headers, method=method)
        
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as e:
            if e.code == 401:
                # Token 过期，重新登录
                self._login()
                headers["Authorization"] = self.token
                req = urllib.request.Request(url, data=body, headers=headers, method=method)
                with urllib.request.urlopen(req, timeout=30) as resp:
                    return json.loads(resp.read().decode("utf-8"))
            raise
    
    def _login(self):
        """登录获取 token"""
        data = {"mobile": self.account, "password": self.password}
        result = self._request("POST", "v2/authentication/accountLogin", data, auth=False)
        
        if result.get("code") == 200000:
            self.token = result.get("data", {}).get("token")
            self._save_token()
        else:
            raise Exception(f"登录失败: {result.get('message', '未知错误')}")
    
    def ensure_login(self):
        """确保已登录"""
        if not self.token:
            self._login()
    
    # ========== 路演纪要 ==========
    
    def get_roadshow_list(self, page: int = 1, page_size: int = 20) -> dict:
        """获取路演纪要列表"""
        self.ensure_login()
        return self._request("POST", "reading/roadshow/summary/list", {
            "pageNum": page, "pageSize": page_size
        })
    
    def get_roadshow_detail(self, roadshow_id: str) -> dict:
        """获取路演纪要详情（含AI摘要）"""
        self.ensure_login()
        result = self._request("GET", f"reading/roadshow/summary/detail?id={roadshow_id}")
        return result.get("data", {})
    
    def search_roadshow(self, keyword: str, page: int = 1, page_size: int = 20) -> dict:
        """搜索路演纪要"""
        self.ensure_login()
        return self._request("POST", "reading/roadshow/summary/search", {
            "keyword": keyword, "pageNum": page, "pageSize": page_size
        })
    
    def get_hot_roadshow(self) -> dict:
        """获取热门路演推荐"""
        self.ensure_login()
        result = self._request("GET", "reading/roadshow/summary/hot/recommend")
        return result.get("data", {})
    
    # ========== 点评 ==========
    
    def search_comment(self, keyword: str, page: int = 1, page_size: int = 20) -> dict:
        """搜索点评"""
        self.ensure_login()
        return self._request("POST", "reading/comment/search", {
            "word": keyword, "pageNum": page, "pageSize": page_size
        })
    
    def get_comment_detail(self, comment_id: str) -> dict:
        """获取点评详情"""
        self.ensure_login()
        result = self._request("GET", f"reading/comment/detail?commentId={comment_id}")
        return result.get("data", {})
    
    def get_comment_recommend(self) -> dict:
        """获取推荐点评"""
        self.ensure_login()
        result = self._request("GET", "reading/comment/list/recommend")
        return result.get("data", {})
    
    # ========== 热门话题 ==========
    
    def get_hot_topics(self) -> List[dict]:
        """获取当前热门话题列表"""
        self.ensure_login()
        result = self._request("GET", "mix/hot/topic/current/batch/list")
        return result.get("data", [])
    
    def get_hot_topic_detail(self, topic_id: str) -> dict:
        """获取话题详情"""
        self.ensure_login()
        result = self._request("GET", f"mix/hot/topic/detail?id={topic_id}")
        return result.get("data", {})
    
    def get_hot_topic_agent(self, topic_id: str) -> dict:
        """获取话题AI分析"""
        self.ensure_login()
        result = self._request("GET", f"mix/hot/topic/agent/detail?id={topic_id}")
        return result.get("data", {})
    
    def get_hot_words(self) -> List[str]:
        """获取热搜词"""
        self.ensure_login()
        result = self._request("GET", "reading/hot/word")
        return result.get("data", [])
    
    def get_hot_report(self) -> dict:
        """获取最新热门报告"""
        self.ensure_login()
        result = self._request("GET", "mix/hot/topic/report/latest")
        return result.get("data", {})
    
    def get_hot_stocks(self) -> dict:
        """获取热门股票（机构持仓）"""
        self.ensure_login()
        result = self._request("GET", "mix/hot/topic/stock/list")
        return result.get("data", {})
    
    # ========== 搜索 ==========
    
    def smart_search(self, keyword: str) -> dict:
        """智能搜索（综合）"""
        self.ensure_login()
        result = self._request("POST", "reading/smartSearch", {"word": keyword})
        return result.get("data", {})
    
    def search_stock(self, keyword: str) -> dict:
        """搜索股票"""
        self.ensure_login()
        result = self._request("GET", f"reading/stock/search?keyword={urllib.parse.quote(keyword)}")
        return result.get("data", {})
    
    # ========== 其他 ==========
    
    def get_user_info(self) -> dict:
        """获取用户信息"""
        self.ensure_login()
        result = self._request("GET", "v2/authorization/user/info")
        return result.get("data", {})
    
    def get_today_count(self) -> dict:
        """获取今日内容统计"""
        self.ensure_login()
        result = self._request("GET", "reading/count/today")
        return result.get("data", {})


# 命令行入口（兼容旧用法）
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="AlphaPai Web Client")
    parser.add_argument("command", help="命令")
    parser.add_argument("args", nargs="*", help="参数")
    parser.add_argument("--detail", action="store_true", help="获取详情")
    
    args = parser.parse_args()
    
    # 从环境变量或 .env 读取账号密码
    account = os.environ.get("ALPHAPAI_ACCOUNT")
    password = os.environ.get("ALPHAPAI_PASSWORD")
    
    if not account or not password:
        # 尝试从 .env 读取
        env_path = Path(__file__).parent.parent / ".env"
        if env_path.exists():
            for line in env_path.read_text().splitlines():
                if line.startswith("ALPHAPAI_ACCOUNT="):
                    account = line.split("=", 1)[1].strip().strip('"')
                elif line.startswith("ALPHAPAI_PASSWORD="):
                    password = line.split("=", 1)[1].strip().strip('"')
    
    if not account or not password:
        print("错误: 未配置账号密码，请设置 ALPHAPAI_ACCOUNT 和 ALPHAPAI_PASSWORD")
        exit(1)
    
    client = AlphaPaiWebClient(account, password)
    
    if args.command == "topics":
        topics = client.get_hot_topics()
        for t in topics or []:
            name = (t.get("topicName") or "")[:30]
            summary = (t.get("summary") or "")[:50]
            print(f"{name:30} | {summary}")
    
    elif args.command == "topic":
        if not args.args:
            print("用法: alphapai_web_client.py topic <id>")
        else:
            detail = client.get_hot_topic_detail(args.args[0])
            print(json.dumps(detail, ensure_ascii=False, indent=2))
    
    elif args.command == "roadshow":
        result = client.get_roadshow_list()
        for r in result.get("list", []):
            print(f"{r.get('roadshowDate', ''):12} | {r.get('title', '')[:60]}")
    
    elif args.command == "roadshow-search":
        if not args.args:
            print("用法: alphapai_web_client.py roadshow-search <关键词>")
        else:
            result = client.search_roadshow(args.args[0])
            for r in result.get("list", []):
                print(f"{r.get('roadshowDate', ''):12} | {r.get('title', '')[:60]}")
    
    elif args.command == "roadshow-detail":
        if not args.args:
            print("用法: alphapai_web_client.py roadshow-detail <id>")
        else:
            detail = client.get_roadshow_detail(args.args[0])
            print(f"# {detail.get('title', '')}")
            if detail.get("aiSummaryV3"):
                print(detail["aiSummaryV3"].get("fullTextSummary", ""))
    
    elif args.command == "comment-search":
        if not args.args:
            print("用法: alphapai_web_client.py comment-search <关键词>")
        else:
            result = client.search_comment(args.args[0])
            for r in result.get("list", []):
                print(f"{r.get('time', ''):12} | {r.get('title', '')[:60]}")
    
    elif args.command == "comment-detail":
        if not args.args:
            print("用法: alphapai_web_client.py comment-detail <id>")
        else:
            detail = client.get_comment_detail(args.args[0])
            print(f"# {detail.get('title', '')}")
            print(detail.get("content", "")[:2000])
    
    elif args.command == "hot-word":
        words = client.get_hot_words()
        for i, w in enumerate(words, 1):
            print(f"{i:2}. {w}")
    
    else:
        print(f"未知命令: {args.command}")
        print("可用命令: topics, topic, roadshow, roadshow-search, roadshow-detail, comment-search, comment-detail, hot-word")

#!/usr/bin/env python3
"""
AlphaPai Unified CLI
====================
Alpha派统一命令行工具，融合官方API和Web版功能。

用法：
  python alphapai.py qa "茅台经营情况"
  python alphapai.py recall "新能源" --type report
  python alphapai.py company 300014.SZ
  python alphapai.py roadshow --search "茅台业绩会"
  python alphapai.py hot

认证优先级：
  1. 环境变量 ALPHAPAI_API_KEY
  2. 同级/上级/上上级目录 .env
  3. config.json
  4. 账号密码（ALPHAPAI_ACCOUNT/ALPHAPAI_PASSWORD）
"""

import argparse
import json
import os
import sys
from typing import Any, Dict, List, Optional

# 添加脚本目录到路径
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, SCRIPT_DIR)

# ============================================================
# 配置管理（融合版）
# ============================================================

DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(SCRIPT_DIR), "config.json")


def _parse_env_file(env_path: str) -> Dict[str, str]:
    """解析 .env 文件"""
    result = {}
    if not os.path.exists(env_path):
        return result
    try:
        with open(env_path, "r", encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith("#"):
                    continue
                if "=" in line:
                    key, _, value = line.partition("=")
                    key = key.strip()
                    value = value.strip().strip('"').strip("'")
                    if key:
                        result[key] = value
    except Exception:
        pass
    return result


def _find_env_value(key: str, start_dir: str) -> Optional[str]:
    """从当前目录向上逐级查找 .env 文件中的 key"""
    max_levels = 3
    current_dir = start_dir
    for _ in range(max_levels):
        env_path = os.path.join(current_dir, ".env")
        env_dict = _parse_env_file(env_path)
        if key in env_dict and env_dict[key]:
            return env_dict[key]
        parent = os.path.dirname(current_dir)
        if parent == current_dir:
            break
        current_dir = parent
    return None


def load_config() -> Dict:
    """加载配置，返回包含所有认证信息的字典"""
    skill_dir = os.path.dirname(SCRIPT_DIR)
    
    # 从环境变量或.env读取
    api_key = os.environ.get("ALPHAPAI_API_KEY")
    if not api_key:
        api_key = _find_env_value("ALPHAPAI_API_KEY", skill_dir)
    
    base_url = os.environ.get("ALPHAPAI_BASE_URL")
    if not base_url:
        base_url = _find_env_value("ALPHAPAI_BASE_URL", skill_dir)
    if not base_url:
        base_url = "https://open-api.rabyte.cn"
    
    # 账号密码（Web版）
    account = os.environ.get("ALPHAPAI_ACCOUNT")
    password = os.environ.get("ALPHAPAI_PASSWORD")
    if not account:
        account = _find_env_value("ALPHAPAI_ACCOUNT", skill_dir)
    if not password:
        password = _find_env_value("ALPHAPAI_PASSWORD", skill_dir)
    
    # 从 config.json 读取
    config_path = DEFAULT_CONFIG_PATH
    if os.path.exists(config_path):
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
            if not api_key and config.get("api_key"):
                api_key = config["api_key"]
            if not base_url and config.get("base_url"):
                base_url = config["base_url"]
            if not account and config.get("account"):
                account = config["account"]
            if not password and config.get("password"):
                password = config["password"]
        except Exception:
            pass
    
    return {
        "api_key": api_key,
        "base_url": base_url,
        "account": account,
        "password": password,
        "config_path": config_path,
    }


def save_config(api_key: str = None, base_url: str = None, account: str = None, password: str = None):
    """保存配置到 config.json"""
    config = {}
    config_path = DEFAULT_CONFIG_PATH
    
    # 读取现有配置
    if os.path.exists(config_path):
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                config = json.load(f)
        except Exception:
            pass
    
    # 更新
    if api_key is not None:
        config["api_key"] = api_key
    if base_url is not None:
        config["base_url"] = base_url
    if account is not None:
        config["account"] = account
    if password is not None:
        config["password"] = password
    
    # 写入
    os.makedirs(os.path.dirname(config_path), exist_ok=True)
    with open(config_path, "w", encoding="utf-8") as f:
        json.dump(config, f, indent=2, ensure_ascii=False)


# ============================================================
# 命令处理
# ============================================================

def cmd_config(args):
    """配置管理"""
    if args.set_key:
        save_config(api_key=args.set_key)
        print(f"API Key 已保存")
    elif args.set_account and args.set_password:
        save_config(account=args.set_account, password=args.set_password)
        print(f"账号密码已保存")
    elif args.show:
        config = load_config()
        print(f"API Key  : {config['api_key'][:8] + '****' + config['api_key'][-4:] if config['api_key'] else '未配置'}")
        print(f"Base URL : {config['base_url']}")
        print(f"Account  : {config['account'] or '未配置'}")
        print(f"Password : {'****' if config['password'] else '未配置'}")
        print(f"配置路径 : {config['config_path']}")
    else:
        print("用法: alphapai config --show  或  --set-key KEY  或  --set-account PHONE --set-password PWD")


def cmd_qa(args):
    """投研问答（官方API）"""
    config = load_config()
    if not config["api_key"]:
        print("错误: 未配置 API Key，请运行: alphapai config --set-key YOUR_KEY")
        return
    
    from alphapai_client import AlphaPaiClient
    client = AlphaPaiClient({"api_key": config["api_key"], "base_url": config["base_url"]})
    
    mode = "Think" if args.think else "Flash"
    result = client.qa_text(
        question=args.question,
        mode=mode,
        context=args.context or None,
        web_search=args.web_search,
        deep_reasoning=args.deep_reasoning,
        start_date=args.start,
        end_date=args.end,
    )
    
    if result.get("answer"):
        print(result["answer"])
    else:
        print("未获取到回答")


def cmd_recall(args):
    """数据召回（官方API）"""
    config = load_config()
    if not config["api_key"]:
        print("错误: 未配置 API Key")
        return
    
    from alphapai_client import AlphaPaiClient
    client = AlphaPaiClient({"api_key": config["api_key"], "base_url": config["base_url"]})
    
    recall_types = args.type.split(",") if args.type else None
    result = client.recall_data(
        query=args.query,
        recall_type=recall_types,
        cutoff=not args.no_cutoff,
        start_time=args.start,
        end_time=args.end,
    )
    
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        for item in result.get("data", []):
            print(f"\n### {item.get('type', 'unknown')} | {item.get('contextInfo', '')[:50]}...")
            for chunk in item.get("chunks", [])[:3]:
                print(f"  - {chunk[:200]}...")


def cmd_agent(args):
    """股票Agent（官方API）"""
    config = load_config()
    if not config["api_key"]:
        print("错误: 未配置 API Key")
        return
    
    from alphapai_client import AlphaPaiClient
    client = AlphaPaiClient({"api_key": config["api_key"], "base_url": config["base_url"]})
    
    # 解析股票参数
    stock = None
    if args.stock:
        parts = args.stock.split(":")
        if len(parts) == 2:
            stock = {"code": parts[0], "name": parts[1]}
        else:
            stock = {"code": parts[0], "name": parts[0]}
    
    # 解析股票列表
    stock_list = None
    if getattr(args, 'stock_list', None):
        stock_list = []
        for s in args.stock_list:
            parts = s.split(":")
            if len(parts) == 2:
                stock_list.append({"code": parts[0], "name": parts[1]})
            else:
                stock_list.append({"code": parts[0], "name": parts[0]})
    
    result = client.stock_agent(
        question=args.question,
        agent_mode=args.mode,
        stock=stock,
        input_industry=getattr(args, 'input_industry', None),
        stock_list=stock_list,
        report_date=getattr(args, 'report_date', None),
        fund_type=getattr(args, 'fund_type', None),
        template_text=getattr(args, 'template_text', None) or "",
        template_concern=getattr(args, 'concern', None),
        report_type=getattr(args, 'report_type', None),
        stock_report_id=getattr(args, 'report_id', None),
        stock_report_title=getattr(args, 'report_title', None),
        stock_report_period=getattr(args, 'report_period', None),
        picture_color=getattr(args, 'picture_color', None),
        picture_style=getattr(args, 'picture_style', None),
    )
    
    if result.get("answer"):
        print(result["answer"])


def cmd_report(args):
    """公告列表（官方API）"""
    config = load_config()
    if not config["api_key"]:
        print("错误: 未配置 API Key")
        return
    
    from alphapai_client import AlphaPaiClient
    client = AlphaPaiClient({"api_key": config["api_key"], "base_url": config["base_url"]})
    
    result = client.get_stock_reports(args.code)
    
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        for r in result.get("data", []):
            print(f"{r.get('reportPeriod', ''):15} | {r.get('reportType', ''):6} | {r.get('stockReportTitle', '')[:50]}")


def cmd_image(args):
    """搜图表（官方API）"""
    config = load_config()
    if not config["api_key"]:
        print("错误: 未配置 API Key")
        return
    
    from alphapai_client import AlphaPaiClient
    client = AlphaPaiClient({"api_key": config["api_key"], "base_url": config["base_url"]})
    
    files_range = args.files_range.split(",") if args.files_range else None
    result = client.search_image(
        query=args.query,
        files_range=files_range,
        topk=args.topk,
        start_date=args.start,
        end_date=args.end,
    )
    
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        for img in result.get("data", [])[:10]:
            print(f"{img.get('articleTitle', '')[:40]} | {img.get('imageUrl', '')[:60]}")


# ========== Web版命令 ==========

def _get_web_client():
    """获取Web客户端"""
    config = load_config()
    if not config["account"] or not config["password"]:
        print("错误: 未配置账号密码，请运行: alphapai config --set-account PHONE --set-password PWD")
        return None
    
    from alphapai_web_client import AlphaPaiWebClient
    return AlphaPaiWebClient(config["account"], config["password"])


def cmd_roadshow(args):
    """路演纪要（Web版）"""
    client = _get_web_client()
    if not client:
        return
    
    if args.detail:
        result = client.get_roadshow_detail(args.detail)
        if result:
            print(f"\n# {result.get('title', '')}")
            if result.get("aiSummaryV3"):
                print(result["aiSummaryV3"].get("fullTextSummary", ""))
    elif args.search:
        result = client.search_roadshow(args.search, page=args.page, page_size=args.page_size)
        for r in result.get("list", []):
            print(f"{r.get('roadshowDate', ''):12} | {r.get('title', '')[:60]}")
    else:
        result = client.get_roadshow_list(page=args.page, page_size=args.page_size)
        for r in result.get("list", []):
            print(f"{r.get('roadshowDate', ''):12} | {r.get('title', '')[:60]}")


def cmd_comment(args):
    """点评（Web版）"""
    client = _get_web_client()
    if not client:
        return
    
    if args.detail:
        result = client.get_comment_detail(args.detail)
        if result:
            print(f"\n# {result.get('title', '')}")
            print(result.get("content", "")[:2000])
    elif args.search:
        result = client.search_comment(args.search, page=args.page, page_size=args.page_size)
        for r in result.get("list", []):
            print(f"{r.get('time', ''):12} | {r.get('title', '')[:60]}")
    else:
        result = client.get_comment_recommend()
        for r in result.get("list", []):
            print(f"{r.get('time', ''):12} | {r.get('title', '')[:60]}")


def cmd_hot(args):
    """热门话题（Web版）"""
    client = _get_web_client()
    if not client:
        return
    
    if args.detail:
        result = client.get_hot_topic_detail(args.detail)
        if result:
            print(f"\n# {result.get('topicName', '')}")
            print(result.get("summary", "") or "")
    else:
        result = client.get_hot_topics()
        for t in result or []:
            new_flag = "🆕" if t.get("newFlag") else "  "
            name = (t.get("topicName") or "")[:30]
            summary = (t.get("summary") or "")[:50]
            print(f"{new_flag} {name} | {summary}")


def cmd_hot_word(args):
    """热词（Web版）"""
    client = _get_web_client()
    if not client:
        return

    result = client.get_hot_words()
    for i, w in enumerate(result or [], 1):
        # w 可能是字符串或字典
        if isinstance(w, dict):
            print(f"{i:2}. {w.get('word', w)}")
        else:
            print(f"{i:2}. {w}")


def cmd_hot_report(args):
    """热门报告（Web版）"""
    client = _get_web_client()
    if not client:
        return

    from html.parser import HTMLParser

    def _strip(html: str) -> str:
        class _S(HTMLParser):
            def __init__(self):
                super().__init__()
                self.parts = []
            def handle_data(self, d):
                self.parts.append(d)
        s = _S()
        s.feed(html or "")
        return "".join(s.parts)

    result = client.get_hot_report()
    if result:
        print(f"# {result.get('title', '')}")
        print(f"日期: {result.get('date', '')}")
        print()
        print(_strip(result.get("content", "")))


def cmd_hot_stock(args):
    """热门股票（Web版）"""
    client = _get_web_client()
    if not client:
        return

    result = client.get_hot_stocks()
    if not result:
        print("无数据")
        return

    print(f"更新时间: {result.get('updateTime', '')}\n")
    for key, label in [("publicList", "公募"), ("privateList", "私募"), ("insuranceList", "险资")]:
        items = result.get(key, [])
        if items:
            print(f"## {label} ({len(items)} 只)")
            for item in items[:20]:
                if isinstance(item, dict):
                    name = item.get("stockName", item.get("name", ""))
                    code = item.get("stockCode", item.get("code", ""))
                    print(f"  - {name} ({code})")
            print()


def cmd_search(args):
    """综合搜索（Web版）"""
    client = _get_web_client()
    if not client:
        return

    result = client.smart_search(args.keyword)
    if not result:
        print("无结果")
        return

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return

    # 路演纪要
    rs = result.get("roadshowSummary", {})
    if rs and rs.get("list"):
        print(f"## 路演纪要 (共 {rs.get('total', '?')} 条)")
        for item in rs["list"][:5]:
            print(f"  - {item.get('title', '')[:60]} | {item.get('publishInstitution', '')}")
        print()

    # 股票
    stock = result.get("stock", {})
    if stock and stock.get("list"):
        print(f"## 相关股票")
        for item in stock["list"][:5]:
            if isinstance(item, dict):
                print(f"  - {item.get('name', '')} ({item.get('code', '')})")
        print()


def cmd_user_info(args):
    """用户信息（Web版）"""
    client = _get_web_client()
    if not client:
        return

    result = client.get_user_info()
    print(json.dumps(result, ensure_ascii=False, indent=2))


# ============================================================
# 快捷命令
# ============================================================

def cmd_company(args):
    """公司一页纸（快捷命令）"""
    args.mode = 2
    args.question = f"公司一页纸: {args.code}"
    args.stock = args.code
    cmd_agent(args)


def cmd_earnings(args):
    """业绩点评（快捷命令）"""
    config = load_config()
    if not config["api_key"]:
        print("错误: 未配置 API Key")
        return
    
    from alphapai_client import AlphaPaiClient
    client = AlphaPaiClient({"api_key": config["api_key"], "base_url": config["base_url"]})
    
    # 先查公告
    code = args.code.split(":")[0] if ":" in args.code else args.code
    reports = client.stock_report(code)
    
    if not reports.get("data"):
        print("未找到公告")
        return
    
    # 取最新一份
    latest = reports["data"][0]
    print(f"公告: {latest.get('stockReportTitle', '')}")
    print(f"期间: {latest.get('reportPeriod', '')}")
    print()
    
    # 调用业绩点评
    name = args.code.split(":")[1] if ":" in args.code else code
    result = client.stock_agent(
        question=f"{name}{latest.get('reportPeriod', '')}业绩点评",
        agent_mode=1,
        stock={"code": code, "name": name},
        report_type=latest.get("reportType", ""),
        stock_report_id=latest.get("stockReportId", ""),
        stock_report_title=latest.get("stockReportTitle", ""),
        stock_report_period=latest.get("reportPeriod", ""),
    )
    
    if result.get("answer"):
        print(result["answer"])


def cmd_logic(args):
    """投资逻辑（快捷命令）"""
    args.mode = 7
    args.question = f"{args.code} 投资逻辑"
    args.stock = args.code
    args.input_industry = None
    args.stock_list = None
    args.report_date = None
    args.fund_type = None
    args.template_text = None
    args.concern = None
    args.report_type = None
    args.report_id = None
    args.report_title = None
    args.report_period = None
    args.picture_color = None
    args.picture_style = None
    cmd_agent(args)


def cmd_peers(args):
    """可比公司（快捷命令）"""
    args.mode = 8
    args.question = f"{args.code} 可比公司"
    args.stock = args.code
    args.input_industry = None
    args.stock_list = None
    args.report_date = None
    args.fund_type = None
    args.template_text = None
    args.concern = None
    args.report_type = None
    args.report_id = None
    args.report_title = None
    args.report_period = None
    args.picture_color = None
    args.picture_style = None
    cmd_agent(args)


def cmd_industry(args):
    """行业一页纸（快捷命令）"""
    args.mode = 11
    args.question = f"{args.industry} 行业一页纸"
    args.stock = None
    args.input_industry = args.industry
    args.stock_list = None
    args.report_date = None
    args.fund_type = None
    args.template_text = None
    args.concern = None
    args.report_type = None
    args.report_id = None
    args.report_title = None
    args.report_period = None
    args.picture_color = None
    args.picture_style = None
    cmd_agent(args)


def cmd_screen(args):
    """主题选股（快捷命令）"""
    args.mode = 5
    args.question = args.theme
    args.stock = None
    args.input_industry = None
    args.stock_list = None
    args.report_date = None
    args.fund_type = None
    args.template_text = args.theme
    args.concern = None
    args.report_type = None
    args.report_id = None
    args.report_title = None
    args.report_period = None
    args.picture_color = None
    args.picture_style = None
    cmd_agent(args)


def cmd_challenge(args):
    """观点Challenge（快捷命令）"""
    args.mode = 9
    args.question = f"Challenge该观点: {args.viewpoint}"
    args.stock = None
    args.input_industry = None
    args.stock_list = None
    args.report_date = None
    args.fund_type = None
    args.template_text = args.viewpoint
    args.concern = None
    args.report_type = None
    args.report_id = None
    args.report_title = None
    args.report_period = None
    args.picture_color = None
    args.picture_style = None
    cmd_agent(args)


# ============================================================
# 主入口
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description="Alpha派统一CLI - 融合官方API和Web版功能",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    subparsers = parser.add_subparsers(dest="command", help="命令")
    
    # config
    p_config = subparsers.add_parser("config", help="配置管理")
    p_config.add_argument("--show", action="store_true", help="显示配置")
    p_config.add_argument("--set-key", metavar="KEY", help="设置API Key")
    p_config.add_argument("--set-account", metavar="PHONE", help="设置账号")
    p_config.add_argument("--set-password", metavar="PWD", help="设置密码")
    p_config.set_defaults(func=cmd_config)
    
    # qa
    p_qa = subparsers.add_parser("qa", help="投研问答")
    p_qa.add_argument("question", help="问题")
    p_qa.add_argument("--think", action="store_true", help="Think模式")
    p_qa.add_argument("--web-search", action="store_true", help="联网搜索")
    p_qa.add_argument("--deep-reasoning", action="store_true", help="深度推理")
    p_qa.add_argument("--context", nargs="+", help="多轮对话历史")
    p_qa.add_argument("--start", help="开始日期")
    p_qa.add_argument("--end", help="结束日期")
    p_qa.set_defaults(func=cmd_qa)
    
    # recall
    p_recall = subparsers.add_parser("recall", help="数据召回")
    p_recall.add_argument("query", help="查询")
    p_recall.add_argument("--type", help="数据类型(逗号分隔)")
    p_recall.add_argument("--no-cutoff", action="store_true", help="返回完整内容")
    p_recall.add_argument("--start", help="开始日期")
    p_recall.add_argument("--end", help="结束日期")
    p_recall.add_argument("--json", action="store_true", help="输出JSON")
    p_recall.set_defaults(func=cmd_recall)
    
    # agent
    p_agent = subparsers.add_parser("agent", help="股票Agent")
    p_agent.add_argument("--mode", type=int, required=True, help="模式(1-15)")
    p_agent.add_argument("--question", help="问题")
    p_agent.add_argument("--stock", help="股票 CODE:NAME")
    p_agent.add_argument("--industry", help="行业")
    p_agent.add_argument("--stock-list", nargs="+", help="股票列表")
    p_agent.add_argument("--report-date", help="报告日期")
    p_agent.add_argument("--fund-type", help="基金类型")
    p_agent.add_argument("--template-text", help="模板文本")
    p_agent.add_argument("--concern", help="关注点")
    p_agent.add_argument("--report-type", help="报告类型")
    p_agent.add_argument("--report-id", help="公告ID")
    p_agent.add_argument("--report-title", help="公告标题")
    p_agent.add_argument("--report-period", help="报告期")
    p_agent.add_argument("--picture-color", nargs=2, help="颜色HEX")
    p_agent.add_argument("--picture-style", help="画图风格")
    p_agent.set_defaults(func=cmd_agent)
    
    # report
    p_report = subparsers.add_parser("report", help="公告列表")
    p_report.add_argument("code", help="股票代码")
    p_report.add_argument("--json", action="store_true", help="输出JSON")
    p_report.set_defaults(func=cmd_report)
    
    # image / chart
    p_image = subparsers.add_parser("image", aliases=["chart"], help="搜图表")
    p_image.add_argument("query", help="查询")
    p_image.add_argument("--files-range", help="来源类型")
    p_image.add_argument("--topk", type=int, default=10, help="返回数量")
    p_image.add_argument("--start", help="开始日期")
    p_image.add_argument("--end", help="结束日期")
    p_image.add_argument("--json", action="store_true", help="输出JSON")
    p_image.set_defaults(func=cmd_image)
    
    # roadshow
    p_roadshow = subparsers.add_parser("roadshow", help="路演纪要")
    p_roadshow.add_argument("--list", action="store_true", help="列表")
    p_roadshow.add_argument("--search", help="搜索关键词")
    p_roadshow.add_argument("--detail", help="纪要ID")
    p_roadshow.add_argument("--page", type=int, default=1)
    p_roadshow.add_argument("--page-size", type=int, default=20)
    p_roadshow.set_defaults(func=cmd_roadshow)
    
    # comment
    p_comment = subparsers.add_parser("comment", help="点评")
    p_comment.add_argument("--list", action="store_true", help="推荐列表")
    p_comment.add_argument("--search", help="搜索关键词")
    p_comment.add_argument("--detail", help="点评ID")
    p_comment.add_argument("--page", type=int, default=1)
    p_comment.add_argument("--page-size", type=int, default=20)
    p_comment.set_defaults(func=cmd_comment)
    
    # hot
    p_hot = subparsers.add_parser("hot", help="热门话题")
    p_hot.add_argument("--detail", help="话题ID")
    p_hot.set_defaults(func=cmd_hot)
    
    # hot-word
    p_hw = subparsers.add_parser("hot-word", help="热词")
    p_hw.set_defaults(func=cmd_hot_word)

    # hot-report
    p_hr = subparsers.add_parser("hot-report", help="热门报告")
    p_hr.set_defaults(func=cmd_hot_report)

    # hot-stock
    p_hs = subparsers.add_parser("hot-stock", help="热门股票(机构持仓)")
    p_hs.set_defaults(func=cmd_hot_stock)

    # search (综合搜索)
    p_search = subparsers.add_parser("search", help="综合搜索")
    p_search.add_argument("keyword", help="搜索关键词")
    p_search.add_argument("--json", action="store_true", help="输出JSON")
    p_search.set_defaults(func=cmd_search)

    # user-info
    p_ui = subparsers.add_parser("user-info", help="用户信息")
    p_ui.set_defaults(func=cmd_user_info)

    # 快捷命令
    p_company = subparsers.add_parser("company", help="公司一页纸")
    p_company.add_argument("code", help="股票 CODE:NAME")
    p_company.set_defaults(func=cmd_company)
    
    p_earnings = subparsers.add_parser("earnings", help="业绩点评")
    p_earnings.add_argument("code", help="股票 CODE:NAME")
    p_earnings.set_defaults(func=cmd_earnings)
    
    p_logic = subparsers.add_parser("logic", help="投资逻辑")
    p_logic.add_argument("code", help="股票 CODE:NAME")
    p_logic.set_defaults(func=cmd_logic)
    
    p_peers = subparsers.add_parser("peers", help="可比公司")
    p_peers.add_argument("code", help="股票 CODE:NAME")
    p_peers.set_defaults(func=cmd_peers)
    
    p_ind = subparsers.add_parser("industry", help="行业一页纸")
    p_ind.add_argument("industry", help="行业名称")
    p_ind.set_defaults(func=cmd_industry)
    
    p_screen = subparsers.add_parser("screen", help="主题选股")
    p_screen.add_argument("theme", help="选股主题")
    p_screen.set_defaults(func=cmd_screen)
    
    p_challenge = subparsers.add_parser("challenge", help="观点Challenge")
    p_challenge.add_argument("viewpoint", help="待挑战观点")
    p_challenge.set_defaults(func=cmd_challenge)
    
    # 解析
    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        return

    # 启动时检查配置完整性，给出明确提示
    if args.command not in ("config",):
        config = load_config()
        has_api = bool(config.get("api_key"))
        has_web = bool(config.get("account") and config.get("password"))
        if not has_api and not has_web:
            print("警告: 未配置任何认证信息，所有命令将无法使用。", file=sys.stderr)
            print("  官方API: alphapai config --set-key YOUR_KEY", file=sys.stderr)
            print("  Web版:   alphapai config --set-account PHONE --set-password PWD", file=sys.stderr)
            print("  或设置环境变量: ALPHAPAI_API_KEY / ALPHAPAI_ACCOUNT + ALPHAPAI_PASSWORD", file=sys.stderr)
            sys.exit(1)
        # 针对 API 类命令但缺 API Key 的情况提前警告
        api_commands = {"qa", "recall", "agent", "report", "image", "chart",
                        "company", "earnings", "logic", "peers", "industry", "screen", "challenge"}
        web_commands = {"roadshow", "comment", "hot", "hot-word", "hot-report",
                        "hot-stock", "search", "user-info"}
        if args.command in api_commands and not has_api:
            print(f"错误: 命令 '{args.command}' 需要官方 API Key，但未配置。", file=sys.stderr)
            print("  运行: alphapai config --set-key YOUR_KEY", file=sys.stderr)
            sys.exit(1)
        if args.command in web_commands and not has_web:
            print(f"错误: 命令 '{args.command}' 需要账号密码，但未配置。", file=sys.stderr)
            print("  运行: alphapai config --set-account PHONE --set-password PWD", file=sys.stderr)
            sys.exit(1)

    if hasattr(args, "func"):
        args.func(args)
    else:
        parser.print_help()


if __name__ == "__main__":
   main()