data/fetcher.py

"""
Data fetcher using yfinance (Yahoo Finance) — supports all Taiwan stocks
with .TW (TWSE) and .TWO (TPEX/OTC) suffixes, matching the data shown on
https://tw.stock.yahoo.com/quote/7856.TWO
"""

import logging
import pickle
import subprocess
import sys
from pathlib import Path
from typing import Optional

import pandas as pd
import yfinance as yf

from data.fund_fetcher import is_moneydj_fund, fetch_fund_history, get_fund_info as _get_fund_info_moneydj

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Exchange detection
# ---------------------------------------------------------------------------

_exchange_map: dict[str, str] | None = None


def _build_exchange_map() -> dict[str, str]:
    """Build a stock_no -> exchange mapping from twstock.codes (cached)."""
    global _exchange_map
    if _exchange_map is not None:
        return _exchange_map
    codes = _get_twstock_codes()
    if not codes:
        _exchange_map = {}
        return _exchange_map
    mapping: dict[str, str] = {}
    for code, info in codes.items():
        market = getattr(info, "market", "")
        if market in ("上櫃", "興櫃"):
            mapping[code] = "TPEX"
        else:
            mapping[code] = "TWSE"
    _exchange_map = mapping
    return _exchange_map


def detect_exchange(stock_no: str) -> str:
    """Return 'TPEX' or 'TWSE' based on twstock.codes, suffix, or heuristic."""
    if stock_no.endswith(".TWO"):
        return "TPEX"
    if stock_no.endswith(".TW"):
        return "TWSE"
    bare = stock_no.strip()
    # Exact lookup via twstock.codes — avoids heuristic misses like 4169
    exchange_map = _build_exchange_map()
    if bare in exchange_map:
        return exchange_map[bare]
    # TW ETF with letter suffix (e.g. 00772B) or 6-digit code — always TWSE
    if bare and bare[0].isdigit() and any(c.isalpha() for c in bare):
        return "TWSE"
    if len(bare) == 6 and bare.isdigit():
        return "TWSE"
    # Fallback heuristic for stocks not in twstock database
    if len(bare) == 5 and bare.isdigit():
        return "TPEX"
    if len(bare) == 4 and bare.isdigit():
        v = int(bare)
        if 6000 <= v <= 9999:
            return "TPEX"
    return "TWSE"


def is_us_ticker(stock_no: str) -> bool:
    """Return True if this looks like a US stock/ETF ticker (contains letters, no TW suffix).
    Taiwan ETF codes starting with digits (e.g. 00772B, 006206) are NOT US tickers.
    """
    if stock_no.endswith(".TW") or stock_no.endswith(".TWO"):
        return False
    bare = stock_no.strip()
    # Any code starting with a digit is a Taiwan stock/ETF (may have letter suffix like 00772B)
    if bare and bare[0].isdigit():
        return False
    return any(c.isalpha() for c in bare)


def _normalize_raw_df(raw: pd.DataFrame) -> pd.DataFrame:
    """Flatten and normalize a raw yfinance DataFrame to standard OHLCV format."""
    if isinstance(raw.columns, pd.MultiIndex):
        raw.columns = [col[0].lower() if isinstance(col, tuple) else col.lower()
                       for col in raw.columns]
    else:
        raw.columns = [c.lower() for c in raw.columns]

    df = raw.reset_index()
    df = df.rename(columns={"Date": "date", "index": "date"})
    df["date"] = pd.to_datetime(df["date"]).dt.date

    col_map = {}
    for col in df.columns:
        cl = col.lower()
        if cl == "date":
            col_map[col] = "date"
        elif cl in ("close", "adj close", "adj_close"):
            col_map[col] = "close"
        elif cl == "open":
            col_map[col] = "open"
        elif cl == "high":
            col_map[col] = "high"
        elif cl == "low":
            col_map[col] = "low"
        elif cl == "volume":
            col_map[col] = "volume"

    df = df.rename(columns=col_map)
    needed = ["date", "open", "high", "low", "close", "volume"]
    df = df[[c for c in needed if c in df.columns]]
    df = df.dropna(subset=["close"])
    df = df.sort_values("date").reset_index(drop=True)
    df = df.drop_duplicates(subset=["date"])
    for col in ["open", "high", "low", "close", "volume"]:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce")
    return df


