"""Tests for us2kr.collectors.prices (no network: yfinance.download / yfinance.Ticker are monkeypatched)."""

from __future__ import annotations

from datetime import date, datetime, timedelta
from typing import ClassVar
from zoneinfo import ZoneInfo

import numpy as np
import pandas as pd
import pytest
import yfinance as yf

from us2kr.collectors import prices
from us2kr.collectors.prices import PriceFetchError, fetch_quotes

SESSION = date(2026, 8, 24)  # Monday
ET = ZoneInfo("America/New_York")
FIELDS = ["Open", "High", "Low", "Close", "Adj Close", "Volume"]


# ------------------------------------------------------------------------------------------------
# fixtures / helpers
# ------------------------------------------------------------------------------------------------


def _dates(n_before: int, n_after: int = 0) -> pd.DatetimeIndex:
    """Calendar-day index (weekends included, like the union index yf.download produces)."""
    start = pd.Timestamp(SESSION) - pd.Timedelta(days=n_before)
    return pd.date_range(start, pd.Timestamp(SESSION) + pd.Timedelta(days=n_after), freq="D")


def _ticker_block(idx: pd.DatetimeIndex, closes: dict[date, float], volume: float | None = 1_000.0):
    """Build a 6-column block for one ticker; rows not in ``closes`` are NaN (weekend / no trade)."""
    block = pd.DataFrame(index=idx, columns=FIELDS, dtype="float64")
    for d, c in closes.items():
        ts = pd.Timestamp(d)
        block.loc[ts, "Open"] = c - 1.0
        block.loc[ts, "High"] = c + 2.0
        block.loc[ts, "Low"] = c - 3.0
        block.loc[ts, "Close"] = c
        block.loc[ts, "Adj Close"] = c
        block.loc[ts, "Volume"] = np.nan if volume is None else volume
    return block


def _weekday_closes(last_close: float, n: int, step: float = 1.0, end: date = SESSION) -> dict[date, float]:
    """n weekday closes ending on ``end`` (inclusive), increasing by ``step`` per day: ... , last-step, last."""
    out: dict[date, float] = {}
    d = end
    k = 0
    while len(out) < n:
        if d.weekday() < 5:
            out[d] = last_close - step * k
            k += 1
        d -= timedelta(days=1)
    return out


def _multi(blocks: dict[str, pd.DataFrame]) -> pd.DataFrame:
    """Assemble the (Ticker, Price) MultiIndex frame exactly like yf.download(group_by='ticker')."""
    df = pd.concat(blocks.values(), axis=1, keys=blocks.keys(), names=["Ticker", "Price"])
    df.index.name = "Date"
    return df


class FakeTicker:
    """Stand-in for yfinance.Ticker; ``history`` returns pre-registered frames or raises."""

    frames: ClassVar[dict[str, pd.DataFrame]] = {}
    errors: ClassVar[set[str]] = set()
    calls: ClassVar[list[tuple[str, dict]]] = []

    def __init__(self, ticker: str, session=None):
        self.ticker = ticker

    def history(self, **kwargs):
        FakeTicker.calls.append((self.ticker, kwargs))
        if self.ticker in FakeTicker.errors:
            raise RuntimeError(f"boom {self.ticker}")
        return FakeTicker.frames.get(self.ticker, pd.DataFrame())


@pytest.fixture(autouse=True)
def _reset_fake_ticker(monkeypatch):
    FakeTicker.frames = {}
    FakeTicker.errors = set()
    FakeTicker.calls = []
    monkeypatch.setattr(yf, "Ticker", FakeTicker)
    yield


