"""Offline downloader contract tests with mocked data providers."""

from __future__ import annotations

import numpy as np
import pandas as pd

import data.downloader as downloader
import pipeline.data.downloader as pipeline_downloader
from data.downloader import download_daily, download_daily_batch
from pipeline.common.schema import DATA_COLUMNS
from pipeline.data.downloader import download_universe


class _FakeResult:
    def __init__(self, rows, error_code="0", error_msg=""):
        self.rows = rows
        self.error_code = error_code
        self.error_msg = error_msg
        self._idx = -1

    def next(self):
        self._idx += 1
        return self._idx < len(self.rows)

    def get_row_data(self):
        return self.rows[self._idx]


def _daily_batch_row(
    date: str = "2024-01-02",
    open_: str = "10",
    high: str = "11",
    low: str = "9",
    close: str = "10.5",
    preclose: str = "10",
    volume: str = "1000",
    amount: str = "10500",
) -> list[str]:
    return [
        date,
        open_,
        high,
        low,
        close,
        preclose,
        volume,
        amount,
        "1.23",
        "5.0",
        "1",
        "0",
        "8.0",
        "1.1",
        "2.2",
        "3.3",
    ]


def test_download_daily_uses_baostock_before_akshare_in_auto(monkeypatch):
    calls: list[str] = []
    expected = pd.DataFrame({
        "symbol": ["sh600000"],
        "date": ["2024-01-02"],
        "open": [10.0],
        "high": [11.0],
        "low": [9.0],
        "close": [10.5],
        "volume": [1000.0],
        "amount": [10500.0],
    })

    def fake_baostock(symbol, start, end, adjust):
        calls.append("baostock")
        return expected

    def fake_akshare(symbol, start, end, adjust):
        calls.append("akshare")
        raise AssertionError("akshare should not be called after baostock succeeds")

    monkeypatch.setattr(downloader, "_download_baostock", fake_baostock)
    monkeypatch.setattr(downloader, "_download_akshare", fake_akshare)

    result = download_daily("sh600000", "2024-01-02", "2024-01-02", source="auto")

    assert calls == ["baostock"]
    assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]
    assert result["close"].tolist() == [10.5]


def test_download_daily_falls_back_to_akshare_when_baostock_empty(monkeypatch):
    calls: list[str] = []
    fallback = pd.DataFrame({
        "symbol": ["sz000001"],
        "date": ["2024-01-02"],
        "open": [20.0],
        "high": [21.0],
        "low": [19.0],
        "close": [20.5],
        "volume": [2000.0],
        "amount": [41000.0],
    })

    monkeypatch.setattr(
        downloader,
        "_download_baostock",
        lambda symbol, start, end, adjust: calls.append("baostock") or None,
    )
    monkeypatch.setattr(
        downloader,
        "_download_akshare",
        lambda symbol, start, end, adjust: calls.append("akshare") or fallback,
    )

    result = download_daily("sz000001", "2024-01-02", "2024-01-02", source="auto")

    assert calls == ["baostock", "akshare"]
    assert result["symbol"].tolist() == ["sz000001"]
    assert result["date"].tolist() == [pd.Timestamp("2024-01-02")]


def test_download_daily_batch_maps_rich_schema_and_vwap(monkeypatch):
    query_calls: list[dict] = []
    login_count = 0
    logout_count = 0

    def fake_login():
        nonlocal login_count
        login_count += 1

    def fake_logout():
        nonlocal logout_count
        logout_count += 1

    def fake_query(**kwargs):
        query_calls.append(kwargs)
        rows = [
            _daily_batch_row(volume="1000", amount="10500"),
            _daily_batch_row(date="2024-01-03", volume="0", amount="0"),
        ]
        return _FakeResult(rows)

    monkeypatch.setattr(downloader.bs, "login", fake_login)
    monkeypatch.setattr(downloader.bs, "logout", fake_logout)
    monkeypatch.setattr(downloader.bs, "query_history_k_data_plus", fake_query)

    [(symbol, frame)] = list(
        download_daily_batch(
            ["sh600000"],
            "2024-01-02",
            "2024-01-03",
            adjust="hfq",
        )
    )

    assert symbol == "sh600000"
    assert query_calls[0]["code"] == "sh.600000"
    assert query_calls[0]["adjustflag"] == "1"
    assert login_count == 1
    assert logout_count == 1
    assert frame is not None
    assert frame.columns.tolist() == [
        "symbol", "date", "open", "high", "low", "close", "preclose",
        "volume", "amount", "vwap", "turn", "pctChg", "tradestatus", "isST",
        "peTTM", "pbMRQ", "psTTM", "pcfNcfTTM",
    ]
    assert np.isclose(frame["vwap"].iloc[0], 10.5)
    assert pd.isna(frame["vwap"].iloc[1])
    assert pd.api.types.is_datetime64_any_dtype(frame["date"])
    assert pd.api.types.is_numeric_dtype(frame["tradestatus"])


