"""Price collector: daily OHLCV via one batched yfinance download, plus best-effort after-hours prices.

Semantics
---------
* ``close`` is the Close of the last daily bar dated on/before ``session_date`` (selected by date, never
  by position, because the batch frame's index is the union of every ticker's dates and 24h instruments
  such as ``KRW=X``/``BTC-USD``/futures carry bars for days on which US equities did not trade).
* For fx / futures / crypto, which trade ~24h, "close" therefore means the last available daily bar on or
  before the session date under Yahoo's own trading-day convention, not a 16:00 ET print.
* ``prev_close`` is the Close of the preceding non-NaN daily bar; ``change_pct`` and ``change_abs`` derive
  from those two. For the ``rates`` group (``^IRX ^FVX ^TNX ^TYX``) ``change_abs`` = yield points is the
  primary metric (the renderer picks it via ``symbols.POINT_CHANGE_GROUPS``); ``change_pct`` is filled too.
* ``avg_volume_10d`` averages the up-to-10 daily bars *before* the session bar (excludes the session day).
* ``after_hours_price`` is filled only for plain equities/ETFs from 1-minute bars stamped at/after 16:00 ET
  on the session date. It is best-effort: per-ticker try/except and a total time budget.
* ``as_of`` is the exchange-local calendar date of the daily bar actually used, expressed as midnight
  America/New_York, so a stale bar is visible.

Failure policy: the batch download is retried once; if it still fails, a per-ticker ``Ticker.history`` loop
is attempted for the ``indices`` group only; if nothing at all is obtained ``PriceFetchError`` is raised so
``collectors.base.run_safely`` records a data gap. A single missing/bad ticker is skipped and logged.
"""

from __future__ import annotations

import logging
import math
import time
from datetime import date, datetime, timedelta
from datetime import time as dtime
from zoneinfo import ZoneInfo

import pandas as pd
import yfinance as yf

from us2kr import symbols
from us2kr.models import Quote

log = logging.getLogger(__name__)

ET = ZoneInfo("America/New_York")

HISTORY_DAYS = 21  # calendar days of daily bars to pull (>= 10 sessions + holidays margin)
AVG_VOLUME_WINDOW = 10  # sessions used for avg_volume_10d
MIN_AVG_VOLUME_BARS = 3  # fewer prior bars than this -> avg_volume_10d = None
AFTER_HOURS_BUDGET_SECONDS = 45.0  # total wall-clock budget for the after-hours loop
REGULAR_CLOSE_ET = dtime(16, 0)  # after-hours bars are those stamped at/after this local time
DOWNLOAD_TIMEOUT_SECONDS = 20

_REQUIRED_COLS = ("Close",)
_FALLBACK_GROUP = "indices"


class PriceFetchError(RuntimeError):
    """Raised when no price data at all could be obtained (batch + fallback both failed)."""


# --------------------------------------------------------------------------------------------------
# public API
# --------------------------------------------------------------------------------------------------


def fetch_quotes(session_date: date, tickers: list[str] | None = None) -> list[Quote]:
    """Return one Quote per available ticker for the US session ``session_date``.

    ``tickers`` defaults to ``symbols.all_tickers()``. Output order follows the input order. Tickers that
    yield no usable bars are skipped (logged), never raised. Raises ``PriceFetchError`` only when the
    whole batch and the fallback fail.
    """
    wanted = list(tickers) if tickers is not None else symbols.all_tickers()
    if not wanted:
        return []

    frames = _download_history(session_date, wanted)

    quotes: list[Quote] = []
    for ticker in wanted:
        frame = frames.get(ticker)
        if frame is None or frame.empty:
            log.warning("prices: %s missing from download - skipped", ticker)
            continue
        try:
            quote = _quote_from_history(ticker, frame, session_date)
        except Exception:  # one bad ticker must not sink the batch
            log.exception("prices: %s could not be parsed - skipped", ticker)
            continue
        if quote is None:
            log.warning("prices: %s has no bar on/before %s - skipped", ticker, session_date)
            continue
        quotes.append(quote)

    log.info("prices: %d/%d tickers quoted for %s", len(quotes), len(wanted), session_date)
    _fill_after_hours(quotes, session_date)
    return quotes


# --------------------------------------------------------------------------------------------------
# download layer
# --------------------------------------------------------------------------------------------------