@pytest.fixture
def batch_frame() -> pd.DataFrame:
    idx = _dates(n_before=20, n_after=1)  # includes the day AFTER the session (24h instruments)
    gspc = _weekday_closes(6500.0, 12, step=10.0)  # 6500 on SESSION, 6490 on Friday, ...
    tnx = _weekday_closes(4.25, 12, step=-0.05)  # 4.25 on SESSION, 4.30 on Friday (yield fell 5bp)
    aapl = _weekday_closes(200.0, 12, step=2.0)
    krw = _weekday_closes(1380.0, 12, step=5.0)
    krw[SESSION + timedelta(days=1)] = 1399.0  # bar AFTER the session must be ignored
    blocks = {
        "^GSPC": _ticker_block(idx, gspc, volume=4.0e9),
        "^TNX": _ticker_block(idx, tnx, volume=0.0),
        "AAPL": _ticker_block(idx, aapl, volume=50e6),
        "KRW=X": _ticker_block(idx, krw, volume=0.0),
        "MISSING": _ticker_block(idx, {}),  # all-NaN column, as yfinance emits for a bad symbol
    }
    # make AAPL's session volume differ from the prior average (prior 10 bars = 50e6, session = 100e6)
    blocks["AAPL"].loc[pd.Timestamp(SESSION), "Volume"] = 100e6
    return _multi(blocks)


@pytest.fixture
def patch_download(monkeypatch):
    def _install(frame_or_exc, *sequence):
        calls: list[dict] = []
        items = [frame_or_exc, *sequence]

        def fake_download(tickers, **kwargs):
            calls.append({"tickers": list(tickers), **kwargs})
            item = items[min(len(calls) - 1, len(items) - 1)]
            if isinstance(item, Exception):
                raise item
            return item

        monkeypatch.setattr(yf, "download", fake_download)
        return calls

    return _install


def _by_symbol(quotes):
    return {q.symbol: q for q in quotes}


# ------------------------------------------------------------------------------------------------
# tests: per-group computations
# ------------------------------------------------------------------------------------------------


def test_batch_download_called_with_date_window(batch_frame, patch_download):
    calls = patch_download(batch_frame)
    fetch_quotes(SESSION, ["^GSPC", "AAPL"])
    assert len(calls) == 1
    call = calls[0]
    assert call["tickers"] == ["^GSPC", "AAPL"]
    assert call["group_by"] == "ticker"
    assert call["auto_adjust"] is False
    assert call["interval"] == "1d"
    assert call["progress"] is False
    assert call["start"] == (SESSION - timedelta(days=prices.HISTORY_DAYS)).isoformat()
    assert call["end"] == (SESSION + timedelta(days=1)).isoformat()  # end is exclusive


def test_index_quote_fields(batch_frame, patch_download):
    patch_download(batch_frame)
    q = _by_symbol(fetch_quotes(SESSION, ["^GSPC"]))["^GSPC"]
    assert q.group == "indices"
    assert q.name == "S&P 500"
    assert q.close == pytest.approx(6500.0)
    assert q.prev_close == pytest.approx(6490.0)
    assert q.change_abs == pytest.approx(10.0)
    assert q.change_pct == pytest.approx((6500.0 / 6490.0 - 1) * 100)
    assert q.day_high == pytest.approx(6502.0)
    assert q.day_low == pytest.approx(6497.0)
    assert q.volume == pytest.approx(4.0e9)
    assert q.avg_volume_10d == pytest.approx(4.0e9)
    assert q.as_of == datetime(2026, 8, 24, tzinfo=ET)
    assert q.source == "yfinance"
    assert q.after_hours_price is None  # indices never get after-hours lookups


def test_rates_change_abs_in_yield_points(batch_frame, patch_download):
    patch_download(batch_frame)
    q = _by_symbol(fetch_quotes(SESSION, ["^TNX"]))["^TNX"]
    assert q.group == "rates"
    assert q.close == pytest.approx(4.25)
    assert q.prev_close == pytest.approx(4.30)
    assert q.change_abs == pytest.approx(-0.05)
    assert q.change_pct == pytest.approx((4.25 / 4.30 - 1) * 100)
    assert q.volume is None  # zero volume is reported as None
    assert q.avg_volume_10d is None


def test_fx_close_selected_by_date_not_last_row(batch_frame, patch_download):
    """KRW=X has a bar dated after the session; close must be the session-date bar, prev the one before."""
    patch_download(batch_frame)
    q = _by_symbol(fetch_quotes(SESSION, ["KRW=X"]))["KRW=X"]
    assert q.group == "fx"
    assert q.close == pytest.approx(1380.0)
    assert q.prev_close == pytest.approx(1375.0)
    assert q.change_abs == pytest.approx(5.0)