def _to_yf_ticker(stock_no: str) -> str:
    """
    Convert bare stock number or suffixed symbol to yfinance ticker.
    Examples:
        '7856'      -> '7856.TWO'  (TPEX heuristic)
        '7856.TWO'  -> '7856.TWO'
        '2330.TW'   -> '2330.TW'
        '2330'      -> '2330.TW'   (TWSE heuristic)
    """
    if stock_no.endswith(".TW") or stock_no.endswith(".TWO"):
        return stock_no
    exchange = detect_exchange(stock_no)
    suffix = ".TWO" if exchange == "TPEX" else ".TW"
    return stock_no + suffix


# ---------------------------------------------------------------------------
# History fetcher
# ---------------------------------------------------------------------------

def _download_raw(
    ticker: str,
    period: str,
    *,
    interval: str = "1d",
    timeout_seconds: int = 20,
) -> pd.DataFrame:
    """Download from yfinance; terminate blocked SSL reads after a hard timeout."""

    worker = Path(__file__).with_name("yfinance_download_worker.py")
    try:
        completed = subprocess.run(
            [
                sys.executable,
                str(worker),
                "--ticker",
                ticker,
                "--period",
                period,
                "--interval",
                interval,
                "--request-timeout",
                str(max(1, timeout_seconds - 3)),
            ],
            check=False,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=timeout_seconds,
        )
    except subprocess.TimeoutExpired:
        logger.warning("yfinance download hard timed out for %s (%ss)", ticker, timeout_seconds)
        return pd.DataFrame()
    except Exception as exc:
        logger.warning("yfinance download failed for %s: %s", ticker, exc)
        return pd.DataFrame()
    if completed.returncode != 0:
        detail = completed.stderr.decode("utf-8", errors="replace").strip()
        logger.warning("yfinance download failed for %s: %s", ticker, detail or f"worker exit {completed.returncode}")
        return pd.DataFrame()
    try:
        raw = pickle.loads(completed.stdout)
    except Exception as exc:
        logger.warning("yfinance download decode failed for %s: %s", ticker, exc)
        return pd.DataFrame()
    return raw if raw is not None else pd.DataFrame()


def fetch_history(stock_no: str, months: int = 24) -> pd.DataFrame:
    """
    Fetch historical daily OHLCV for a Taiwan stock via Yahoo Finance.
    Tries the heuristic exchange first, then falls back to the other exchange
    so stocks like 4169 (TPEX) are found even when detected as TWSE.

    Returns a DataFrame with columns:
        date (datetime.date), open, high, low, close, volume (float)
    Sorted ascending by date, no duplicates.
    """
    if is_moneydj_fund(stock_no):
        return fetch_fund_history(stock_no, months)

    if is_us_ticker(stock_no):
        ticker = stock_no.upper().strip()
        period = f"{months}mo"
        raw = _download_raw(ticker, period)
        if raw.empty:
            raise ValueError(f"No data found for {ticker}")
        df = _normalize_raw_df(raw)
        logger.info("Fetched %d rows for US ticker %s", len(df), ticker)
        return df

    bare = stock_no.replace(".TWO", "").replace(".TW", "")
    period = f"{months}mo"

    # Build ordered list of tickers to try: heuristic first, then the other
    primary = _to_yf_ticker(stock_no)
    alt = bare + (".TW" if primary.endswith(".TWO") else ".TWO")
    tickers_to_try = [primary, alt]

    raw = pd.DataFrame()
    used_ticker = primary
    for ticker in tickers_to_try:
        raw = _download_raw(ticker, period)
        if not raw.empty:
            used_ticker = ticker
            logger.info("Fetched data for %s using ticker %s", stock_no, ticker)
            break

    if raw.empty:
        raise ValueError(
            f"No data found for stock {bare} (.TW or .TWO). "
            "The stock may be delisted, newly listed, or the number is incorrect."
        )

    df = _normalize_raw_df(raw)
    logger.info("Fetched %d rows for %s", len(df), used_ticker)
    return df


# ---------------------------------------------------------------------------
# Stock info helpers
# ---------------------------------------------------------------------------

def _resolve_ticker(stock_no: str) -> tuple[str, str]:
    """
    Return (ticker, exchange) by trying both .TWO and .TW suffixes.
    Uses whichever has live price data.
    """
    bare = stock_no.replace(".TWO", "").replace(".TW", "")
    primary = _to_yf_ticker(stock_no)
    alt = bare + (".TW" if primary.endswith(".TWO") else ".TWO")

    for ticker in [primary, alt]:
        try:
            info = yf.Ticker(ticker).info
            if info.get("regularMarketPrice") or info.get("previousClose") or info.get("currentPrice"):
                exchange = "TPEX" if ticker.endswith(".TWO") else "TWSE"
                return ticker, exchange
        except Exception:
            pass
    return primary, detect_exchange(stock_no)


