Rqdata Python

#!/usr/bin/env python3
"""
API 索引管理器 - 快速定位 API 在文档中的位置
"""

from pathlib import Path
from typing import Dict, List, Optional
import re
import logging

# 配置日志（INFO级别）
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("APIIndexManager")


class APIIndexManager:
    """API 索引管理器"""

    def __init__(self, api_index_dir: Optional[str] = None):
        """
        初始化 API 索引管理器

        Args:
            api_index_dir: api_index 目录路径
        """
        if api_index_dir is None:
            skill_root = Path(__file__).parent.parent
            self.api_index_dir = skill_root / "cache" / "api_index"
        else:
            self.api_index_dir = Path(api_index_dir)

        logger.info(
            f"API Index Manager initialized with directory: {self.api_index_dir}"
        )

        # 内存缓存：{document_name: {api_name: line_info}}
        self.memory_cache: Dict[str, Dict[str, dict]] = {}

    def _parse_api_index_file(self, document_name: str) -> Dict[str, dict]:
        """解析 api_index 文件，构建 API 到行号的映射"""
        logger.info(f"Parsing API index for document: {document_name}")

        # 将文档名转换为索引文件名
        index_file_name = f"{document_name.replace('.md', '')}_index.md"
        index_file_path = self.api_index_dir / index_file_name

        if not index_file_path.exists():
            error_msg = f"API index file not found: {index_file_path}"
            logger.error(error_msg)
            raise FileNotFoundError(
                f"{error_msg}\nPlease run init_skill.py to generate API indices."
            )

        # 读取索引文件
        try:
            with open(index_file_path, "r", encoding="utf-8") as f:
                content = f.read()
        except Exception as e:
            error_msg = f"Error reading index file {index_file_path}: {e}"
            logger.error(error_msg)
            raise IOError(error_msg)

        # 解析表格中的 API 信息
        api_mapping = {}
        table_pattern = re.compile(r"\|\s*`([^`]+)`\s*\|\s*([^\|]+)\|\s*(\d+)\s*\|")

        for line in content.splitlines():
            match = table_pattern.match(line)
            if match:
                api_name = match.group(1).strip()
                description = match.group(2).strip()
                line_number = int(match.group(3).strip())

                api_mapping[api_name] = {
                    "line_number": line_number,
                    "description": description,
                    "document_name": document_name,
                }

        logger.info(f"Parsed {len(api_mapping)} APIs from {document_name}")
        return api_mapping

    def get_api_location(self, api_name: str, document_name: str) -> dict:
        """获取 API 在文档中的位置（单个 API）"""
        logger.info(
            f"Getting location for API '{api_name}' in document '{document_name}'"
        )

        # 检查内存缓存
        if document_name not in self.memory_cache:
            self.memory_cache[document_name] = self._parse_api_index_file(document_name)

        doc_cache = self.memory_cache[document_name]

        # 查找 API
        if api_name in doc_cache:
            api_info = doc_cache[api_name]
            logger.info(f"Found API '{api_name}' at line {api_info['line_number']}")
            return api_info
        else:
            # 获取可用 API 列表（前 10 个）
            available_apis = list(doc_cache.keys())[:10]
            error_msg = (
                f"API '{api_name}' not found in document index '{document_name}'.\n"
                f"Available APIs in this document: {available_apis}...\n"
                f"Please check the API name or use full document search."
            )
            logger.error(error_msg)
            raise ValueError(error_msg)

    def get_batch_api_locations(
        self, api_names: List[str], document_name: str
    ) -> Dict[str, dict]:
        """批量获取多个 API 的位置"""
        logger.info(
            f"Batch getting locations for {len(api_names)} APIs in document '{document_name}'"
        )

        # 确保文档索引已加载
        if document_name not in self.memory_cache:
            self.memory_cache[document_name] = self._parse_api_index_file(document_name)

        doc_cache = self.memory_cache[document_name]
        result = {}
        missing_apis = []

        for api_name in api_names:
            if api_name in doc_cache:
                result[api_name] = doc_cache[api_name]
            else:
                missing_apis.append(api_name)

        if missing_apis:
            available_apis = list(doc_cache.keys())[:10]
            error_msg = (
                f"APIs not found in document index '{document_name}': {missing_apis}\n"
                f"Available APIs in this document: {available_apis}...\n"
                f"Please check the API names or use full document search."
            )
            logger.error(error_msg)
            raise ValueError(error_msg)

        logger.info(f"Successfully found {len(result)} APIs")
        return result

    def list_apis(self, document_name: str) -> List[str]:
        """列出指定文档中的所有 API"""
        logger.info(f"Listing APIs for document '{document_name}'")

        if document_name not in self.memory_cache:
            self.memory_cache[document_name] = self._parse_api_index_file(document_name)

        apis = list(self.memory_cache[document_name].keys())
        logger.info(f"Found {len(apis)} APIs in document '{document_name}'")
        return apis

    def clear_cache(self, document_name: Optional[str] = None):
        """清除内存缓存"""
        if document_name:
            logger.info(f"Clearing cache for document '{document_name}'")
            self.memory_cache.pop(document_name, None)
        else:
            logger.info("Clearing all cache")
            self.memory_cache.clear()


if __name__ == "__main__":
    # 测试代码
    manager = APIIndexManager()

    try:
        # 测试单个 API 查找
        location = manager.get_api_location("get_price", "generic-api.md")
        print(f"API 'get_price' found at line {location['line_number']}")
        print(f"Description: {location['description']}")

        # 测试批量查找
        batch_locations = manager.get_batch_api_locations(
            ["get_price", "get_ticks", "current_snapshot"], "generic-api.md"
        )
        print(f"\nBatch lookup found {len(batch_locations)} APIs")
        for api_name, info in batch_locations.items():
            print(f"  {api_name}: line {info['line_number']}")

        # 测试列出所有 API
        apis = manager.list_apis("generic-api.md")
        print(f"\nTotal APIs in generic-api.md: {len(apis)}")

    except Exception as e:
        print(f"Error: {e}")

#!/usr/bin/env python3
"""
RQData文档缓存管理器
提供缓存机制，优化文档访问速度
"""

import io
import re
import subprocess
import time
from pathlib import Path
from typing import Optional
from urllib.parse import unquote
import logging

# 配置日志（INFO级别）
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("RQDataCacheManager")


INDEX_URL = "https://www.ricequant.com/doc/document-index.txt"
DEFAULT_CACHE_DAYS = 7