def test_equity_volume_ratio_inputs(batch_frame, patch_download):
    patch_download(batch_frame)
    q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
    assert q.group == "megacap"
    assert q.name == "애플"
    assert q.volume == pytest.approx(100e6)
    assert q.avg_volume_10d == pytest.approx(50e6)  # prior 10 sessions only, session day excluded


def test_missing_ticker_skipped_and_order_preserved(batch_frame, patch_download, caplog):
    patch_download(batch_frame)
    with caplog.at_level("WARNING", logger="us2kr.collectors.prices"):
        quotes = fetch_quotes(SESSION, ["^GSPC", "MISSING", "AAPL", "NOTINFRAME"])
    assert [q.symbol for q in quotes] == ["^GSPC", "AAPL"]
    assert "MISSING" in caplog.text
    assert "NOTINFRAME" in caplog.text


def test_unknown_ticker_gets_other_group(patch_download):
    idx = _dates(10)
    frame = _multi({"ZZZ": _ticker_block(idx, _weekday_closes(10.0, 5))})
    patch_download(frame)
    q = fetch_quotes(SESSION, ["ZZZ"])[0]
    assert q.group == "other"
    assert q.name == "ZZZ"


def test_default_tickers_is_symbol_universe(batch_frame, patch_download):
    calls = patch_download(batch_frame)
    fetch_quotes(SESSION)
    from us2kr import symbols

    assert calls[0]["tickers"] == symbols.all_tickers()


def test_empty_ticker_list_returns_empty_without_download(patch_download):
    calls = patch_download(pd.DataFrame())
    assert fetch_quotes(SESSION, []) == []
    assert calls == []


# ------------------------------------------------------------------------------------------------
# tests: session_date edge cases
# ------------------------------------------------------------------------------------------------


def test_session_row_missing_uses_last_prior_bar(patch_download, caplog):
    """If the session-date bar is absent (not yet published), fall back to the last bar before it."""
    idx = _dates(10)
    friday = SESSION - timedelta(days=3)
    closes = _weekday_closes(100.0, 5, step=1.0, end=friday)  # 100 on Fri, 99 on Thu
    frame = _multi({"XLK": _ticker_block(idx, closes)})
    patch_download(frame)
    with caplog.at_level("INFO", logger="us2kr.collectors.prices"):
        q = fetch_quotes(SESSION, ["XLK"])[0]
    assert q.close == pytest.approx(100.0)
    assert q.prev_close == pytest.approx(99.0)
    assert q.as_of.date() == friday
    assert "no 2026-08-24 bar" in caplog.text


def test_only_one_bar_gives_no_change(patch_download):
    idx = _dates(3)
    frame = _multi({"XLF": _ticker_block(idx, {SESSION: 50.0})})
    patch_download(frame)
    q = fetch_quotes(SESSION, ["XLF"])[0]
    assert q.close == pytest.approx(50.0)
    assert q.prev_close is None
    assert q.change_pct is None
    assert q.change_abs is None
    assert q.avg_volume_10d is None


def test_nan_close_rows_are_ignored_for_prev_close(patch_download):
    idx = _dates(10)
    closes = _weekday_closes(100.0, 5, step=1.0)
    block = _ticker_block(idx, closes)
    friday = SESSION - timedelta(days=3)
    block.loc[pd.Timestamp(friday), "Close"] = np.nan  # broken row must be skipped, not used as prev
    frame = _multi({"XLE": block})
    patch_download(frame)
    q = fetch_quotes(SESSION, ["XLE"])[0]
    assert q.close == pytest.approx(100.0)
    assert q.prev_close == pytest.approx(98.0)  # Thursday


def test_nan_volume_becomes_none(patch_download):
    idx = _dates(10)
    frame = _multi({"XLV": _ticker_block(idx, _weekday_closes(100.0, 5), volume=None)})
    patch_download(frame)
    q = fetch_quotes(SESSION, ["XLV"])[0]
    assert q.volume is None
    assert q.avg_volume_10d is None


def test_all_bars_after_session_date_skips_ticker(patch_download):
    idx = _dates(0, n_after=2)
    frame = _multi({"BTC-USD": _ticker_block(idx, {SESSION + timedelta(days=1): 1.0})})
    patch_download(frame)
    assert fetch_quotes(SESSION, ["BTC-USD"]) == []