def _download_history(session_date: date, tickers: list[str]) -> dict[str, pd.DataFrame]:
    """Batched daily download (retry once), then per-ticker fallback for indices. Never returns empty."""
    start = session_date - timedelta(days=HISTORY_DAYS)
    end = session_date + timedelta(days=1)  # yfinance ``end`` is exclusive
    errors: list[str] = []

    for attempt in (1, 2):
        try:
            raw = yf.download(
                tickers,
                start=start.isoformat(),
                end=end.isoformat(),
                interval="1d",
                group_by="ticker",
                auto_adjust=False,
                threads=True,
                progress=False,
                timeout=DOWNLOAD_TIMEOUT_SECONDS,
            )
        except Exception as exc:  # noqa: BLE001 - network / parsing errors from yfinance
            errors.append(f"batch attempt {attempt}: {type(exc).__name__}: {exc}")
            log.warning("prices: batch download attempt %d failed: %s", attempt, exc)
            continue
        frames = _split_frame(raw, tickers)
        if frames:
            return frames
        errors.append(f"batch attempt {attempt}: empty frame")
        log.warning("prices: batch download attempt %d returned no data", attempt)

    fallback_tickers = [t for t in tickers if symbols.group_of(t) == _FALLBACK_GROUP]
    log.warning("prices: batch failed twice; falling back to per-ticker history for %s", fallback_tickers)
    frames: dict[str, pd.DataFrame] = {}
    for ticker in fallback_tickers:
        try:
            hist = yf.Ticker(ticker).history(
                start=start.isoformat(),
                end=end.isoformat(),
                interval="1d",
                auto_adjust=False,
                actions=False,
                timeout=DOWNLOAD_TIMEOUT_SECONDS,
            )
        except Exception as exc:  # noqa: BLE001
            errors.append(f"{ticker}: {type(exc).__name__}: {exc}")
            log.warning("prices: fallback history for %s failed: %s", ticker, exc)
            continue
        cleaned = _clean_single(hist)
        if cleaned is not None:
            frames[ticker] = cleaned
    if frames:
        return frames
    raise PriceFetchError("no price data obtained; " + "; ".join(errors)[:500])


def _split_frame(raw: pd.DataFrame | None, tickers: list[str]) -> dict[str, pd.DataFrame]:
    """Turn a yf.download result into {ticker: single-level OHLCV frame}, dropping all-NaN tickers."""
    if raw is None or not isinstance(raw, pd.DataFrame) or raw.empty:
        return {}

    frames: dict[str, pd.DataFrame] = {}
    if isinstance(raw.columns, pd.MultiIndex):
        # Expected layout (group_by="ticker"): level0 = ticker, level1 = price field.
        # Defensive: if level0 holds price fields instead (group_by="column" layout), swap.
        level0 = set(raw.columns.get_level_values(0))
        level1 = set(raw.columns.get_level_values(1))
        if "Close" in level0 and "Close" not in level1:
            raw = raw.swaplevel(0, 1, axis=1)
            level0 = set(raw.columns.get_level_values(0))
        for ticker in tickers:
            key = ticker if ticker in level0 else ticker.upper() if ticker.upper() in level0 else None
            if key is None:
                continue
            cleaned = _clean_single(raw[key])
            if cleaned is not None:
                frames[ticker] = cleaned
        return frames

    # Single-level columns: only meaningful when exactly one ticker was requested.
    if len(tickers) == 1:
        cleaned = _clean_single(raw)
        if cleaned is not None:
            frames[tickers[0]] = cleaned
    else:
        log.warning("prices: single-level columns for %d tickers - cannot attribute, ignoring", len(tickers))
    return frames


def _clean_single(frame: pd.DataFrame | None) -> pd.DataFrame | None:
    """Normalise a per-ticker OHLCV frame: naive date index, rows with a Close only; None if unusable."""
    if frame is None or not isinstance(frame, pd.DataFrame) or frame.empty:
        return None
    if isinstance(frame.columns, pd.MultiIndex):
        frame = frame.copy()
        frame.columns = frame.columns.get_level_values(-1)
    if any(col not in frame.columns for col in _REQUIRED_COLS):
        return None
    idx = pd.to_datetime(frame.index)
    if idx.tz is not None:
        # yfinance stamps daily bars at local midnight in the exchange tz (Europe/London for KRW=X/JPY=X,
        # UTC for BTC-USD, America/New_York or America/Chicago for US indices). Converting to ET would move
        # midnight-London/UTC bars to the previous evening ET and relabel them one day early, so just drop
        # the tz and keep the exchange-local date, consistent with the tz-naive batch download.
        idx = idx.tz_localize(None)
    frame = frame.copy()
    frame.index = idx.normalize()
    frame = frame[frame["Close"].notna()]
    if frame.empty:
        return None
    return frame.sort_index()


# --------------------------------------------------------------------------------------------------
# quote computation
# --------------------------------------------------------------------------------------------------