def get_stock_info(stock_no: str) -> dict:
    """
    Return basic stock info: name, exchange, stock_no.
    Uses yfinance Ticker.info for metadata, with exchange auto-detection fallback.
    """
    if is_moneydj_fund(stock_no):
        return _get_fund_info_moneydj(stock_no)

    if is_us_ticker(stock_no):
        ticker = stock_no.upper().strip()
        try:
            info = yf.Ticker(ticker).info
            name = info.get("longName") or info.get("shortName") or ticker
            exchange = info.get("exchange") or "US"
            return {"stock_no": ticker, "name": name, "exchange": exchange, "ticker": ticker}
        except Exception as exc:
            logger.warning("Could not fetch info for %s: %s", ticker, exc)
            return {"stock_no": ticker, "name": ticker, "exchange": "US", "ticker": ticker}

    bare = stock_no.replace(".TWO", "").replace(".TW", "")
    ticker, exchange = _resolve_ticker(stock_no)

    # Try twstock first (has accurate Chinese names)
    name = bare
    codes = _get_twstock_codes()
    if bare in codes:
        name = codes[bare].name

    # If no Chinese name found, fall back to yfinance (English name)
    if name == bare:
        try:
            info = yf.Ticker(ticker).info
            name = (
                info.get("longName")
                or info.get("shortName")
                or info.get("displayName")
                or bare
            )
        except Exception as exc:
            logger.warning("Could not fetch yfinance info for %s: %s", ticker, exc)

    return {
        "stock_no": bare,
        "name": name,
        "exchange": exchange,
        "ticker": ticker,
    }


# ---------------------------------------------------------------------------
# Realtime quote cache — stores last successful quote per stock so that
# non-trading-hours requests (twstock price=None) return stale but valid data.
# ---------------------------------------------------------------------------

from cachetools import TTLCache
_realtime_cache: TTLCache = TTLCache(maxsize=50, ttl=30)


