"""Simplified Tushare client with rate limiting and retry logic."""

import time
import pandas as pd
import tushare as ts

from typing import Optional
from src.data.rate_limiter import TokenBucketRateLimiter
from src.data.sync_config import get_rate_limit
from src.config.settings import get_settings


class TushareClient:
    """Tushare API client with rate limiting and retry."""

    # 类级别限流器缓存（按接口名称存储）
    _limiter_cache: dict[str, TokenBucketRateLimiter] = {}
    _cached_rate_limits: dict[str, int] = {}

    def __init__(
        self, token: Optional[str] = None, interface_name: Optional[str] = None
    ):
        """Initialize client.

        Args:
            token: Tushare API token (auto-loaded from config if not provided)
            interface_name: 接口名称，用于获取特定的速率限制配置
        """
        cfg = get_settings()
        token = token or cfg.tushare_token

        if not token:
            raise ValueError("Tushare token is required")

        self.token = token
        self.config = cfg
        self.interface_name = interface_name or "default"

        # 获取或创建限流器
        self.rate_limiter = self._get_or_create_limiter(self.interface_name)

        self._api = None

    def _get_or_create_limiter(self, interface_name: str) -> TokenBucketRateLimiter:
        """获取或创建指定接口的限流器。

        如果接口的 rate_limit 配置发生变化，会重新创建限流器。

        Args:
            interface_name: 接口名称

        Returns:
            TokenBucketRateLimiter 实例
        """
        # 获取当前配置的 rate_limit
        current_rate_limit = get_rate_limit(interface_name)

        # 检查是否需要创建新的限流器
        if (
            interface_name not in TushareClient._limiter_cache
            or TushareClient._cached_rate_limits.get(interface_name)
            != current_rate_limit
        ):
            # 创建新的限流器
            TushareClient._limiter_cache[interface_name] = TokenBucketRateLimiter(
                rate_limit=current_rate_limit,
            )
            TushareClient._cached_rate_limits[interface_name] = current_rate_limit
            min_interval = 60.0 / current_rate_limit
            print(
                f"[TushareClient] Initialized rate limiter for '{interface_name}': "
                f"rate={current_rate_limit}/min, interval={min_interval:.2f}s"
            )

        return TushareClient._limiter_cache[interface_name]

    def get_current_rate_limit(self) -> int:
        """获取当前接口的速率限制。

        Returns:
            每分钟请求数
        """
        return get_rate_limit(self.interface_name)

    def _get_api(self):
        """Get Tushare API instance."""
        if self._api is None:
            self._api = ts.pro_api(self.token)
        return self._api

    def query(self, api_name: str, timeout: float = None, **params) -> pd.DataFrame:
        """Execute API query with rate limiting and retry.

        Args:
            api_name: API name ('daily', 'pro_bar', etc.)
            timeout: Timeout for rate limiting (None = wait indefinitely)
            **params: API parameters

        Returns:
            DataFrame with query results
        """
        timeout = timeout if timeout is not None else float("inf")
        max_retries = 3
        retry_delays = [1, 3, 10]

        for attempt in range(max_retries):
            # Acquire rate limit token before each attempt (including retries)
            success, wait_time = self.rate_limiter.acquire(timeout=timeout)

            if not success:
                raise RuntimeError(f"Rate limit exceeded after {timeout}s timeout")

            try:
                # pro_bar uses ts.pro_bar() instead of api.query()
                if api_name == "pro_bar":
                    # pro_bar parameters: ts_code, start_date, end_date, adj, freq, factors, ma, adjfactor
                    data = ts.pro_bar(
                        ts_code=params.get("ts_code"),
                        start_date=params.get("start_date"),
                        end_date=params.get("end_date"),
                        adj=params.get("adj"),
                        freq=params.get("freq", "D"),
                        factors=params.get(
                            "factors"
                        ),  # factors should be a list like ['tor', 'vr']
                        ma=params.get("ma"),
                        adjfactor=params.get("adjfactor"),
                    )
                    # Handle None response (e.g., delisted stock)
                    if data is None:
                        data = pd.DataFrame()
                else:
                    api = self._get_api()
                    data = api.query(api_name, **params)

                return data

            except Exception as e:
                error_msg = str(e)
                if attempt < max_retries - 1:
                    delay = retry_delays[attempt]
                    # 如果触发 Tushare 限流，增加等待时间避开惩罚期
                    if "最多访问该接口" in error_msg:
                        delay = max(delay, 60)
                        print(
                            f"[RateLimit] {api_name} hit Tushare limit, waiting {delay}s..."
                        )

                    print(
                        f"[Retry] {api_name} failed (attempt {attempt + 1}): {error_msg}, retry in {delay}s"
                    )
                    time.sleep(delay)
                else:
                    raise RuntimeError(
                        f"API call failed after {max_retries} attempts: {e}"
                    )

        return pd.DataFrame()

    def close(self):
        """Close client."""
        pass
