"""Offline tests for baostock-backed universe helpers."""

from __future__ import annotations

import pandas as pd

import data.universe as universe


class _FakeResult:
    def __init__(self, rows, fields=None):
        self.rows = rows
        self.fields = fields or ["code", "name", "date"]
        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 test_index_constituent_helpers_normalize_dotted_codes(monkeypatch):
    calls: list[str] = []

    monkeypatch.setattr(universe.bs, "login", lambda: calls.append("login"))
    monkeypatch.setattr(universe.bs, "logout", lambda: calls.append("logout"))
    monkeypatch.setattr(
        universe.bs,
        "query_hs300_stocks",
        lambda: _FakeResult([
            ["sh.600000", "浦发银行", "2024-01-12"],
            ["sz.000001", "平安银行", "2024-01-12"],
        ]),
    )
    monkeypatch.setattr(
        universe.bs,
        "query_zz500_stocks",
        lambda: _FakeResult([
            ["sh.600006", "东风汽车", "2024-01-12"],
        ]),
    )

    hs300 = universe.get_hs300_stocks()
    zz500 = universe.get_zz500_stocks()

    assert calls == ["login", "logout", "login", "logout"]
    assert hs300["code"].tolist() == ["sh600000", "sz000001"]
    assert zz500["code"].tolist() == ["sh600006"]
    assert hs300["name"].tolist() == ["浦发银行", "平安银行"]


def test_get_all_stocks_walks_back_and_filters_to_listed_a_shares(monkeypatch):
    fields = ["code", "tradeStatus", "code_name"]
    responses = [
        _FakeResult([], fields=fields),
        _FakeResult(
            [
                ["sh.600000", "1", "浦发银行"],
                ["sh.688001", "1", "华兴源创"],
                ["sz.000001", "1", "平安银行"],
                ["sz.300750", "1", "宁德时代"],
                ["sz.399001", "1", "深证成指"],
                ["sz.200001", "1", "深物业B"],
                ["bj.430047", "1", "北交所样本"],
            ],
            fields=fields,
        ),
    ]
    query_days: list[str] = []

    monkeypatch.setattr(universe.bs, "login", lambda: None)
    monkeypatch.setattr(universe.bs, "logout", lambda: None)

    def fake_query_all_stock(day):
        query_days.append(day)
        return responses.pop(0)

    monkeypatch.setattr(universe.bs, "query_all_stock", fake_query_all_stock)

    result = universe.get_all_stocks("2024-01-07")

    assert query_days == ["2024-01-07", "2024-01-06"]
    assert result.columns.tolist() == ["code", "name"]
    assert result["code"].tolist() == [
        "sh600000",
        "sh688001",
        "sz000001",
        "sz300750",
    ]
    assert result["name"].tolist() == ["浦发银行", "华兴源创", "平安银行", "宁德时代"]