def test_download_daily_batch_relogs_and_retries_session_loss(monkeypatch):
    responses = [
        _FakeResult([], error_code="10002007", error_msg="用户未登录"),
        _FakeResult([_daily_batch_row()]),
    ]
    login_count = 0
    logout_count = 0

    def fake_login():
        nonlocal login_count
        login_count += 1

    def fake_logout():
        nonlocal logout_count
        logout_count += 1

    monkeypatch.setattr(downloader.bs, "login", fake_login)
    monkeypatch.setattr(downloader.bs, "logout", fake_logout)
    monkeypatch.setattr(
        downloader.bs,
        "query_history_k_data_plus",
        lambda **kwargs: responses.pop(0),
    )

    [(symbol, frame)] = list(download_daily_batch(["sh600000"], "2024-01-02", "2024-01-02"))

    assert symbol == "sh600000"
    assert frame is not None
    assert len(frame) == 1
    assert login_count == 2
    assert logout_count == 2


def test_download_daily_batch_uses_akshare_fallback_when_enabled(monkeypatch):
    fallback = pd.DataFrame({
        "symbol": ["sh600000"],
        "date": ["2024-01-02"],
        "open": [10.0],
        "high": [11.0],
        "low": [9.0],
        "close": [10.5],
        "volume": [1000.0],
        "amount": [10500.0],
    })

    monkeypatch.setattr(downloader.bs, "login", lambda: None)
    monkeypatch.setattr(downloader.bs, "logout", lambda: None)
    monkeypatch.setattr(
        downloader.bs,
        "query_history_k_data_plus",
        lambda **kwargs: _FakeResult([], error_code="1", error_msg="no data"),
    )
    monkeypatch.setattr(
        downloader,
        "_download_akshare",
        lambda symbol, start, end, adjust: fallback.copy(),
    )

    [(symbol, frame)] = list(
        download_daily_batch(
            ["sh600000"],
            "2024-01-02",
            "2024-01-02",
            akshare_fallback=True,
        )
    )

    assert symbol == "sh600000"
    assert frame is not None
    assert frame["date"].tolist() == [pd.Timestamp("2024-01-02")]
    assert frame["close"].tolist() == [10.5]


def test_download_universe_writes_daily_partitions_from_mock_batch(tmp_path, monkeypatch):
    batch_frame = pd.DataFrame({
        "symbol": ["sh600000", "sh600000"],
        "date": pd.to_datetime(["2024-01-02", "2024-02-01"]),
        "open": [10.0, 11.0],
        "high": [11.0, 12.0],
        "low": [9.0, 10.0],
        "close": [10.5, 11.5],
        "preclose": [10.0, 10.5],
        "volume": [1000.0, 1200.0],
        "amount": [10500.0, 13800.0],
        "vwap": [10.5, 11.5],
        "turn": [1.0, 1.1],
        "pctChg": [5.0, 9.5],
        "tradestatus": [1, 1],
        "isST": [0, 0],
        "peTTM": [8.0, 8.1],
        "pbMRQ": [1.1, 1.2],
        "psTTM": [2.1, 2.2],
        "pcfNcfTTM": [3.1, 3.2],
    })

    monkeypatch.setattr(
        pipeline_downloader,
        "_resolve_universe",
        lambda universe, max_symbols=0: pd.DataFrame({
            "symbol_id": ["sh600000", "sz000001"],
            "symbol_name": ["PF Bank", "Ping An Bank"],
        }),
    )

    def fake_batch(symbols, start, end, adjust="qfq"):
        assert symbols == ["sh600000", "sz000001"]
        assert adjust == "qfq"
        yield "sh600000", batch_frame
        yield "sz000001", None

    monkeypatch.setattr(pipeline_downloader, "download_daily_batch", fake_batch)

    stats = download_universe(
        universe="toy",
        start_date="2024-01-02",
        end_date="2024-02-01",
        output_dir=str(tmp_path),
        chunk_size=1,
    )

    dataset_path = tmp_path / "toy"
    written = pd.read_parquet(dataset_path).sort_values(["date", "symbol_id"]).reset_index(drop=True)
    assert stats == {
        "dataset_path": str(dataset_path),
        "n_symbols": 1,
        "n_requested": 2,
        "n_rows": 2,
        "date_min": "2024-01-02",
        "date_max": "2024-02-01",
    }
    assert (dataset_path / "month=2024-01").exists()
    assert (dataset_path / "month=2024-02").exists()
    assert written[DATA_COLUMNS].columns.tolist() == DATA_COLUMNS
    assert written["symbol_name"].tolist() == ["PF Bank", "PF Bank"]