class RQDataCacheManager:
    """RQData文档缓存管理器"""

    def __init__(self, cache_dir: Optional[str] = None):
        """
        初始化缓存管理器

        Args:
            cache_dir: 缓存目录路径，默认为skill的cache/api_docs目录
        """
        if cache_dir is None:
            skill_root = Path(__file__).parent.parent
            self.cache_dir = skill_root / "cache" / "api_docs"
        else:
            self.cache_dir = Path(cache_dir)

        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _extract_filename_from_url(self, url: str) -> str:
        """从URL提取文件名"""
        filename = url.split("/")[-1]
        if not filename:
            raise ValueError(f"URL does not contain a filename: {url}")

        clean_filename = filename.split("?")[0].split("#")[0]
        decoded_filename = unquote(clean_filename)

        if not re.match(r"^[a-zA-Z0-9\-_.]+$", decoded_filename):
            raise ValueError(
                f"Invalid filename contains special characters: {decoded_filename}"
            )

        return decoded_filename

    def _is_cache_expired(
        self, cache_path: Path, max_age_days: int = DEFAULT_CACHE_DAYS
    ) -> bool:
        """检查缓存文件是否过期（基于文件修改时间）"""
        if not cache_path.exists():
            return True

        file_mtime = cache_path.stat().st_mtime
        file_age = time.time() - file_mtime
        max_age_seconds = max_age_days * 24 * 60 * 60

        return file_age > max_age_seconds

    def _get_cache_path(self, url: str) -> Path:
        """根据URL生成缓存文件路径，使用真实文件名"""
        decoded_filename = self._extract_filename_from_url(url)
        cache_path = self.cache_dir / decoded_filename
        if cache_path.exists():
            raise FileExistsError(f"Cache file already exists: {cache_path}")
        return cache_path

    def _get_cache_path_for_read(self, url: str) -> Optional[Path]:
        """获取缓存文件路径（用于读取，不检查冲突）"""
        try:
            decoded_filename = self._extract_filename_from_url(url)
        except ValueError:
            return None

        cache_path = self.cache_dir / decoded_filename
        if not cache_path.exists():
            return None

        return cache_path

    def get_cached_content(self, url: str) -> Optional[str]:
        """
        获取缓存的文档内容

        Args:
            url: 文档URL

        Returns:
            缓存的文档内容，如果缓存不存在则返回None
        """
        cache_path = self._get_cache_path_for_read(url)
        if cache_path is None:
            return None

        try:
            with open(cache_path, "r", encoding="utf-8") as f:
                return f.read()
        except IOError:
            return None

    def save_to_cache(
        self, url: str, content: str, allow_overwrite: bool = False
    ) -> None:
        """
        保存文档内容到缓存

        Args:
            url: 文档URL
            content: 文档内容
            allow_overwrite: 是否允许覆盖现有文件
        """
        decoded_filename = self._extract_filename_from_url(url)
        cache_path = self.cache_dir / decoded_filename

        if cache_path.exists():
            if not allow_overwrite:
                raise FileExistsError(f"Cache file already exists: {cache_path}")
            cache_path.unlink()

        with open(cache_path, "w", encoding="utf-8") as f:
            f.write(content)

    def fetch_document(self, url: str, timeout: int = 60, retries: int = 3) -> str:
        """
        使用curl获取文档内容，支持重试

        Args:
            url: 文档URL
            timeout: 超时时间（秒），默认60秒
            retries: 重试次数，默认3次

        Returns:
            文档内容

        Raises:
            RuntimeError: 如果获取失败
        """
        last_error = None

        for attempt in range(retries):
            try:
                result = subprocess.run(
                    ["curl", "-s", "-L", "--max-time", str(timeout), url],
                    capture_output=True,
                    timeout=timeout + 5,
                    encoding="utf-8",
                    errors="replace",
                )

                if (
                    result.returncode == 0
                    and result.stdout
                    and len(result.stdout) > 100
                ):
                    return result.stdout
                else:
                    last_error = (
                        f"curl failed with code {result.returncode} or empty response"
                    )
                    if attempt < retries - 1:
                        continue
                    raise RuntimeError(f"Failed to fetch document: {url}")

            except subprocess.TimeoutExpired:
                last_error = f"Timeout after {timeout}s"
                if attempt < retries - 1:
                    continue
                raise RuntimeError(f"Timeout fetching document: {url}")
            except Exception as e:
                last_error = str(e)
                if attempt < retries - 1:
                    continue
                raise RuntimeError(f"Error fetching document: {url}, {str(e)}")

        raise RuntimeError(f"Failed after {retries} attempts: {last_error}")

    def _parse_index_content(self, content: str) -> list[str]:
        """解析索引文档内容，提取URL列表"""
        import re

        urls = []
        url_pattern = re.compile(
            r"https://www\.ricequant\.com/doc/sources/rqdata/python/[^\)]+\.md"
        )

        for line in content.splitlines():
            matches = url_pattern.findall(line)
            for url in matches:
                if url not in urls:
                    urls.append(url)

        return urls

    def fetch_document_index(self, max_age_days: int = DEFAULT_CACHE_DAYS) -> list[str]:
        """
        获取文档索引列表（带过期检查）

        Args:
            max_age_days: 索引缓存的最大天数

        Returns:
            URL列表

        Raises:
            RuntimeError: 如果获取失败且没有可用缓存
        """
        cache_path = self._get_cache_path_for_read(INDEX_URL)

        if cache_path and not self._is_cache_expired(cache_path, max_age_days):
            try:
                with open(cache_path, "r", encoding="utf-8") as f:
                    index_content = f.read()
                return self._parse_index_content(index_content)
            except IOError:
                pass

        try:
            index_content = self.fetch_document(INDEX_URL)
            self.save_to_cache(INDEX_URL, index_content, allow_overwrite=True)
            return self._parse_index_content(index_content)
        except RuntimeError as e:
            if cache_path and cache_path.exists():
                try:
                    with open(cache_path, "r", encoding="utf-8") as f:
                        index_content = f.read()
                    print(f"Warning: Using expired index cache due to fetch error: {e}")
                    return self._parse_index_content(index_content)
                except IOError:
                    pass
            raise

    def _build_doc_name_map(self, max_age_days: int = DEFAULT_CACHE_DAYS) -> dict:
        """构建文档名到URL的映射（从索引动态生成）"""
        urls = self.fetch_document_index(max_age_days)
        mapping = {}
        for url in urls:
            filename = url.split("/")[-1]
            if filename:
                mapping[filename] = url
        return mapping

    def get_document(
        self,
        url: str,
        force_refresh: bool = False,
        max_age_days: int = DEFAULT_CACHE_DAYS,
    ) -> str:
        """
        获取文档内容（优先使用缓存，带过期检查）

        Args:
            url: 文档URL
            force_refresh: 是否强制刷新缓存
            max_age_days: 缓存的最大天数

        Returns:
            文档内容
        """
        cache_path = self._get_cache_path_for_read(url)

        if (
            not force_refresh
            and cache_path
            and not self._is_cache_expired(cache_path, max_age_days)
        ):
            try:
                with open(cache_path, "r", encoding="utf-8") as f:
                    return f.read()
            except IOError:
                pass

        try:
            content = self.fetch_document(url)
            self.save_to_cache(url, content, allow_overwrite=True)
            return content
        except RuntimeError as e:
            if cache_path and cache_path.exists():
                try:
                    with open(cache_path, "r", encoding="utf-8") as f:
                        content = f.read()
                    print(f"Warning: Using cached version due to fetch error: {e}")
                    return content
                except IOError:
                    pass
            raise

    def get_document_by_name(
        self,
        doc_name: str,
        force_refresh: bool = False,
        max_age_days: int = DEFAULT_CACHE_DAYS,
    ) -> str:
        """
        根据文档名获取文档内容（带过期检查）

        Args:
            doc_name: 文档名（如 "stock-mod.md"）
            force_refresh: 是否强制刷新缓存
            max_age_days: 缓存的最大天数

        Returns:
            文档内容
        """
        doc_map = self._build_doc_name_map(max_age_days)

        if doc_name not in doc_map:
            available = list(doc_map.keys())
            raise ValueError(
                f"Unknown document: {doc_name}\n"
                f"Available documents: {', '.join(available)}"
            )

        url = doc_map[doc_name]
        return self.get_document(url, force_refresh, max_age_days)

    def list_documents(self, max_age_days: int = DEFAULT_CACHE_DAYS) -> list[str]:
        """
        获取所有可用文档名称列表

        Args:
            max_age_days: 索引缓存的最大天数

        Returns:
            文档名称列表
        """
        doc_map = self._build_doc_name_map(max_age_days)
        return list(doc_map.keys())

    def clear_cache(self, url: Optional[str] = None) -> int:
        """
        清理缓存

        Args:
            url: 如果指定，只清理该URL的缓存；否则清理所有缓存

        Returns:
            清理的文件数量
        """
        if url:
            cache_path = self._get_cache_path_for_read(url)
            if cache_path is not None and cache_path.exists():
                cache_path.unlink()
                return 1
            return 0
        else:
            count = 0
            for cache_file in self.cache_dir.glob("*"):
                if cache_file.is_file():
                    cache_file.unlink()
                    count += 1
            return count

    def clear_all_cache(self) -> int:
        """
        清理所有缓存文件

        Returns:
            清理的文件数量
        """
        return self.clear_cache()

    def get_cache_info(self) -> dict:
        """
        获取缓存统计信息

        Returns:
            缓存统计信息字典
        """
        cache_files = [f for f in self.cache_dir.glob("*") if f.is_file()]
        total_count = len(cache_files)
        total_size = 0

        for cache_file in cache_files:
            try:
                file_size = cache_file.stat().st_size
                total_size += file_size
            except IOError:
                continue

        return {
            "cache_dir": str(self.cache_dir),
            "total_count": total_count,
            "total_size_bytes": total_size,
            "total_size_mb": round(total_size / 1024 / 1024, 2),
        }

    def read_document_lines(
        self, document_name: str, start_line: int, end_line: Optional[int] = None
    ) -> str:
        """读取文档的特定行范围"""
        logger.info(
            f"Reading lines {start_line}-{end_line or start_line} from document '{document_name}'"
        )

        cache_path = self.cache_dir / document_name

        if not cache_path.exists():
            error_msg = f"Document not found in cache: {document_name}"
            logger.error(error_msg)
            raise FileNotFoundError(f"{error_msg}\nCache directory: {self.cache_dir}")

        try:
            with open(cache_path, "r", encoding="utf-8") as f:
                lines = f.readlines()
        except Exception as e:
            error_msg = f"Error reading document '{document_name}': {e}"
            logger.error(error_msg)
            raise IOError(error_msg)

        start_idx = max(0, start_line - 1)
        end_idx = end_line if end_line else start_line
        end_idx = min(len(lines), end_idx)

        selected_lines = lines[start_idx:end_idx]
        content = "".join(selected_lines)

        logger.info(f"Extracted {len(selected_lines)} lines")
        return content

    def get_api_definition(
        self,
        api_name: str,
        document_name: str,
        context_lines: int = 100,
        api_index_manager: Optional[object] = None,
    ) -> dict:
        """获取 API 定义（使用 api_index 优化）"""
        logger.info(
            f"Getting API definition for '{api_name}' in document '{document_name}'"
        )

        if api_index_manager is None:
            from api_index_manager import APIIndexManager

            api_index_manager = APIIndexManager()

        try:
            # type: ignore - APIIndexManager 会在运行时正确导入
            api_info = api_index_manager.get_api_location(api_name, document_name)  # type: ignore
            line_number = api_info["line_number"]
            logger.info(f"API '{api_name}' found at line {line_number}")
        except (ValueError, FileNotFoundError) as e:
            logger.error(f"Failed to get API location: {e}")
            raise

        start_line = max(1, line_number - context_lines // 2)
        end_line = line_number + context_lines // 2

        try:
            content = self.read_document_lines(document_name, start_line, end_line)
            logger.info(
                f"Successfully read {len(content)} characters for API '{api_name}'"
            )
        except (FileNotFoundError, IOError) as e:
            logger.error(f"Failed to read document lines: {e}")
            raise

        return {
            "api_name": api_name,
            "document_name": document_name,
            "line_number": line_number,
            "description": api_info.get("description", ""),
            "content": content,
            "context_range": (start_line, end_line),
        }

    def find_factor_download_link(
        self, doc_content: str, api_start_line: int = 0
    ) -> Optional[str]:
        """从API文档的特定API部分查找宏观因子xlsx下载链接"""
        lines = doc_content.split("\n")
        search_start = api_start_line if api_start_line > 0 else 0

        for i, line in enumerate(lines[search_start:], search_start):
            if ("factors" in line.lower() or "宏观因子" in line) and ".xlsx" in line:
                match = re.search(r'https://[^\s"\'<>]+\.xlsx', line)
                if match:
                    url = match.group(0)
                    logger.info(f"Found factor download link at line {i + 1}: {url}")
                    return url

        logger.warning("No factor download link found in document")
        return None

    def fetch_binary_file(
        self, url: str, output_name: str, timeout: int = 60, retries: int = 3
    ) -> Path:
        """下载二进制文件（xlsx等）"""
        output_path = self.cache_dir / output_name

        last_error = None
        for attempt in range(retries):
            try:
                result = subprocess.run(
                    [
                        "curl",
                        "-s",
                        "-L",
                        "--max-time",
                        str(timeout),
                        "-o",
                        str(output_path),
                        url,
                    ],
                    capture_output=True,
                    timeout=timeout + 5,
                    encoding="utf-8",
                    errors="replace",
                )

                if (
                    result.returncode == 0
                    and output_path.exists()
                    and output_path.stat().st_size > 0
                ):
                    logger.info(
                        f"Successfully downloaded {output_name} ({output_path.stat().st_size} bytes)"
                    )
                    return output_path
                else:
                    last_error = f"curl failed with code {result.returncode}"
                    if output_path.exists():
                        output_path.unlink()
                    if attempt < retries - 1:
                        continue

            except subprocess.TimeoutExpired:
                last_error = f"Timeout after {timeout}s"
                if output_path.exists():
                    output_path.unlink()
                if attempt < retries - 1:
                    continue
            except Exception as e:
                last_error = str(e)
                if output_path.exists():
                    output_path.unlink()
                if attempt < retries - 1:
                    continue

        raise RuntimeError(
            f"Failed to download binary file after {retries} attempts: {last_error}"
        )

    def convert_xlsx_to_csv(
        self, xlsx_path: Path, csv_name: Optional[str] = None
    ) -> Path:
        """将xlsx文件转换为CSV"""
        if not xlsx_path.exists():
            raise FileNotFoundError(f"Xlsx file not found: {xlsx_path}")

        if csv_name is None:
            csv_name = xlsx_path.stem + ".csv"
        csv_path = self.cache_dir / csv_name

        try:
            import pandas as pd

            df = pd.read_excel(xlsx_path, engine="openpyxl")
            df.to_csv(csv_path, index=False, encoding="utf-8-sig")
            logger.info(
                f"Converted {xlsx_path.name} to {csv_name} ({csv_path.stat().st_size} bytes)"
            )
            return csv_path

        except ImportError:
            raise ImportError("pandas and openpyxl are required for xlsx conversion")
        except Exception as e:
            raise RuntimeError(f"Failed to convert xlsx to csv: {e}")

    def download_and_convert_factor_file(
        self, doc_name: str = "macro-economy.md"
    ) -> Optional[Path]:
        """下载并转换宏观因子名称文件"""
        cache_path = self.cache_dir / doc_name

        if not cache_path.exists():
            logger.warning(f"Document not found: {doc_name}")
            return None

        with open(cache_path, "r", encoding="utf-8") as f:
            content = f.read()

        link = self.find_factor_download_link(content, api_start_line=87)
        if not link:
            logger.warning("No factor download link found in macro-economy.md")
            return None

        xlsx_path = self.fetch_binary_file(link, "macro_factor_names.xlsx")
        csv_path = self.convert_xlsx_to_csv(xlsx_path, "macro_factor_names.csv")

        if xlsx_path.exists():
            xlsx_path.unlink()
            logger.info(f"Removed temporary xlsx file: {xlsx_path.name}")

        return csv_path


if __name__ == "__main__":
    cache_mgr = RQDataCacheManager()

    test_url = "https://www.ricequant.com/doc/sources/rqdata/python/stock-mod.md"

    content = cache_mgr.get_document(test_url)
    print(f"Document length: {len(content)}")

    info = cache_mgr.get_cache_info()
    print(f"\nCache Info:")
    print(f"  Total files: {info['total_count']}")
    print(f"  Total size: {info['total_size_mb']} MB")

    cleared = cache_mgr.clear_cache()
    print(f"\nCleared {cleared} cache files")

import rqdatac

"""检查 RQData License 是否有效"""
try:
    rqdatac.init()
    print("✓ RQData License 有效")
except rqdatac.RQDataError:
    print("❌ RQData License 未激活或已过期")
    print("申请试用: https://www.ricequant.com/welcome/trial/rqsdk-cloud")

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
资产代码索引管理器
提供股票代码搜索功能，支持通过代码或名称模糊匹配查找股票
"""

import time
from pathlib import Path
from typing import Optional

INDEX_CACHE_DAYS = 1


class CodeIndexManager:
    """资产代码索引管理器"""

    def __init__(self, cache_dir: Optional[str] = None):
        """
        初始化索引管理器

        Args:
            cache_dir: 缓存目录路径，默认为 skill 的 cache/code_index 目录
        """
        if cache_dir is None:
            skill_root = Path(__file__).parent.parent
            self.cache_dir = skill_root / "cache" / "code_index"
        else:
            self.cache_dir = Path(cache_dir)

        self.cache_dir.mkdir(parents=True, exist_ok=True)

    def _get_index_path(self, market: str, asset_type: str = "CS") -> Path:
        """获取索引文件路径"""
        filename = f"{market}_{asset_type.lower()}_code_index.md"
        return self.cache_dir / filename

    def _is_index_expired(self, index_path: Path) -> bool:
        """检查索引是否过期"""
        if not index_path.exists():
            return True

        file_mtime = index_path.stat().st_mtime
        file_age = time.time() - file_mtime
        max_age_seconds = INDEX_CACHE_DAYS * 24 * 60 * 60

        return file_age > max_age_seconds

    def build_index(self, market: str = "cn", asset_type: str = "CS") -> bool:
        """
        构建资产代码索引

        Args:
            market: 市场代码，'cn' - A股，'hk' - 港股
            asset_type: 资产类型，'CS' - 股票

        Returns:
            是否成功构建索引
        """
        import rqdatac

        index_path = self._get_index_path(market, asset_type)

        try:
            df = rqdatac.all_instruments(type=asset_type, market=market)

            if df is None or df.empty:
                print(f"[WARN] No data returned for {market}/{asset_type}")
                return False

            rows = []
            for _, row in df.iterrows():
                order_book_id = row.get("order_book_id", "")
                symbol = row.get("symbol", "")
                abbrev_symbol = row.get("abbrev_symbol", "")

                if order_book_id:
                    rows.append(
                        {
                            "order_book_id": order_book_id,
                            "symbol": symbol,
                            "abbrev_symbol": abbrev_symbol,
                        }
                    )

            lines = [
                "| order_book_id | symbol | abbrev_symbol |",
                "|---------------|--------|----------------|",
            ]

            for r in rows:
                symbol_escaped = r["symbol"].replace("|", "\\|")
                abbrev_escaped = (
                    r["abbrev_symbol"].replace("|", "\\|") if r["abbrev_symbol"] else ""
                )
                lines.append(
                    f"| {r['order_book_id']} | {symbol_escaped} | {abbrev_escaped} |"
                )

            content = "\n".join(lines)

            with open(index_path, "w", encoding="utf-8") as f:
                f.write(content)

            return True

        except Exception as e:
            print(f"[FAIL] Failed to build index: {e}")
            return False

    def get_index(
        self, market: str = "cn", asset_type: str = "CS", force_refresh: bool = False
    ) -> Optional[list[dict]]:
        """
        获取资产代码索引

        Args:
            market: 市场代码
            asset_type: 资产类型
            force_refresh: 是否强制刷新索引

        Returns:
            索引数据列表，每个元素为 {'order_book_id', 'symbol', 'abbrev_symbol'}
        """
        index_path = self._get_index_path(market, asset_type)

        if force_refresh or self._is_index_expired(index_path):
            if not self.build_index(market, asset_type):
                return None

        return self._parse_index(index_path)

    def _parse_index(self, index_path: Path) -> Optional[list[dict]]:
        """解析索引文件"""
        if not index_path.exists():
            return None

        try:
            with open(index_path, "r", encoding="utf-8") as f:
                lines = f.readlines()

            records = []
            in_table = False

            for line in lines:
                line = line.strip()

                if line.startswith("| order_book_id"):
                    in_table = True
                    continue

                if in_table and line.startswith("|"):
                    if line == "|---|---|---|" or not line.strip():
                        continue

                    parts = [p.strip() for p in line.split("|")]
                    if len(parts) >= 4:
                        order_book_id = parts[1]
                        symbol = parts[2]
                        abbrev_symbol = parts[3] if len(parts) > 3 else ""

                        if order_book_id and order_book_id != "order_book_id":
                            records.append(
                                {
                                    "order_book_id": order_book_id,
                                    "symbol": symbol,
                                    "abbrev_symbol": abbrev_symbol,
                                }
                            )

            return records

        except Exception as e:
            print(f"[FAIL] Failed to parse index: {e}")
            return None

    def search(
        self, query: str, market: str = "cn", limit: int = 10, asset_type: str = "CS"
    ) -> list[dict]:
        """
        搜索资产代码

        Args:
            query: 查询字符串（代码或名称）
            market: 市场代码
            limit: 返回结果数量限制
            asset_type: 资产类型

        Returns:
            匹配结果列表，每个元素为 {'order_book_id', 'symbol', 'abbrev_symbol', 'match_type'}
        """
        records = self.get_index(market, asset_type)

        if not records:
            return []

        query = query.strip()
        if not query:
            return []

        results = []
        query_lower = query.lower()

        for r in records:
            match_type = None

            order_book_id = r["order_book_id"]
            symbol = r["symbol"]
            abbrev_symbol = r.get("abbrev_symbol", "")

            code_without_suffix = (
                order_book_id.split(".")[0] if "." in order_book_id else order_book_id
            )

            if order_book_id.lower() == query_lower:
                match_type = "code_exact"
            elif symbol.lower() == query_lower:
                match_type = "name_exact"
            elif code_without_suffix == query:
                match_type = "code_prefix"
            elif query_lower in symbol.lower():
                match_type = "name_contains"
            elif abbrev_symbol and query_lower == abbrev_symbol.lower():
                match_type = "abbrev_exact"
            elif abbrev_symbol and query_lower in abbrev_symbol.lower():
                match_type = "abbrev_contains"

            if match_type:
                results.append(
                    {
                        "order_book_id": order_book_id,
                        "symbol": symbol,
                        "abbrev_symbol": abbrev_symbol,
                        "match_type": match_type,
                    }
                )

        results.sort(
            key=lambda x: (
                0 if x["match_type"] == "code_exact" else 1,
                0 if x["match_type"] == "name_exact" else 1,
                0 if x["match_type"] == "code_prefix" else 1,
                0 if x["match_type"] == "abbrev_exact" else 1,
                0 if x["match_type"] == "name_contains" else 1,
                0 if x["match_type"] == "abbrev_contains" else 1,
            )
        )

        return results[:limit]


def resolve_stock_code(
    query: str, market: str = "cn", limit: int = 10, asset_type: str = "CS"
) -> list[dict]:
    """
    通过股票代码或公司名称查询真实股票代码

    支持输入:
    - 纯数字: 600519, 000001
    - 带后缀: 600519.SH, 000001.SZ
    - 公司名称: 贵州茅台, 智谱AI

    Args:
        query: 查询字符串（代码或名称）
        market: 市场代码，'cn' - A股，'hk' - 港股
        limit: 返回结果数量限制
        asset_type: 资产类型，默认 'CS' (股票)

    Returns:
        匹配结果列表，每个元素为:
        {
            'order_book_id': '600519.XSHG',
            'symbol': '贵州茅台',
            'abbrev_symbol': 'GZMT',
            'match_type': 'name_exact'
        }
    """
    manager = CodeIndexManager()
    return manager.search(query, market, limit, asset_type)


if __name__ == "__main__":
    import argparse
    import rqdatac

    parser = argparse.ArgumentParser(
        description="资产代码搜索工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
允许的市场代码:
  cn    A股 (中国内地市场)
  hk    港股 (香港市场)

允许的资产类型:
  CS          股票 (Common Stock)
  ETF         交易所交易基金
  Future      期货
  Option      期权
  Convertible 可转债
  INDX        指数
  LOF         上市型开放式基金
  FUND        基金

使用示例:
  python code_index_manager.py -q "贵州茅台" -m cn -t CS
  python code_index_manager.py --query "600519" --market cn
  python code_index_manager.py --query "腾讯" --market hk --type CS
        """,
    )
    parser.add_argument("--query", "-q", required=True, help="查询字符串（代码或名称）")
    parser.add_argument(
        "--market",
        "-m",
        default="cn",
        choices=["cn", "hk"],
        help="市场代码: cn (A股), hk (港股) [default: cn]",
    )
    parser.add_argument(
        "--type",
        "-t",
        default="CS",
        help="资产类型: CS(股票), ETF, Future, Option, Convertible, INDX [default: CS]",
    )
    parser.add_argument(
        "--limit",
        "-l",
        type=int,
        default=10,
        help="返回结果数量 [default: 10]",
    )

    args = parser.parse_args()

    rqdatac.init()

    manager = CodeIndexManager()
    results = manager.search(args.query, args.market, args.limit, args.type)

    if results:
        for r in results:
            print(f"{r['order_book_id']} | {r['symbol']} | {r['match_type']}")
    else:
        print("未找到匹配结果")

