# sanguo_data/datafeed.py
import pandas as pd
import urllib.request
import json
import time
import logging
from datetime import datetime, timedelta
from multiprocessing import Process, Queue
from typing import Optional
from sanguo_data.config import DataConfig

logger = logging.getLogger(__name__)

def fetch_with_fallback(symbol, start, end, sources: list[str]) -> pd.DataFrame:
    fetchers = {
        "eastmoney": _fetch_eastmoney,
        "baostock": lambda s, st, e: _fetch_baostock_with_timeout(s, st, e, timeout=30),
        "tencent": _fetch_tencent,
    }
    last_err = None
    for name in sources:
        try:
            df = fetchers[name](symbol, start, end)
            if df is not None and len(df) > 0:
                return df
        except Exception as e:
            last_err = e
            continue
    raise RuntimeError(f"all sources failed: {last_err}")

def fetch_daily(symbol, start, end, cfg: DataConfig) -> pd.DataFrame:
    sources = [s["name"] for s in cfg.data_sources.get("daily", []) if s.get("enabled", True)]
    return fetch_with_fallback(symbol, start, end, sources)

# Worker function for multiprocessing (must be at module level to be picklable)
def _baostock_worker(symbol, start, end, result_queue):
    try:
        result = _fetch_baostock_raw(symbol, start, end)
        result_queue.put(result)
    except Exception as e:
        result_queue.put(e)

# Test helper for timeout testing (simulates hanging BaoStock call)
def _hanging_worker_for_test(symbol, start, end, result_queue):
    """Test helper: simulates a hanging BaoStock call (60s sleep)"""
    import time as _time
    _time.sleep(60)
    result_queue.put(pd.DataFrame({"date": ["2026-01-01"], "open": [10.0]}))

def _fetch_baostock_with_timeout(symbol, start, end, timeout=30):
    """子进程隔离 BaoStock（修复 v1 无超时卡死坑）"""
    result_queue = Queue()
    p = Process(target=_baostock_worker, args=(symbol, start, end, result_queue))
    p.start()
    p.join(timeout)
    if p.is_alive():
        p.terminate()
        p.join()
        raise TimeoutError(f"baostock timeout after {timeout}s")

    res = result_queue.get()
    if isinstance(res, Exception):
        raise res
    return res

def _get_em_secid(code: str) -> str:
    if code.startswith(("60", "68", "51")):
        return f"1.{code}"
    return f"0.{code}"

def _parse_em_klines(klines: list) -> Optional[pd.DataFrame]:
    """解析东方财富K线数据（日线和15min通用）"""
    if not klines:
        return None
    rows = []
    for line in klines:
        parts = line.split(",")
        if len(parts) < 7:
            continue
        rows.append({
            "date": parts[0],
            "open": float(parts[1]),
            "close": float(parts[2]),
            "high": float(parts[3]),
            "low": float(parts[4]),
            "volume": float(parts[5]),
            "amount": float(parts[6]),
        })
    if not rows:
        return None
    return pd.DataFrame(rows)

def _fetch_baostock_raw(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
    """BaoStock日线：全量历史，无反爬，amount真实，T+1延迟

    Copied from v1 data_platform/daily_all_update.py:fetch_baostock_daily (lines 242-270)
    """
    try:
        import baostock as bs
    except ImportError:
        return None

    # 转换代码格式：600000 -> sh.600000
    code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "")
    if code.startswith(("60", "68", "51")):
        bs_code = f"sh.{code}"
    else:
        bs_code = f"sz.{code}"

    try:
        rs = bs.query_history_k_data_plus(
            bs_code,
            "date,open,high,low,close,volume,amount",
            start_date=start_date.replace("-", ""),
            end_date=end_date.replace("-", ""),
            frequency="d",
            adjustflag="2",
        )
        rows = []
        while (rs.error_code == "0") and rs.next():
            rows.append(rs.get_row_data())
        if not rows:
            return None
        df = pd.DataFrame(rows, columns=["date", "open", "high", "low", "close", "volume", "amount"])
        for c in ["open", "high", "low", "close", "volume", "amount"]:
            df[c] = pd.to_numeric(df[c], errors="coerce")
        df = df.dropna(subset=["close"])
        if df.empty:
            return None
        return df
    except Exception as e:
        logger.debug("BaoStock日线失败 %s: %s", symbol, e)
        return None