def get_realtime_quote(stock_no: str) -> dict:
    """
    Return real-time quote dict: {price, change, change_pct, open, high, low, source, is_realtime}.
    For Taiwan stocks: uses twstock.realtime (direct TWSE/TPEX API), falls back to yfinance.
    For US stocks: uses yfinance 5-day download (last 2 rows → today's change).
    Returns cached stale data (with is_realtime=False) during non-trading hours.
    Returns empty dict on complete failure.
    """
    def _safe_float(s) -> Optional[float]:
        try:
            v = float(s)
            return v if v > 0 else None
        except (TypeError, ValueError):
            return None

    def _cache_and_return(key: str, result: dict, is_realtime: bool = True) -> dict:
        result["is_realtime"] = is_realtime
        if result.get("price") is not None:
            _realtime_cache[key] = result
        return result

    def _get_stale(key: str) -> dict:
        cached = _realtime_cache.get(key)
        if cached:
            stale = dict(cached)
            stale["is_realtime"] = False
            return stale
        return {}

    try:
        if is_moneydj_fund(stock_no):
            info = _get_fund_info_moneydj(stock_no)
            price = info.get("current_price")
            if price:
                return _cache_and_return(stock_no, {"price": price, "change": None, "change_pct": None, "source": "moneydj"})
            return _get_stale(stock_no) or {}

        if is_us_ticker(stock_no):
            ticker = stock_no.upper().strip()
            raw = _download_raw(ticker, "5d")
            if raw.empty:
                return _get_stale(ticker) or {}
            df = _normalize_raw_df(raw)
            if len(df) < 2:
                return _get_stale(ticker) or {}
            today = float(df.iloc[-1]["close"])
            yesterday = float(df.iloc[-2]["close"])
            change = today - yesterday
            change_pct = change / yesterday * 100 if yesterday else 0
            row = df.iloc[-1]
            return _cache_and_return(ticker, {
                "price": today,
                "change": round(change, 2),
                "change_pct": round(change_pct, 2),
                "open": float(row["open"]) if "open" in row else None,
                "high": float(row["high"]) if "high" in row else None,
                "low": float(row["low"]) if "low" in row else None,
                "source": "yfinance",
            })

        # Taiwan stock — try twstock.realtime first
        bare = stock_no.replace(".TWO", "").replace(".TW", "")
        try:
            import twstock
            data = twstock.realtime.get(bare)
            if data and data.get("success"):
                rt = data.get("realtime", {})
                price = _safe_float(rt.get("latest_trade_price"))
                yest  = _safe_float(rt.get("yesterday_price"))
                if price and yest:
                    change = round(price - yest, 2)
                    change_pct = round(change / yest * 100, 2)
                    return _cache_and_return(bare, {
                        "price": price,
                        "change": change,
                        "change_pct": change_pct,
                        "open": _safe_float(rt.get("open")),
                        "high": _safe_float(rt.get("high")),
                        "low": _safe_float(rt.get("low")),
                        "source": "twstock",
                    })
                elif price:
                    return _cache_and_return(bare, {"price": price, "change": None, "change_pct": None, "source": "twstock"})
                else:
                    # Non-trading hours: price is None — return cached stale data
                    stale = _get_stale(bare)
                    if stale:
                        logger.info("twstock price=None for %s (non-trading hours), using cached quote", bare)
                        return stale
        except Exception as exc:
            logger.warning("twstock.realtime failed for %s: %s", bare, exc)

        # Fallback: yfinance 5-day download
        ticker = _to_yf_ticker(stock_no)
        raw = _download_raw(ticker, "5d")
        if raw.empty:
            alt = bare + (".TW" if ticker.endswith(".TWO") else ".TWO")
            raw = _download_raw(alt, "5d")
        if not raw.empty:
            df = _normalize_raw_df(raw)
            if len(df) >= 2:
                today = float(df.iloc[-1]["close"])
                yesterday = float(df.iloc[-2]["close"])
                change = today - yesterday
                change_pct = change / yesterday * 100 if yesterday else 0
                row = df.iloc[-1]
                return _cache_and_return(bare, {
                    "price": today,
                    "change": round(change, 2),
                    "change_pct": round(change_pct, 2),
                    "open": float(row["open"]) if "open" in row else None,
                    "high": float(row["high"]) if "high" in row else None,
                    "low": float(row["low"]) if "low" in row else None,
                    "source": "yfinance",
                })
            elif len(df) == 1:
                return _cache_and_return(bare, {"price": float(df.iloc[-1]["close"]), "change": None, "change_pct": None, "source": "yfinance"})

        # All live sources failed — try stale cache
        return _get_stale(bare)
    except Exception as exc:
        logger.warning("get_realtime_quote error for %s: %s", stock_no, exc)
    return _get_stale(stock_no.replace(".TWO", "").replace(".TW", ""))


def get_current_price(stock_no: str) -> Optional[float]:
    """Return the most recent closing price (delegates to get_realtime_quote)."""
    q = get_realtime_quote(stock_no)
    return q.get("price") if q else None


# ---------------------------------------------------------------------------
# Search helper  (TPEX + TWSE company lists via yfinance lookup)
# ---------------------------------------------------------------------------

_twstock_codes_cache = None


def _get_twstock_codes():
    """Lazily load twstock codes dict; cached globally after first call."""
    global _twstock_codes_cache
    if _twstock_codes_cache is not None:
        return _twstock_codes_cache
    import concurrent.futures
    try:
        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(__import__, "twstock")
            twstock = future.result(timeout=5)
            _twstock_codes_cache = twstock.codes
            return _twstock_codes_cache
    except Exception as exc:
        logger.warning("twstock unavailable: %s", exc)
        return {}