# ------------------------------------------------------------------------------------------------
# tests: frame layouts
# ------------------------------------------------------------------------------------------------


def test_single_level_frame_for_single_ticker(patch_download):
    idx = _dates(10)
    frame = _ticker_block(idx, _weekday_closes(30.0, 5, step=0.5))
    patch_download(frame)
    q = fetch_quotes(SESSION, ["^SOX"])[0]
    assert q.close == pytest.approx(30.0)
    assert q.prev_close == pytest.approx(29.5)


def test_group_by_column_layout_is_handled(batch_frame, patch_download):
    swapped = batch_frame.swaplevel(0, 1, axis=1).sort_index(axis=1)
    patch_download(swapped)
    q = _by_symbol(fetch_quotes(SESSION, ["^GSPC"]))["^GSPC"]
    assert q.close == pytest.approx(6500.0)


def test_tz_aware_index_is_normalised(patch_download):
    idx = pd.date_range(pd.Timestamp(SESSION) - pd.Timedelta(days=6), pd.Timestamp(SESSION), freq="D")
    block = _ticker_block(idx, _weekday_closes(10.0, 4))
    block.index = block.index.tz_localize(ET)
    patch_download(_multi({"SMH": block}))
    q = fetch_quotes(SESSION, ["SMH"])[0]
    assert q.close == pytest.approx(10.0)
    assert q.prev_close == pytest.approx(9.0)


@pytest.mark.parametrize("tz", ["Europe/London", "UTC", "America/New_York", "America/Chicago"])
def test_tz_aware_daily_index_keeps_exchange_local_date(tz):
    """London/UTC-stamped daily bars (KRW=X, BTC-USD via Ticker.history) must not shift a day earlier."""
    idx = pd.DatetimeIndex(pd.to_datetime(["2026-08-20", "2026-08-21", "2026-08-24", "2026-08-25"]))
    frame = pd.DataFrame(
        {"Close": [1380.0, 1385.0, 1390.0, 1400.0], "Volume": [0.0] * 4}, index=idx.tz_localize(tz)
    )
    cleaned = prices._clean_single(frame)
    assert cleaned is not None and cleaned.index.tz is None
    dates = [d.date().isoformat() for d in cleaned.index]
    assert dates == ["2026-08-20", "2026-08-21", "2026-08-24", "2026-08-25"]
    q = prices._quote_from_history("KRW=X", cleaned, SESSION)
    assert q.close == pytest.approx(1390.0)
    assert q.prev_close == pytest.approx(1385.0)
    assert q.as_of == datetime(2026, 8, 24, tzinfo=ET)


# ------------------------------------------------------------------------------------------------
# tests: failure policy
# ------------------------------------------------------------------------------------------------


def test_batch_retried_once_then_succeeds(batch_frame, patch_download):
    calls = patch_download(RuntimeError("yahoo down"), batch_frame)
    quotes = fetch_quotes(SESSION, ["^GSPC"])
    assert len(calls) == 2
    assert quotes[0].close == pytest.approx(6500.0)


def test_batch_failure_falls_back_to_indices_only(patch_download):
    patch_download(RuntimeError("down"), RuntimeError("still down"))
    idx = _dates(10)
    FakeTicker.frames["^GSPC"] = _ticker_block(idx, _weekday_closes(6000.0, 5, step=10.0))
    FakeTicker.frames["AAPL"] = _ticker_block(idx, _weekday_closes(200.0, 5))
    quotes = fetch_quotes(SESSION, ["^GSPC", "AAPL", "^VIX"])
    # only indices are attempted in fallback; AAPL is not requested, ^VIX yields nothing
    assert [q.symbol for q in quotes] == ["^GSPC"]
    assert quotes[0].close == pytest.approx(6000.0)
    history_calls = [c for c in FakeTicker.calls if c[1].get("interval") == "1d"]
    assert sorted(t for t, _ in history_calls) == ["^GSPC", "^VIX"]


def test_empty_batch_frames_count_as_failure(patch_download):
    patch_download(pd.DataFrame(), None)
    FakeTicker.frames["^DJI"] = _ticker_block(_dates(10), _weekday_closes(40000.0, 5))
    quotes = fetch_quotes(SESSION, ["^DJI"])
    assert quotes[0].close == pytest.approx(40000.0)