def _fetch_eastmoney(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
    """东方财富日线：当天实时，amount真实，4s限频

    Copied from v1 data_platform/daily_all_update.py:fetch_eastmoney_daily (lines 339-374)
    """
    try:
        import requests as _requests
    except ImportError:
        return None

    code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "")
    secid = _get_em_secid(code)
    ts = str(int(time.time() * 1000))
    url = (
        f"https://push2his.eastmoney.com/api/qt/stock/kline/get?"
        f"secid={secid}&klt=101&fqt=1&"
        f"beg={start_date.replace('-', '')}&end={end_date.replace('-', '')}&"
        f"fields1=f1,f2,f3,f4,f5,f6,f7,f8&"
        f"fields2=f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61&"
        f"ut=b2884a393a59ad64002292a3e90d46a5&lmt=10000&"
        f"cb=jQuery_em_{ts}&_={ts}"
    )

    headers_em = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/124.0.0.0 Safari/537.36",
        "Referer": "https://quote.eastmoney.com/",
        "Accept": "*/*",
        "Accept-Language": "zh-CN,zh;q=0.9",
    }

    session = _requests.Session()
    session.trust_env = False
    try:
        r = session.get(url, headers=headers_em, timeout=15, verify=False)
        if r.status_code != 200:
            return None
        text = r.text
        data = json.loads(text[text.index("(") + 1:text.rindex(")")])
        if data.get("rc") != 0:
            return None
        klines = data.get("data", {}).get("klines", [])
        df = _parse_em_klines(klines)
        if df is None:
            return None
        df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d")
        mask = (df["date"] >= start_date) & (df["date"] <= end_date)
        result = df.loc[mask, ["date", "open", "high", "low", "close", "volume", "amount"]]
        return result if not result.empty else None
    except Exception as e:
        logger.debug("东方财富日线失败 %s: %s", symbol, e)
        return None

def _fetch_tencent(symbol: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
    """腾讯日线：amount有时为0

    Copied from v1 data_platform/fallback.py:_fetch_tencent_daily (lines 66-104)
    """
    code = symbol.replace("SH", "").replace("SZ", "").replace("sh", "").replace("sz", "")
    if code.startswith(("6", "5", "1")):
        prefix = "sh"
    else:
        prefix = "sz"
    tq_symbol = f"{prefix}{code}"

    days = (datetime.strptime(end_date, "%Y-%m-%d") - datetime.strptime(start_date, "%Y-%m-%d")).days + 10
    url = f"https://web.ifzq.gtimg.cn/appstock/app/fqkline/get?param={tq_symbol},day,{start_date},,{days},"

    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        with opener.open(req, timeout=10) as r:
            resp = json.loads(r.read())
        d = resp.get("data")
        if not isinstance(d, dict):
            return None
        klines = d.get(tq_symbol, {}).get("day", [])
        if not klines:
            return None
        df = pd.DataFrame(klines)
        ncols = len(df.columns)
        if ncols >= 7:
            df.columns = ["date", "open", "close", "high", "low", "volume", "amount"][:ncols]
        else:
            df.columns = ["date", "open", "close", "high", "low", "volume"][:ncols]
        if "amount" not in df.columns:
            df["amount"] = 0.0
        for c in ["open", "close", "high", "low", "volume", "amount"]:
            df[c] = pd.to_numeric(df[c], errors="coerce").fillna(0)
        df["date"] = pd.to_datetime(df["date"]).dt.strftime("%Y-%m-%d")
        mask = (df["date"] >= start_date) & (df["date"] <= end_date)
        return df.loc[mask, ["date", "open", "high", "low", "close", "volume", "amount"]]
    except Exception as e:
        logger.debug("腾讯日线失败 %s: %s", symbol, e)
        return None