from __future__ import annotations

import importlib
import os
from datetime import datetime

import pandas as pd

from market.provider_base import MarketDataProvider


class CmesCnProvider(MarketDataProvider):
    def __init__(self, token: str | None = None) -> None:
        self._token = token or os.getenv("CMES_TOKEN", "").strip()
        self._module = importlib.import_module("cmesdata")
        self._login_once()

    @property
    def provider_name(self) -> str:
        return "cmesdata"

    @property
    def channel(self) -> str:
        return "cn"

    def _login_once(self) -> None:
        if not self._token:
            raise RuntimeError("CMES_TOKEN 未配置，无法使用 cmesdata 通道")
        self._module.login(self._token)

    @staticmethod
    def _to_prefixed_code(code: str) -> str:
        raw = code.strip().upper().replace(".", "")
        if raw.startswith("SH") or raw.startswith("SZ"):
            return f"{raw[:2]}.{raw[2:]}"
        if raw.isdigit() and len(raw) == 6:
            if raw.startswith(("6", "9")):
                return f"SH.{raw}"
            return f"SZ.{raw}"
        raise RuntimeError("cmesdata 通道仅支持 6 位 A 股代码或 SH./SZ. 前缀代码")

    def fetch_spot(self) -> pd.DataFrame:
        raise RuntimeError("cmesdata 不支持全市场快照拉取，请通过精确代码查询")

    def search_spot(self, query: str, limit: int) -> pd.DataFrame:
        _ = limit
        code = self._to_prefixed_code(query)
        df = self._module.get_real_hq([code])
        if df is None or df.empty:
            raise RuntimeError(f"未获取到实时行情: {code}")
        return df

    def fetch_daily_kline(self, code: str, start: datetime, end: datetime) -> pd.DataFrame:
        prefixed = self._to_prefixed_code(code)
        df = self._module.get_history_data(
            prefixed,
            start.strftime("%Y-%m-%d"),
            end.strftime("%Y-%m-%d"),
            "D",
        )
        if df is None or df.empty:
            raise RuntimeError(f"未获取到历史 K 线: {prefixed}")
        frame = pd.DataFrame(
            {
                "date": pd.to_datetime(df["时间"]),
                "close": pd.to_numeric(df["收盘价"], errors="coerce"),
                "volume": pd.to_numeric(df["成交量"], errors="coerce"),
            }
        ).dropna()
        return frame.sort_values("date").reset_index(drop=True)