#!/usr/bin/env python3
"""
RQData文档索引获取脚本
从官方获取文档索引并保存为Markdown文件
"""

import requests
from pathlib import Path
from datetime import datetime
from typing import Optional


class DocumentIndexFetcher:
    """文档索引获取器"""

    def __init__(self):
        self.index_url = "https://www.ricequant.com/doc/document-index.txt"

    def fetch(self, timeout: int = 60) -> str:
        """获取文档索引内容"""
        response = requests.get(self.index_url, timeout=timeout)
        response.raise_for_status()
        return response.text

    def save(self, content: str, output_path: Optional[Path] = None) -> Path:
        """保存内容到文件"""
        if output_path is None:
            skill_root = Path(__file__).parent.parent
            output_path = skill_root / "cache" / "document_index.md"

        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text(content, encoding="utf-8")
        return output_path

    def run(self, output_path: Optional[Path] = None) -> Path:
        """执行获取并保存"""
        print(f"正在获取文档索引: {self.index_url}")
        content = self.fetch()
        print(f"获取成功，内容长度: {len(content)} 字符")

        saved_path = self.save(content, output_path)
        print(f"已保存到: {saved_path}")
        return saved_path


def main():
    import argparse

    parser = argparse.ArgumentParser(description="RQData文档索引获取脚本")
    parser.add_argument("-o", "--output", type=str, help="输出文件路径")
    args = parser.parse_args()

    output_path = Path(args.output) if args.output else None
    fetcher = DocumentIndexFetcher()

    try:
        saved_path = fetcher.run(output_path)
        print(f"\n成功: {saved_path}")
        return 0
    except Exception as e:
        print(f"\n错误: {e}")
        return 1