def search_stocks(query: str) -> list[dict]:
    """
    Search Taiwan stocks by Chinese name or stock number.
    Uses twstock local database for Chinese name matching (fast, accurate),
    falls back to yfinance for English / number queries.
    """
    q = query.strip()
    results = []

    # ── Chinese name search via twstock ──────────────────────────────────────
    has_chinese = any('\u4e00' <= c <= '\u9fff' for c in q)
    if has_chinese:
        codes = _get_twstock_codes()
        VALID_TYPES = {'股票', 'ETF'}
        q_lower = q.lower()
        for code, info in codes.items():
            if info.type not in VALID_TYPES:
                continue
            if q_lower not in info.name.lower() and q_lower not in code.lower():
                continue
            exchange = 'TPEX' if info.market in ('上櫃', '興櫃') else 'TWSE'
            results.append({'stock_no': code, 'name': info.name, 'exchange': exchange})
            if len(results) >= 20:
                break
        return results

    # ── Number search: try twstock first, then yfinance ───────────────────────
    if q.isdigit():
        codes = _get_twstock_codes()
        if q in codes:
            info = codes[q]
            exchange = 'TPEX' if info.market in ('上櫃', '興櫃') else 'TWSE'
            results.append({'stock_no': q, 'name': info.name, 'exchange': exchange})
        if not results:
            for suffix, exchange in [('.TWO', 'TPEX'), ('.TW', 'TWSE')]:
                try:
                    import yfinance as yf
                    info_yf = yf.Ticker(q + suffix).info
                    if info_yf.get('regularMarketPrice') or info_yf.get('previousClose'):
                        name = info_yf.get('longName') or info_yf.get('shortName') or q
                        results.append({'stock_no': q, 'name': name, 'exchange': exchange})
                        break
                except Exception:
                    pass
        return results[:5]

    # ── English keyword: yfinance search ─────────────────────────────────────
    try:
        import yfinance as yf
        for item in yf.Search(q + ' Taiwan', max_results=10).quotes:
            sym = item.get('symbol', '')
            if not (sym.endswith('.TW') or sym.endswith('.TWO')):
                continue
            bare = sym.replace('.TWO', '').replace('.TW', '')
            exchange = 'TPEX' if sym.endswith('.TWO') else 'TWSE'
            name = item.get('longname') or item.get('shortname') or bare
            results.append({'stock_no': bare, 'name': name, 'exchange': exchange})
    except Exception as exc:
        logger.warning("yfinance search error: %s", exc)

    return results[:20]


def search_us(query: str) -> list[dict]:
    """Search US stocks and ETFs via yfinance."""
    try:
        results = []
        for item in yf.Search(query, max_results=15).quotes:
            sym = item.get("symbol", "")
            if not sym or sym.endswith(".TW") or sym.endswith(".TWO"):
                continue
            name = item.get("longname") or item.get("shortname") or sym
            exchange = item.get("exchange") or "US"
            results.append({"stock_no": sym, "name": name, "exchange": exchange})
        return results[:10]
    except Exception as exc:
        logger.warning("US search error: %s", exc)
        return []


# ---------------------------------------------------------------------------
# Cross-asset data for Taiwan stocks (TAIEX + USD/TWD)
# ---------------------------------------------------------------------------

_cross_asset_cache: dict[str, tuple] = {}  # key → (series, fetched_at)
_CROSS_ASSET_CACHE_TTL = 3600  # 1 hour

def fetch_cross_asset_tw(start_date: str, end_date: str) -> tuple:
    """
    Fetch TAIEX (^TWII), USD/TWD (TWD=X), SOX (^SOX), and US 10Y yield (^TNX)
    daily close prices via yfinance.

    Returns:
        (taiex_close, usdtwd_close, sox_close, tnx_close) — pd.Series indexed
        by 'YYYY-MM-DD' strings. Any element may be None if the fetch fails.
    """
    import time as _time
    cache_key = f"{start_date}:{end_date}"
    cached = _cross_asset_cache.get(cache_key)
    if cached is not None:
        data, fetched_at = cached
        if _time.time() - fetched_at < _CROSS_ASSET_CACHE_TTL:
            return data

    taiex_close = None
    usdtwd_close = None
    sox_close = None
    tnx_close = None
    try:
        raw = yf.download(
            ["^TWII", "TWD=X", "^SOX", "^TNX"],
            start=start_date, end=end_date,
            auto_adjust=True, progress=False,
        )
        if not raw.empty:
            close = raw["Close"] if "Close" in raw.columns else raw.get("close")
            if close is not None and not close.empty:
                close.index = pd.to_datetime(close.index).strftime("%Y-%m-%d")
                if "^TWII" in close.columns:
                    taiex_close = close["^TWII"].dropna()
                if "TWD=X" in close.columns:
                    usdtwd_close = close["TWD=X"].dropna()
                if "^SOX" in close.columns:
                    sox_close = close["^SOX"].dropna()
                if "^TNX" in close.columns:
                    tnx_close = close["^TNX"].dropna()
    except Exception as exc:
        logger.warning("fetch_cross_asset_tw failed: %s", exc)

    result = (taiex_close, usdtwd_close, sox_close, tnx_close)
    _cross_asset_cache[cache_key] = (result, _time.time())
    return result