def _quote_from_history(ticker: str, frame: pd.DataFrame, session_date: date) -> Quote | None:
    """Build a Quote from a cleaned daily frame using the last bar on/before ``session_date``."""
    hist = frame[frame.index <= pd.Timestamp(session_date)]
    if hist.empty:
        return None

    row = hist.iloc[-1]
    bar_date: date = hist.index[-1].date()
    if bar_date != session_date:
        log.info("prices: %s has no %s bar; using last available bar %s", ticker, session_date, bar_date)

    close = _num(row.get("Close"))
    if close is None:
        return None
    prev_close = _num(hist.iloc[-2].get("Close")) if len(hist) >= 2 else None

    change_abs = change_pct = None
    if prev_close is not None:
        change_abs = close - prev_close
        if prev_close != 0:
            change_pct = (close / prev_close - 1.0) * 100.0

    volume = _num(row.get("Volume"))
    if volume is not None and volume <= 0:
        volume = None  # fx / rates / indices without volume report 0

    prior_vol = hist["Volume"].iloc[-(AVG_VOLUME_WINDOW + 1) : -1] if "Volume" in hist.columns else None
    avg_volume: float | None = None
    if prior_vol is not None:
        prior_vol = pd.to_numeric(prior_vol, errors="coerce").dropna()
        prior_vol = prior_vol[prior_vol > 0]
        if len(prior_vol) >= MIN_AVG_VOLUME_BARS:
            avg_volume = float(prior_vol.mean())

    return Quote(
        symbol=ticker,
        name=symbols.display_name(ticker),
        group=symbols.group_of(ticker) or "other",
        close=close,
        prev_close=prev_close,
        change_pct=change_pct,
        change_abs=change_abs,
        day_high=_num(row.get("High")),
        day_low=_num(row.get("Low")),
        volume=volume,
        avg_volume_10d=avg_volume,
        as_of=datetime.combine(bar_date, dtime.min, tzinfo=ET),
        source="yfinance",
    )


def _num(value: object) -> float | None:
    """Coerce a pandas scalar to float, mapping NaN/None/non-numeric to None."""
    if value is None:
        return None
    try:
        out = float(value)  # type: ignore[arg-type]
    except (TypeError, ValueError):
        return None
    return None if math.isnan(out) or math.isinf(out) else out


# --------------------------------------------------------------------------------------------------
# after-hours (best effort)
# --------------------------------------------------------------------------------------------------


def _is_equity_like(ticker: str) -> bool:
    """Plain equities/ETFs only (no indices ``^``, futures ``=F``, fx ``=X``, crypto ``-USD``, ``DX-Y.NYB``)."""
    return ticker.isascii() and ticker.isalpha()


def _fill_after_hours(quotes: list[Quote], session_date: date) -> None:
    """Populate after_hours_price / after_hours_change_pct in place; never raises."""
    deadline = time.monotonic() + AFTER_HOURS_BUDGET_SECONDS
    attempted = filled = 0
    for quote in quotes:
        if quote.close is None or not _is_equity_like(quote.symbol):
            continue
        if time.monotonic() >= deadline:
            log.warning("prices: after-hours budget exhausted after %d tickers", attempted)
            break
        attempted += 1
        try:
            price = _after_hours_price(quote.symbol, session_date)
        except Exception as exc:  # noqa: BLE001 - best effort only
            log.debug("prices: after-hours lookup failed for %s: %s", quote.symbol, exc)
            continue
        if price is None:
            continue
        quote.after_hours_price = price
        quote.after_hours_change_pct = (price / quote.close - 1.0) * 100.0 if quote.close else None
        filled += 1
    if attempted:
        log.info("prices: after-hours filled for %d/%d equity tickers", filled, attempted)


def _after_hours_price(ticker: str, session_date: date) -> float | None:
    """Last 1-minute Close stamped at/after 16:00 ET on ``session_date``; None when unavailable."""
    hist = yf.Ticker(ticker).history(
        start=session_date.isoformat(),
        end=(session_date + timedelta(days=1)).isoformat(),
        interval="1m",
        prepost=True,
        auto_adjust=False,
        actions=False,
        timeout=DOWNLOAD_TIMEOUT_SECONDS,
    )
    if hist is None or hist.empty or "Close" not in hist.columns:
        return None
    idx = pd.to_datetime(hist.index)
    idx = idx.tz_convert(ET) if idx.tz is not None else idx.tz_localize(ET)
    same_day = idx.date == session_date  # type: ignore[attr-defined]
    post = hist[same_day & (idx.time >= REGULAR_CLOSE_ET)]  # type: ignore[attr-defined]
    post = post[post["Close"].notna()]
    if post.empty:
        return None
    return _num(post["Close"].iloc[-1])