if __name__ == "__main__":
    import sys

    sys.exit(main())

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RQData API Index Generator

Scans all markdown documentation files in the cache/api_docs directory and generates
individual API index files for each source file.
"""

import re
import sys
from pathlib import Path
from typing import List, Dict, Optional


class APIIndexGenerator:
    """Generates API index files from markdown documentation"""

    def __init__(
        self, api_docs_dir: Optional[Path] = None, output_dir: Optional[Path] = None
    ):
        """
        Initialize the generator

        Args:
            api_docs_dir: Directory containing source markdown files
            output_dir: Directory to output index files
        """
        if api_docs_dir is None:
            # Default to cache/api_docs relative to script location
            script_dir = Path(__file__).parent
            skill_dir = script_dir.parent
            self.api_docs_dir = skill_dir / "cache" / "api_docs"
        else:
            self.api_docs_dir = Path(api_docs_dir)

        if output_dir is None:
            # Default to cache/api_index relative to script location
            skill_dir = Path(__file__).parent.parent
            self.output_dir = skill_dir / "cache" / "api_index"
        else:
            self.output_dir = Path(output_dir)

        # Patterns to exclude (non-API entries)
        self.exclude_patterns = [
            re.compile(r"^[A-Z]+$"),  # All caps (API, FAQ, etc.)
            re.compile(r"^\d+$"),  # Pure numbers
            re.compile(r"^[\u4e00-\u9fff]+$"),  # Pure Chinese
        ]

        # Regex to match all ## headings (both top-level and API)
        self.heading_pattern = re.compile(r"^#{2,3}\s+(.+?)\s*(\{#.+})?$")

        # Regex to match API headings with function name and description
        # Format: ## or ### function_name - description {#xxx-API-anchor}
        # Matches any anchor with -API- in it (rqdata-API, stock-API, etc.)
        self.api_pattern = re.compile(
            r"^#{2,3}\s+([\w\.]+)\s*[-–]\s*(.+?)\s*\{#[^}]+-API-[^}]+\}\s*$"
        )

        self.max_paragraphs = (
            5  # Maximum paragraphs to extract for detailed description
        )

    def _is_valid_api(self, api_name: str) -> bool:
        """
        Check if the name is a valid API function

        Args:
            api_name: The API name to check

        Returns:
            True if valid API, False otherwise
        """
        # Must contain underscore or dot (module.function or function_name)
        # OR be a valid alphanumeric name (at least 2 chars, starting with letter)
        if "_" not in api_name and "." not in api_name:
            # Check if it's a valid alphanumeric name (e.g., instruments, get_price)
            if not re.match(r"^[a-zA-Z][a-zA-Z0-9]+$", api_name):
                return False

        # Check against exclude patterns
        for pattern in self.exclude_patterns:
            if pattern.match(api_name):
                return False

        return True

    def _extract_detailed_description(
        self, lines: List[str], start_line: int, end_line: int
    ) -> str:
        """
        Extract detailed description from lines between API header and next section.

        Args:
            lines: All lines from the file
            start_line: Line number after the API header (0-indexed)
            end_line: Line number before the next API header (0-indexed)

        Returns:
            Detailed description as a string, or empty string if not found
        """
        paragraphs = []
        prev_was_empty = False
        table_line_count = 0
        in_code_block = False

        for i in range(start_line, min(end_line, len(lines))):
            line = lines[i].rstrip("\n")

            # Handle code blocks
            if line.strip().startswith("```"):
                in_code_block = not in_code_block
                continue

            if in_code_block:
                continue

            # Check for table (two consecutive lines starting with |)
            if line.strip().startswith("|"):
                table_line_count += 1
                if table_line_count >= 2:
                    break  # Stop at table
                continue

            # Reset table counter if line doesn't start with |
            if not line.strip().startswith("|"):
                table_line_count = 0

            # Check for new API header
            if line.strip().startswith("##"):
                break

            # Collect paragraphs (non-empty lines)
            if line.strip():
                if not prev_was_empty or not paragraphs:
                    if not paragraphs:
                        paragraphs.append(line.strip())
                    else:
                        paragraphs[-1] += " " + line.strip()
                else:
                    paragraphs.append(line.strip())
                prev_was_empty = False
            else:
                prev_was_empty = True
                # Stop if we have enough paragraphs
                if len(paragraphs) >= self.max_paragraphs:
                    break

        return "<br/>".join(paragraphs) if paragraphs else ""

    def extract_apis_from_file(self, file_path: Path) -> List[Dict]:
        """
        Extract API definitions from a markdown file

        Args:
            file_path: Path to the markdown file

        Returns:
            List of API dictionaries with name, description, line_number, end_line_number
        """
        apis = []

        try:
            with open(file_path, "r", encoding="utf-8") as f:
                lines = f.readlines()

            total_lines = len(lines)
            api_positions = []

            for line_num, line in enumerate(lines):
                line = line.strip()

                match = self.api_pattern.match(line)
                if match:
                    api_name = match.group(1)
                    description = match.group(2).strip()

                    if self._is_valid_api(api_name):
                        api_positions.append(
                            {
                                "api_name": api_name,
                                "description": description,
                                "line_number": line_num + 1,
                            }
                        )

            for i, api in enumerate(api_positions):
                if i + 1 < len(api_positions):
                    api["end_line_number"] = api_positions[i + 1]["line_number"] - 1
                else:
                    api["end_line_number"] = total_lines

                # Extract detailed description
                start_idx = api["line_number"]  # 0-indexed, line_number is 1-indexed
                end_idx = api["end_line_number"]

                detailed_desc = self._extract_detailed_description(
                    lines, start_idx, end_idx
                )
                api["detailed_description"] = detailed_desc

                apis.append(api)

        except Exception as e:
            print(f"  Warning: Error processing {file_path.name}: {e}")

        return apis

    def _extract_titles_from_file(self, file_path: Path) -> List[str]:
        """
        Extract all titles from a markdown file.

        Args:
            file_path: Path to the markdown file

        Returns:
            List of title strings (top-level headings and API descriptions without function names)
        """
        titles = []

        try:
            with open(file_path, "r", encoding="utf-8") as f:
                lines = f.readlines()

            for line in lines:
                line = line.strip()

                # Check if it's a heading
                match = self.heading_pattern.match(line)
                if not match:
                    continue

                heading_text = match.group(1).strip()

                # Check if it's an API heading (contains function name)
                api_match = self.api_pattern.match(line)
                if api_match:
                    # Extract only the description part (after the dash)
                    description = api_match.group(2).strip()
                    titles.append(description)
                else:
                    # It's a top-level heading, use as-is
                    titles.append(heading_text)

        except Exception as e:
            print(f"  Warning: Error extracting titles from {file_path.name}: {e}")

        return titles

    def generate_doc_index(self) -> Optional[Path]:
        """
        Generate api_doc_index.md - a combined index of all API documents.

        Returns:
            Path to the generated file, or None if failed
        """
        exclude_files = {"changelogs.md", "manual.md"}

        output_file = self.api_docs_dir.parent / "api_doc_index.md"

        source_files = [
            f for f in self.api_docs_dir.glob("*.md") if f.name not in exclude_files
        ]

        if not source_files:
            print(f"Error: No markdown files found in {self.api_docs_dir}")
            return None

        # Check if regeneration is needed based on source file freshness
        if output_file.exists():
            output_mtime = output_file.stat().st_mtime
            # Only regenerate if any source file is newer than output
            if all(f.stat().st_mtime <= output_mtime for f in source_files):
                print(f"Skipped: {output_file.name} (up to date)")
                return output_file

        md_files = sorted(self.api_docs_dir.glob("*.md"))

        doc_descriptions = []

        for md_file in md_files:
            if md_file.name in exclude_files:
                continue

            titles = self._extract_titles_from_file(md_file)
            if titles:
                description = "。".join(titles)
                doc_descriptions.append(
                    {"filename": md_file.name, "description": description}
                )

        if not doc_descriptions:
            print("Warning: No document descriptions found")
            return None

        with open(output_file, "w", encoding="utf-8") as f:
            f.write("# API Doc Index\n\n")
            f.write("| Document | Description |\n")
            f.write("|----------|-------------|\n")

            for doc in doc_descriptions:
                escaped_desc = doc["description"].replace("|", "\\|")
                f.write(f"| {doc['filename']} | {escaped_desc} |\n")

        print(f"Generated: {output_file.name}")
        return output_file

    def generate_index_file(
        self, source_file: Path, apis: List[Dict]
    ) -> Optional[Path]:
        """
        Generate an index file for a source markdown file

        Args:
            source_file: Path to the source markdown file
            apis: List of API dictionaries

        Returns:
            Path to the generated index file, or None if no APIs found
        """
        if not apis:
            print(f"  No APIs found in {source_file.name}, skipping...")
            return None

        # Sort by line number
        apis.sort(key=lambda x: x["line_number"])

        # Create output file path
        source_name = source_file.stem  # filename without extension
        output_file = self.output_dir / f"{source_name}_index.md"

        # Check freshness: skip if output exists and source hasn't changed
        if output_file.exists():
            output_mtime = output_file.stat().st_mtime
            source_mtime = source_file.stat().st_mtime
            if source_mtime <= output_mtime:
                print(f"  Skipped: {output_file.name} (up to date)")
                return None

        # Ensure output directory exists
        self.output_dir.mkdir(parents=True, exist_ok=True)

        # Write markdown content
        with open(output_file, "w", encoding="utf-8") as f:
            # Header
            f.write(f"# API Index for {source_file.name}\n\n")

            # Summary
            f.write("## Summary\n\n")
            f.write(f"- Source File: {source_file.name}\n")
            f.write(f"- Total APIs: {len(apis)}\n\n")

            # API Definitions
            f.write("## API Definitions (Sorted by Line Number)\n\n")
            f.write("| API Name | Description | Line Range |\n")
            f.write("|----------|-------------|------------|\n")

            for api in apis:
                # Build description: combine short description and detailed description
                short_desc = api["description"]
                detailed_desc = api.get("detailed_description", "")

                if detailed_desc:
                    full_description = f"{short_desc}<br/>{detailed_desc}"
                else:
                    full_description = short_desc

                # Escape pipe characters in description
                full_description = full_description.replace("|", "\\|")
                # Replace newlines with <br/>
                full_description = full_description.replace("\n", "<br/>")

                line_range = f"{api['line_number']}-{api['end_line_number']}"
                f.write(
                    f"| `{api['api_name']}` | {full_description} | {line_range} |\n"
                )

            f.write("\n")

        return output_file

    def run(self) -> int:
        """
        Run the index generator

        Returns:
            Number of index files generated
        """
        print(f"API docs directory: {self.api_docs_dir}")
        print(f"Output directory: {self.output_dir}")
        print()

        # Check if api_docs directory exists
        if not self.api_docs_dir.exists():
            print(f"Error: API docs directory not found: {self.api_docs_dir}")
            return 0

        # Find all markdown files
        md_files = sorted(self.api_docs_dir.glob("*.md"))

        if not md_files:
            print(f"Error: No markdown files found in {self.api_docs_dir}")
            return 0

        print(f"Found {len(md_files)} markdown files\n")

        # Files to exclude from API index generation
        exclude_from_api_index = {"changelogs.md", "manual.md", "api_doc_index.md"}

        # Process each file
        index_count = 0
        total_apis = 0

        for md_file in md_files:
            if md_file.name in exclude_from_api_index:
                print(f"Skipping: {md_file.name} (excluded)")
                continue

            print(f"Processing: {md_file.name}")

            try:
                # Extract APIs from file
                apis = self.extract_apis_from_file(md_file)

                if apis:
                    # Generate index file
                    output_file = self.generate_index_file(md_file, apis)
                    if output_file:
                        print(f"  Generated: {output_file.name} ({len(apis)} APIs)")
                        index_count += 1
                        total_apis += len(apis)
                else:
                    print(f"  No APIs found, skipping...")

            except Exception as e:
                print(f"  Error: {e}")
                continue

        print()
        print("=" * 50)
        print(f"Summary:")
        print(f"  - Index files generated: {index_count}")
        print(f"  - Total APIs indexed: {total_apis}")
        print(f"  - Output directory: {self.output_dir}")
        print("=" * 50)
        print()

        print("Generating API doc index...")
        doc_index_file = self.generate_doc_index()
        if doc_index_file:
            print(f"  Success: {doc_index_file.name}")
        else:
            print("  Warning: Failed to generate API doc index")

        return index_count


def main():
    """Main entry point"""
    print("RQData API Index Generator")
    print("=" * 50)
    print()

    generator = APIIndexGenerator()
    index_count = generator.run()

    if index_count > 0:
        print(f"\nSuccess! Generated {index_count} index file(s).")
        return 0
    else:
        print("\nNo index files generated.")
        return 1


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RQData文档批量缓存初始化脚本
首次运行时缓存所有RQData Python API文档

Usage:
    python init_cache.py [--force-refresh]
"""

