backend/data/market_data.py

"""
Market Data — yfinance wrapper with smart Parquet caching.
Fetches OHLCV data for individual tickers and batches with rate limiting.
"""
import logging
import time
from datetime import datetime, timedelta

import pandas as pd
import yfinance as yf

from backend.data.store import get_store
from config import PARQUET_DIR

logger = logging.getLogger(__name__)

# Rate limiting: max 5 requests per second
_last_request_time = 0.0
_REQUEST_INTERVAL = 0.25  # 250ms between requests


def _rate_limit():
    global _last_request_time
    elapsed = time.time() - _last_request_time
    if elapsed < _REQUEST_INTERVAL:
        time.sleep(_REQUEST_INTERVAL - elapsed)
    _last_request_time = time.time()


def fetch_ohlcv(ticker: str, period: str = "1y", force_refresh: bool = False) -> pd.DataFrame:
    """
    Fetch OHLCV data for a single ticker with Parquet caching.
    Only fetches new data since last cached date (incremental).

    Args:
        ticker: Stock symbol (e.g., 'TCS.NS', 'AAPL')
        period: yfinance period string (default '1y')
        force_refresh: If True, re-download everything

    Returns:
        DataFrame with columns: Open, High, Low, Close, Volume
    """
    store = get_store()

    if not force_refresh:
        cached = store.load_ohlcv(ticker)
        if cached is not None and not cached.empty:
            last_date = cached.index[-1]
            # If data is recent enough (within 1 day for weekdays), return cache
            days_stale = (datetime.now() - last_date.to_pydatetime().replace(tzinfo=None)).days
            if days_stale <= 1:
                logger.debug(f"{ticker}: Using cached data ({len(cached)} rows, last: {last_date.date()})")
                return cached

            # Incremental: fetch only new data
            start_date = (last_date + timedelta(days=1)).strftime("%Y-%m-%d")
            try:
                _rate_limit()
                new_data = yf.download(
                    ticker, start=start_date, progress=False,
                    auto_adjust=True, threads=False
                )
                if new_data is not None and not new_data.empty:
                    # Handle MultiIndex columns from yfinance
                    if isinstance(new_data.columns, pd.MultiIndex):
                        new_data.columns = new_data.columns.get_level_values(0)
                    combined = pd.concat([cached, new_data])
                    combined = combined[~combined.index.duplicated(keep='last')]
                    combined.sort_index(inplace=True)
                    store.save_ohlcv(ticker, combined)
                    logger.info(f"{ticker}: Updated with {len(new_data)} new rows (total: {len(combined)})")
                    return combined
                else:
                    logger.debug(f"{ticker}: No new data since {last_date.date()}")
                    return cached
            except Exception as e:
                logger.warning(f"{ticker}: Incremental update failed ({e}), using cache")
                return cached

    # Full download
    try:
        _rate_limit()
        logger.info(f"{ticker}: Downloading {period} of OHLCV data...")
        df = yf.download(ticker, period=period, progress=False, auto_adjust=True, threads=False)

        if df is None or df.empty:
            logger.warning(f"{ticker}: No data returned from yfinance")
            return pd.DataFrame()

        # Handle MultiIndex columns from yfinance
        if isinstance(df.columns, pd.MultiIndex):
            df.columns = df.columns.get_level_values(0)

        # Cache to Parquet
        store.save_ohlcv(ticker, df)
        logger.info(f"{ticker}: Cached {len(df)} rows ({df.index[0].date()} to {df.index[-1].date()})")
        return df

    except Exception as e:
        logger.error(f"{ticker}: Failed to fetch OHLCV: {e}")
        # Try returning stale cache as fallback
        cached = store.load_ohlcv(ticker)
        if cached is not None:
            logger.warning(f"{ticker}: Serving stale cache as fallback")
            return cached
        return pd.DataFrame()


def fetch_batch(tickers: list[str], period: str = "1y") -> dict[str, pd.DataFrame]:
    """
    Fetch OHLCV for multiple tickers with rate limiting.
    Returns dict of {ticker: DataFrame}.
    """
    results = {}
    total = len(tickers)

    for i, ticker in enumerate(tickers):
        try:
            df = fetch_ohlcv(ticker, period=period)
            if not df.empty:
                results[ticker] = df
            if (i + 1) % 10 == 0:
                logger.info(f"Batch progress: {i + 1}/{total} tickers fetched")
        except Exception as e:
            logger.error(f"{ticker}: Batch fetch error: {e}")

    logger.info(f"Batch complete: {len(results)}/{total} tickers successful")
    return results


def get_latest_price(ticker: str) -> dict:
    """Get the latest available price data for a ticker."""
    df = fetch_ohlcv(ticker, period="5d")
    if df.empty:
        return {}
    last = df.iloc[-1]
    prev = df.iloc[-2] if len(df) > 1 else last
    return {
        "ticker": ticker,
        "close": round(float(last["Close"]), 2),
        "prev_close": round(float(prev["Close"]), 2),
        "change_pct": round((float(last["Close"]) - float(prev["Close"])) / float(prev["Close"]) * 100, 2),
        "volume": int(last["Volume"]),
        "high": round(float(last["High"]), 2),
        "low": round(float(last["Low"]), 2),
        "date": str(last.name.date()),
    }


def is_indian_ticker(ticker: str) -> bool:
    """Check if ticker is an Indian stock (NSE/BSE)."""
    return ticker.endswith(".NS") or ticker.endswith(".BO")