def test_total_failure_raises_price_fetch_error(patch_download):
    patch_download(RuntimeError("down"))
    FakeTicker.errors = {"^GSPC"}
    with pytest.raises(PriceFetchError) as exc_info:
        fetch_quotes(SESSION, ["^GSPC", "AAPL"])
    assert "down" in str(exc_info.value)
    assert "^GSPC" in str(exc_info.value)


def test_run_safely_records_gap_on_total_failure(patch_download):
    from us2kr.collectors.base import run_safely

    patch_download(RuntimeError("down"))
    gaps: list[str] = []
    out = run_safely("prices", lambda: fetch_quotes(SESSION, ["AAPL"]), [], gaps)
    assert out == []
    assert gaps and gaps[0].startswith("prices: PriceFetchError")


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


def _minute_frame(closes_by_time: dict[str, float], day: date = SESSION) -> pd.DataFrame:
    stamps = [pd.Timestamp(f"{day.isoformat()} {t}", tz=ET) for t in closes_by_time]
    return pd.DataFrame(
        {"Open": list(closes_by_time.values()), "Close": list(closes_by_time.values()), "Volume": 0},
        index=pd.DatetimeIndex(stamps, name="Datetime"),
    )


def test_after_hours_price_from_post_close_minute_bars(batch_frame, patch_download):
    patch_download(batch_frame)
    FakeTicker.frames["AAPL"] = _minute_frame(
        {"15:58": 199.0, "15:59": 200.0, "16:00": 201.0, "19:59": 204.0}
    )
    q = _by_symbol(fetch_quotes(SESSION, ["AAPL", "^GSPC", "KRW=X"]))
    assert q["AAPL"].after_hours_price == pytest.approx(204.0)
    assert q["AAPL"].after_hours_change_pct == pytest.approx(2.0)
    assert q["^GSPC"].after_hours_price is None
    assert q["KRW=X"].after_hours_price is None
    # only the equity was queried, with prepost 1m bars for the session day
    assert [t for t, _ in FakeTicker.calls] == ["AAPL"]
    kwargs = FakeTicker.calls[0][1]
    assert kwargs["interval"] == "1m" and kwargs["prepost"] is True
    assert kwargs["start"] == SESSION.isoformat()


def test_after_hours_skipped_when_only_regular_bars(batch_frame, patch_download):
    patch_download(batch_frame)
    FakeTicker.frames["AAPL"] = _minute_frame({"15:58": 199.0, "15:59": 200.0})
    q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
    assert q.after_hours_price is None
    assert q.after_hours_change_pct is None


def test_after_hours_ignores_bars_from_other_days(batch_frame, patch_download):
    patch_download(batch_frame)
    FakeTicker.frames["AAPL"] = _minute_frame({"17:00": 150.0}, day=SESSION - timedelta(days=3))
    q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
    assert q.after_hours_price is None


def test_after_hours_failure_is_isolated(batch_frame, patch_download):
    patch_download(batch_frame)
    FakeTicker.errors = {"AAPL"}
    q = _by_symbol(fetch_quotes(SESSION, ["AAPL", "^GSPC"]))
    assert q["AAPL"].close == pytest.approx(200.0)
    assert q["AAPL"].after_hours_price is None
    assert q["^GSPC"].close == pytest.approx(6500.0)


def test_after_hours_respects_time_budget(batch_frame, patch_download, monkeypatch):
    patch_download(batch_frame)
    FakeTicker.frames["AAPL"] = _minute_frame({"16:30": 210.0})
    monkeypatch.setattr(prices, "AFTER_HOURS_BUDGET_SECONDS", 0.0)
    q = _by_symbol(fetch_quotes(SESSION, ["AAPL"]))["AAPL"]
    assert q.after_hours_price is None
    assert FakeTicker.calls == []


def test_is_equity_like_classification():
    assert prices._is_equity_like("AAPL")
    assert prices._is_equity_like("EWY")
    for t in ("^GSPC", "ES=F", "KRW=X", "DX-Y.NYB", "BTC-USD", "^TNX"):
        assert not prices._is_equity_like(t), t