import sys
import io
import os
import argparse

os.environ["PYTHONIOENCODING"] = "utf-8"
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace")

from pathlib import Path
from cache_manager import RQDataCacheManager


def main():
    """批量缓存所有RQData文档"""
    parser = argparse.ArgumentParser(description="RQData文档缓存初始化")
    parser.add_argument(
        "--force-refresh",
        action="store_true",
        help="强制重新下载所有文档（忽略缓存）",
    )
    args = parser.parse_args()

    force_refresh = args.force_refresh

    print("=" * 80)
    print("RQData文档批量缓存初始化")
    if force_refresh:
        print("[强制刷新模式]")
    print("=" * 80)
    print()

    cache_mgr = RQDataCacheManager()

    info = cache_mgr.get_cache_info()
    print(f"缓存目录: {info['cache_dir']}")
    print(f"当前缓存文件数: {info['total_count']}")
    print(f"缓存大小: {info['total_size_mb']} MB")
    print()

    print("获取文档索引...")
    try:
        urls = cache_mgr.fetch_document_index()
        print(f"找到 {len(urls)} 个文档")
    except RuntimeError as e:
        print(f"[X] 无法获取文档索引: {e}")
        return 1

    print()
    print(f"准备缓存 {len(urls)} 个文档...")
    print()

    success_count = 0
    skip_count = 0
    fail_count = 0

    for i, url in enumerate(urls, 1):
        doc_name = url.split("/")[-1]
        print(f"[{i}/{len(urls)}] {doc_name}...", end=" ")

        try:
            content = cache_mgr.get_document(url, force_refresh=force_refresh)
            success_count += 1

        except RuntimeError as e:
            cache_path = cache_mgr._get_cache_path_for_read(url)
            if cache_path and cache_path.exists():
                print(f"[WARN] 使用缓存版本")
                skip_count += 1
            else:
                print(f"[FAIL] {str(e)}")
                fail_count += 1

    print()
    print("=" * 80)
    print("缓存完成")
    print("=" * 80)
    print(f"成功下载: {success_count}")
    print(f"使用缓存: {skip_count}")
    print(f"失败: {fail_count}")
    print()

    final_info = cache_mgr.get_cache_info()
    print(f"最终缓存文件数: {final_info['total_count']}")
    print(f"总缓存大小: {final_info['total_size_mb']} MB")
    print()

    if force_refresh:
        print("=" * 80)
        print("强制刷新模式：下载宏观因子名称文件...")
        print("=" * 80)
    else:
        print("=" * 80)
        print("下载宏观因子名称文件...")
        print("=" * 80)
    try:
        csv_path = cache_mgr.download_and_convert_factor_file("macro-economy.md")
        if csv_path and csv_path.exists():
            print(f"[OK] 宏观因子名称已保存到: {csv_path.name}")
        else:
            print("[WARN] 宏观因子名称文件下载失败")
    except Exception as e:
        print(f"[WARN] 宏观因子名称文件处理异常: {e}")
    print()

    if fail_count > 0 and success_count + skip_count == 0:
        print("[WARN] 所有文档缓存失败，请检查网络连接后重试")
        return 1
    elif fail_count > 0:
        print("[WARN] 部分文档缓存失败，请检查网络连接后重试")
    return 0


