"""麦蕊智数（mairui）数据源。

提供：
  - 日 K 线（`hsstock/history/{code}.{ex}/d/n/{licence}`）
  - 5 分钟 K 线（`hsstock/history/{code}.{ex}/5/n/{licence}`）
  - 指数日 K 线（`hsindex/history/{code}.{ex}/d/{licence}`）
  - 资金流向（`hsstock/history/transaction/{code}.{ex}/{licence}`）

限速：1分钟300次（默认保守到 5 RPS = 1分钟300次）
凭证：licence（无需登录态，从环境变量 MAIRUI_LICENCE 读取）
"""
from __future__ import annotations

import json
import os
import threading
import time
import urllib.request
from typing import Any, Optional

import pandas as pd

from app.core.datasource.base import DataSource
from app.core.datasource.utils import (
    code6_to_mairui,
    filter_date_range,
    normalize_kline,
)


class MairuiSource(DataSource):
    key = "datasource_mairui"
    name = "麦蕊智数（mairui.club）"
    provides = ["kline_daily", "kline_5min", "index_daily", "stock_basic", "moneyflow"]
    requires_credential = True
    credential_key = "MAIRUI_LICENCE"

    _HEADERS = {
        "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7)",
        "Accept": "application/json",
    }
    _BASE_URL = "https://api.mairuiapi.com"

    # mairui 免费 licence 1 分钟 300 次 = 5 RPS（**单 licence**硬上限）。
    # 5 workers × 1 RPS/worker = 5 RPS 总 = 刚好踩满上限，避免风控。
    # 用 thread-local 计时：每个 worker 独立计自己的 last_ts，
    # 避免 N 个 worker 串行抢一把锁退化成 1 worker。
    _RPS_LIMIT = 1.0
    _rps_local = threading.local()

    @classmethod
    def _wait_rps(cls):
        last = getattr(cls._rps_local, "last_ts", 0.0)
        now = time.time()
        elapsed = now - last
        min_interval = 1.0 / cls._RPS_LIMIT
        if elapsed < min_interval:
            time.sleep(min_interval - elapsed)
            now = time.time()
        cls._rps_local.last_ts = now

    def _get_licence(self) -> Optional[str]:
        return os.environ.get(self.credential_key, "").strip() or None

    def is_available(self) -> tuple[bool, str]:
        if not self._get_licence():
            return False, f"未配置 {self.credential_key}"
        return True, "ok"

    def health_check(self) -> dict[str, Any]:
        lic = self._get_licence()
        if not lic:
            return {"success": False, "message": f"{self.credential_key} 未配置"}
        try:
            # 拿 1 条日线数据作为连通性 + 凭证测试
            url = f"{self._BASE_URL}/hsstock/history/600519.SH/d/n/{lic}?st=20260609&et=20260609"
            req = urllib.request.Request(url, headers=self._HEADERS)
            with urllib.request.urlopen(req, timeout=15) as resp:
                raw = resp.read().decode("utf-8", errors="replace")
            data = json.loads(raw)
            if isinstance(data, list) and data:
                return {"success": True, "message": f"连接成功，获取到 {len(data)} 条 K 线"}
            if isinstance(data, dict) and data.get("error"):
                return {"success": False, "message": f"麦蕊返回: {data['error']}"}
            return {"success": False, "message": f"麦蕊返回空/异常: {raw[:200]}"}
        except Exception as e:
            return {"success": False, "message": f"麦蕊连接失败: {e}"}

    def _fetch(self, path: str, params: dict[str, Any], retry: int = 2) -> Any:
        """通用 GET，含限速 + 重试。"""
        self._wait_rps()
        qs = "&".join(f"{k}={v}" for k, v in params.items() if v)
        url = f"{self._BASE_URL}{path}"
        if qs:
            url = f"{url}?{qs}"
        last_err = None
        for attempt in range(retry + 1):
            try:
                req = urllib.request.Request(url, headers=self._HEADERS)
                with urllib.request.urlopen(req, timeout=20) as resp:
                    # 先按 gbk 试（中文乱码返回），再 fallback utf-8
                    raw_bytes = resp.read()
                    # 用 latin-1 永不失败地把 bytes 转成字符串（每字节 1 字符）
                    raw = raw_bytes.decode("latin-1")
                # mairui 风控时返回 "接收数据异常，请稍后再试"（gbk 编码）
                if "请稍后再试" in raw or "codec can't decode" in raw or raw.startswith("'utf-8'") or "Error -3 while decompressing" in raw:
                    raise RuntimeError(f"mairui 返回风控提示: {raw[:80]}")
                # 试 JSON parse；如果是 gbk 编码的中文提示，要先解码
                if raw.startswith("'") and raw.endswith("'"):
                    # gbk 编码的 Python repr 字符串，如 "'接收数据异常，请稍后再试'"
                    try:
                        decoded = raw[1:-1].encode("latin-1").decode("gbk")
                        if "请稍后再试" in decoded:
                            raise RuntimeError(f"mairui 返回风控: {decoded}")
                    except Exception:
                        pass
                return json.loads(raw)
            except Exception as e:
                last_err = e
                if attempt < retry:
                    time.sleep(1.0 * (attempt + 1))  # 风控重试退避长一点
        from app.core.utils.logging import get_logger
        get_logger("mairui").debug(f"mairui fetch {url} failed: {last_err}")
        return []

    # ── 股票列表 ───────────────────────────────────────────

    def fetch_stock_list(self) -> list[dict]:
        """全市场沪深 A 股基础列表。

        API: GET /hslt/list/{licence}
        返回字段：dm (代码.交易所，如 "000001.SZ"), mc (名称), jys (交易所 SZ/SH)
        """
        lic = self._get_licence()
        if not lic:
            return []
        data = self._fetch(f"/hslt/list/{lic}", {})
        if not isinstance(data, list):
            return []
        rows = []
        for item in data:
            try:
                dm = str(item.get("dm", "")).strip()
                mc = str(item.get("mc", "")).strip()
                # Mairui 在 2 字简称中间填了空格（如 "万 科Ａ"），去掉多余空格
                mc = "".join(mc.split())
                jys = str(item.get("jys", "")).strip().upper()
                # dm 格式: "000001.SZ" → code6="000001", exchange="SZ"
                if "." in dm:
                    code6, exch = dm.split(".", 1)
                else:
                    code6, exch = dm, jys
                if not code6 or len(code6) != 6 or not code6.isdigit():
                    continue
                exchange = exch.upper() or jys
                if exchange not in ("SH", "SZ"):
                    continue
                rows.append({
                    "code": f"{exchange}{code6}",
                    "name": mc,
                    "exchange": exchange,
                    "list_date": "",  # hslt/list 不返回 IPO 日期
                    "listing_status": "st" if (mc.startswith("ST") or mc.startswith("*ST")) else "normal",
                })
            except Exception:
                continue
        return rows

    # ── K 线 fetch ─────────────────────────────────────────

    def fetch_kline_daily(self, code6: str, start: str, end: str) -> pd.DataFrame:
        lic = self._get_licence()
        if not lic:
            return pd.DataFrame()
        symbol = code6_to_mairui(code6)
        path = f"/hsstock/history/{symbol}/d/n/{lic}"
        params = {
            "st": (start or "").replace("-", ""),
            "et": (end or "").replace("-", ""),
        }
        data = self._fetch(path, params)
        if not isinstance(data, list):
            return pd.DataFrame()
        records = []
        for item in data:
            try:
                # t 字段：日线 "2026-06-09 00:00:00"，取日期部分
                t = item.get("t", "")
                d = t.split(" ")[0] if isinstance(t, str) else ""
                if not d:
                    continue
                records.append({
                    "trade_date": d,
                    "open": float(item["o"]),
                    "high": float(item["h"]),
                    "low": float(item["l"]),
                    "close": float(item["c"]),
                    "volume": float(item.get("v") or 0),
                })
            except (KeyError, ValueError, TypeError):
                continue
        return filter_date_range(normalize_kline(pd.DataFrame(records)), start, end)

    def fetch_kline_5min(self, code6: str, start: str, end: str) -> pd.DataFrame:
        """5 分钟 K 线。返回标准列：bar_time, open, high, low, close, volume, amount。

        mairui 用 level=5（数字，不是 "5m"）。
        """
        lic = self._get_licence()
        if not lic:
            return pd.DataFrame()
        symbol = code6_to_mairui(code6)
        path = f"/hsstock/history/{symbol}/5/n/{lic}"
        params = {
            "st": (start or "").replace("-", ""),
            "et": (end or "").replace("-", ""),
        }
        data = self._fetch(path, params)
        if not isinstance(data, list):
            return pd.DataFrame()
        records = []
        for item in data:
            try:
                t = item.get("t", "")
                # 分钟级时间格式 "2026-06-09 14:50:00" → 替换空格为 T 让 pandas 识别
                bar_time = t.replace(" ", "T") if isinstance(t, str) else None
                if not bar_time:
                    continue
                records.append({
                    "bar_time": bar_time,
                    "open": float(item["o"]),
                    "high": float(item["h"]),
                    "low": float(item["l"]),
                    "close": float(item["c"]),
                    "volume": float(item.get("v") or 0),
                    "amount": float(item.get("a") or 0),
                })
            except (KeyError, ValueError, TypeError):
                continue
        if not records:
            return pd.DataFrame()
        df = pd.DataFrame(records)
        df["bar_time"] = pd.to_datetime(df["bar_time"], errors="coerce")
        df = df.dropna(subset=["bar_time"]).sort_values("bar_time").reset_index(drop=True)
        # 过滤日期范围（按日期部分，不含时分）
        if start:
            df = df[df["bar_time"] >= pd.Timestamp(start)]
        if end:
            # end 包含整天
            df = df[df["bar_time"] < pd.Timestamp(end) + pd.Timedelta(days=1)]
        return df

    def fetch_index_daily(self, index_code: str, start: str, end: str) -> pd.DataFrame:
        """指数日 K。index_code 用 mairui 格式（如 '000300.SH'）。"""
        lic = self._get_licence()
        if not lic:
            return pd.DataFrame()
        path = f"/hsindex/history/{index_code}/d/{lic}"
        params = {
            "st": (start or "").replace("-", ""),
            "et": (end or "").replace("-", ""),
        }
        data = self._fetch(path, params)
        if not isinstance(data, list):
            return pd.DataFrame()
        records = []
        for item in data:
            try:
                t = item.get("t", "")
                d = t.split(" ")[0] if isinstance(t, str) else ""
                if not d:
                    continue
                records.append({
                    "trade_date": d,
                    "open": float(item["o"]),
                    "high": float(item["h"]),
                    "low": float(item["l"]),
                    "close": float(item["c"]),
                    "volume": float(item.get("v") or 0),
                })
            except (KeyError, ValueError, TypeError):
                continue
        return filter_date_range(normalize_kline(pd.DataFrame(records)), start, end)

    # ── 资金流向 ───────────────────────────────────────────

    def fetch_moneyflow(self, code6: str, start: str, end: str) -> pd.DataFrame:
        """个股资金流向（主力/大/中/小单 净额）。

        API: GET /hsstock/history/transaction/{code}.{ex}/{licence}?st=YYYYMMDD&et=YYYYMMDD
        返回字段（按买卖方向 × 单型 4×4 矩阵 + 主买/主卖/被动买/被动卖）：
          zmbtdcje 主买特大单成交额
          zmbddcje 主买大单成交额
          zmbzdcje 主买中单成交额
          zmbxdcje 主买小单成交额
          zmstdcje 主卖特大单成交额
          zmsddcje 主卖大单成交额
          zmszdcje 主卖中单成交额
          zmsxdcje 主卖小单成交额
          bdmbtdcje 被动买特大单成交额
          bdmbddcje 被动买大单成交额
          bdmbzdcje 被动买中单成交额
          bdmbxdcje 被动买小单成交额
          bdmstdcje 被动卖特大单成交额
          bdmsddcje 被动卖大单成交额
          bdmszdcje 被动卖中单成交额
          bdmsxdcje 被动卖小单成交额
          ... 以及对应的成交量/笔数字段（zmbtdcjl 等），本接口暂只取成交额

        单型口径（mairui 文档）：
          特大单：成交额 ≥ 100 万 或 成交量 ≥ 5000 手
          大单  ：成交额 ≥ 20 万  或 成交量 ≥ 1000 手
          中单  ：成交额 ≥ 4 万   或 成交量 ≥ 200 手
          小单  ：其他

        输出字段：trade_date, main_net_inflow, large_net_inflow, medium_net_inflow, small_net_inflow
        净额口径（与 akshare stock_individual_fund_flow 保持一致）：
          主力净流入 = Σ主买四型 - Σ主卖四型       (zmb{t,d,z,x} - zms{t,d,z,x})
          大单净额   = 主买大   - 主卖大           (zmbddcje - zmsddcje)
          中单净额   = 主买中   - 主卖中           (zmbzdcje - zmszdcje)
          小单净额   = 主买小   - 主卖小           (zmbxdcje - zmsxdcje)
        注：不能加被动买/卖 —— 主动买 ≡ 被动卖、主动卖 ≡ 被动买
        （同一笔成交记在两边），加起来恒等 0。
        """
        lic = self._get_licence()
        if not lic:
            return pd.DataFrame()
        symbol = code6_to_mairui(code6)
        path = f"/hsstock/history/transaction/{symbol}/{lic}"
        params = {
            "st": (start or "").replace("-", ""),
            "et": (end or "").replace("-", ""),
        }
        data = self._fetch(path, params)
        if not isinstance(data, list):
            return pd.DataFrame()

        def f(item, k) -> float:
            try:
                v = item.get(k)
                return float(v) if v not in (None, "", "-") else 0.0
            except (TypeError, ValueError):
                return 0.0

        records = []
        for item in data:
            t = item.get("t", "")
            d = t.split(" ")[0] if isinstance(t, str) else ""
            if not d:
                continue
            # 主买四型 / 主卖四型
            main_buy = (
                f(item, "zmbtdcje") + f(item, "zmbddcje")
                + f(item, "zmbzdcje") + f(item, "zmbxdcje")
            )
            main_sell = (
                f(item, "zmstdcje") + f(item, "zmsddcje")
                + f(item, "zmszdcje") + f(item, "zmsxdcje")
            )
            # 大/中/小 净额：只看主动方（不包含被动方，否则恒等 0）
            large_net = f(item, "zmbddcje") - f(item, "zmsddcje")
            medium_net = f(item, "zmbzdcje") - f(item, "zmszdcje")
            small_net = f(item, "zmbxdcje") - f(item, "zmsxdcje")
            records.append({
                "trade_date": d,
                "main_net_inflow": round(main_buy - main_sell, 2),
                "large_net_inflow": round(large_net, 2),
                "medium_net_inflow": round(medium_net, 2),
                "small_net_inflow": round(small_net, 2),
            })
        if not records:
            return pd.DataFrame()
        df = pd.DataFrame(records)
        # 日期过滤
        if start:
            df = df[df["trade_date"] >= start]
        if end:
            df = df[df["trade_date"] <= end]
        return df.sort_values("trade_date").reset_index(drop=True)