if __name__ == "__main__":
    sys.exit(main())

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RQData Skill Initialization Script

Performs all necessary setup before using the RQData skill:
1. Verify RQData license
2. Check if cache needs refresh (any api_doc expired)
3. Initialize document cache (only if needed)
4. Generate API indices (only if cache was refreshed)
5. Generate macro factor file (only if cache was refreshed)

Exit codes:
- 0: Success
- 1: License check failed
- 2: Cache initialization failed
- 3: Both failed
"""

import sys
import subprocess
import os
import time
from pathlib import Path
from typing import Optional

DEFAULT_CACHE_DAYS = 7


def print_header():
    """Print script header"""
    pass  # Suppressed for clean output on success


def check_license():
    """Check RQData license by calling rqdatac.init()"""
    try:
        import rqdatac

        rqdatac.init()
        return True
    except ImportError:
        print("[FAIL] rqdatac not installed", file=sys.stderr)
        return False
    except Exception as e:
        print(f"[FAIL] RQData license invalid: {e}", file=sys.stderr)
        return False


def check_api_docs_expired(cache_dir: Optional[Path] = None) -> bool:
    """Check if any api_docs file is expired (older than DEFAULT_CACHE_DAYS)"""
    if cache_dir is None:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        skill_root = Path(script_dir).parent
        cache_dir = skill_root / "cache" / "api_docs"

    if not cache_dir.exists():
        return True

    max_age_seconds = DEFAULT_CACHE_DAYS * 24 * 60 * 60
    current_time = time.time()

    md_files = list(cache_dir.glob("*.md"))
    if not md_files:
        return True

    for md_file in md_files:
        file_age = current_time - md_file.stat().st_mtime
        if file_age > max_age_seconds:
            return True

    return False


def run_cache_init(force_refresh: bool = False):
    """Run cache initialization (calls init_cache.py)"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    init_cache_path = os.path.join(script_dir, "init_cache.py")

    if not os.path.exists(init_cache_path):
        print(f"[FAIL] init_cache.py not found at: {init_cache_path}", file=sys.stderr)
        return False

    args = [sys.executable, init_cache_path]
    if force_refresh:
        args.append("--force-refresh")

    try:
        result = subprocess.run(
            args,
            capture_output=True,
            timeout=300,  # 5 minutes timeout when refreshing all
            encoding="utf-8",
            errors="replace",
        )

        if result.returncode == 0:
            return True
        else:
            print("[FAIL] Cache initialization failed", file=sys.stderr)
            if result.stdout:
                print("Output:", file=sys.stderr)
                print(result.stdout, file=sys.stderr)
            if result.stderr:
                print("Errors:", file=sys.stderr)
                print(result.stderr, file=sys.stderr)
            return False

    except subprocess.TimeoutExpired:
        print("[FAIL] Cache initialization timed out", file=sys.stderr)
        return False
    except Exception as e:
        print(f"[FAIL] Cache initialization error: {e}", file=sys.stderr)
        return False


def generate_api_indices():
    """生成 API 索引文件"""

    script_dir = os.path.dirname(os.path.abspath(__file__))
    generate_script = os.path.join(script_dir, "generate_api_index.py")

    if not os.path.exists(generate_script):
        error_msg = f"generate_api_index.py not found at: {generate_script}"
        print(f"[FAIL] {error_msg}", file=sys.stderr)
        return False

    try:
        result = subprocess.run(
            [sys.executable, generate_script],
            capture_output=True,
            timeout=180,
            encoding="utf-8",
            errors="replace",
        )

        if result.returncode == 0:
            return True
        else:
            error_msg = "API indices generation failed"
            print(f"[FAIL] {error_msg}", file=sys.stderr)
            if result.stdout:
                print("Output:", file=sys.stderr)
                print(result.stdout, file=sys.stderr)
            if result.stderr:
                print("Errors:", file=sys.stderr)
                print(result.stderr, file=sys.stderr)
            return False

    except subprocess.TimeoutExpired:
        error_msg = "API indices generation timed out"
        print(f"[FAIL] {error_msg}", file=sys.stderr)
        return False
    except Exception as e:
        error_msg = f"API indices generation error: {e}"
        print(f"[FAIL] {error_msg}", file=sys.stderr)
        return False


def build_code_indices():
    """构建资产代码索引"""

    script_dir = os.path.dirname(os.path.abspath(__file__))
    code_index_script = os.path.join(script_dir, "code_index_manager.py")

    if not os.path.exists(code_index_script):
        print(f"[WARN] code_index_manager.py not found, skipping code index build")
        return True

    try:
        from code_index_manager import CodeIndexManager

        manager = CodeIndexManager()

        cn_ok = manager.build_index("cn", "CS")
        if not cn_ok:
            print("[WARN] CN stock code index build failed")

        return True

    except Exception as e:
        print(f"[WARN] Code index build error: {e}", file=sys.stderr)
        return True


def refresh_macro_factor_file():
    """刷新宏观因子名称文件"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    init_cache_path = os.path.join(script_dir, "init_cache.py")

    if not os.path.exists(init_cache_path):
        print(f"[WARN] init_cache.py not found, skipping macro factor file refresh")
        return True

    try:
        from cache_manager import RQDataCacheManager

        cache_mgr = RQDataCacheManager()
        csv_path = cache_mgr.download_and_convert_factor_file("macro-economy.md")
        if csv_path and csv_path.exists():
            print(f"[OK] Macro factor names refreshed: {csv_path.name}")
            return True
        else:
            print("[WARN] Macro factor names file download failed")
            return False
    except Exception as e:
        print(f"[WARN] Macro factor file refresh error: {e}", file=sys.stderr)
        return False


def main():
    """Main function"""
    print_header()

    license_ok = check_license()
    print(file=sys.stderr)

    if not license_ok:
        print(file=sys.stderr)
        print("[FAIL] License check failed - skill cannot be used", file=sys.stderr)
        return 1

    docs_expired = check_api_docs_expired()
    if docs_expired:
        print("[INFO] API docs expired or missing, refreshing all...")
        cache_ok = run_cache_init(force_refresh=True)
    else:
        cache_ok = True
    print(file=sys.stderr)

    indices_ok = False
    if docs_expired and cache_ok:
        print("[INFO] Regenerating API indices...")
        indices_ok = generate_api_indices()
    elif not docs_expired:
        indices_ok = True
    print(file=sys.stderr)

    macro_ok = False
    if docs_expired and cache_ok:
        print("[INFO] Regenerating macro factor file...")
        macro_ok = refresh_macro_factor_file()
    elif not docs_expired:
        macro_ok = True
    print(file=sys.stderr)

    code_index_ok = build_code_indices()
    print(file=sys.stderr)

    if license_ok and cache_ok and indices_ok and macro_ok and code_index_ok:
        print("Done")
        return 0
    elif license_ok and not cache_ok:
        print("[FAIL] Cache init failed - skill cannot be used", file=sys.stderr)
        return 2
    elif license_ok and cache_ok and not indices_ok:
        print(
            "[WARN] API indices generation failed - skill may still work",
            file=sys.stderr,
        )
        return 0
    else:
        print("[FAIL] Initialization failed - skill cannot be used", file=sys.stderr)
        return 3


if __name__ == "__main__":
    sys.exit(main())