"""hub.py against a mocked transport: no network, no real cache dirs."""
from __future__ import annotations

import io
import json
import urllib.error

import pytest

from align.data import hub


class FakeResp(io.BytesIO):
    def __init__(self, data: bytes, status: int = 200, headers: dict | None = None, fail_after: int | None = None):
        super().__init__(data)
        self.status = status
        self.headers = headers or {}
        self._fail_after = fail_after
        self._served = 0

    def getcode(self):
        return self.status

    def read(self, n=-1):
        if self._fail_after is not None and self._served >= self._fail_after:
            raise ConnectionResetError("boom")
        out = super().read(n if self._fail_after is None else min(n, self._fail_after - self._served))
        self._served += len(out)
        return out


@pytest.fixture(autouse=True)
def _no_sleep(monkeypatch):
    monkeypatch.setattr(hub, "_sleep", lambda s: None)
    monkeypatch.setenv("HF_TOKEN", "tok123")


def test_url_builders():
    assert hub.resolve_url("a/b", "data/x y.jsonl", "abc") == (
        "https://huggingface.co/datasets/a/b/resolve/abc/data/x%20y.jsonl"
    )
    assert hub.tree_url("a/b", "main").endswith("/api/datasets/a/b/tree/main?recursive=true")
    u = hub.rows_url("a/b", "default", "train", 200)
    assert u.startswith("https://datasets-server.huggingface.co/rows?") and "offset=200" in u and "length=100" in u


def test_token_file_and_env(tmp_path, monkeypatch):
    monkeypatch.delenv("HF_TOKEN")
    t = tmp_path / "token"
    t.write_text("  filetok\n")
    assert hub.read_token(t) == "filetok"
    assert hub.read_token(tmp_path / "missing") is None
    monkeypatch.setenv("HF_TOKEN", "envtok")
    assert hub.read_token(t) == "envtok"


def test_token_only_sent_to_hub_hosts(monkeypatch):
    seen = []

    def fake(req, timeout):
        seen.append((req.full_url, req.headers.get("Authorization")))
        return FakeResp(b"x")

    monkeypatch.setattr(hub, "_urlopen", fake)
    hub._open_with_retry("https://huggingface.co/x")
    hub._open_with_retry("https://cdn.example.com/x")
    assert seen[0][1] == "Bearer tok123" and seen[1][1] is None


def test_redirect_strips_auth_across_hosts():
    import urllib.request

    h = hub._StripAuthOnRedirect()
    req = urllib.request.Request("https://huggingface.co/a", headers={"Authorization": "Bearer t"})
    new = h.redirect_request(req, None, 302, "Found", {}, "https://cdn.example.com/b")
    assert "Authorization" not in {k.title() for k in new.headers} | {k.title() for k in new.unredirected_hdrs}
    new2 = h.redirect_request(req, None, 302, "Found", {}, "https://huggingface.co/c")
    assert "Authorization" in {k.title() for k in new2.headers}


def test_stream_resume_uses_range_at_line_boundary(monkeypatch):
    payload = b"".join(json.dumps({"i": i}).encode() + b"\n" for i in range(20))
    calls = []

    def fake(req, timeout):
        rng = req.headers.get("Range")
        calls.append(rng)
        if rng is None:
            return FakeResp(payload, 200, fail_after=len(payload) // 2 + 3)  # dies mid-line
        start = int(rng.split("=")[1].rstrip("-"))
        return FakeResp(payload[start:], 206)

    monkeypatch.setattr(hub, "_urlopen", fake)
    got = [json.loads(x)["i"] for x in hub.iter_lines("https://huggingface.co/f")]
    assert got == list(range(20))  # no dup, no gap, no split record
    assert calls[0] is None and calls[1].startswith("bytes=") and len(calls) == 2
    start = int(calls[1].split("=")[1].rstrip("-"))
    assert payload[start - 1 : start] == b"\n"  # resumed on a line boundary


def test_resume_refuses_if_server_ignores_range(monkeypatch):
    payload = b"a\nb\nc\nd\n"

    def fake(req, timeout):
        if req.headers.get("Range") is None:
            return FakeResp(payload, 200, fail_after=5)
        return FakeResp(payload, 200)  # ignores Range

    monkeypatch.setattr(hub, "_urlopen", fake)
    with pytest.raises(hub.HubError, match="ignored Range"):
        list(hub.iter_lines("https://huggingface.co/f"))


def test_iter_jsonl_skips_blank_and_no_trailing_newline(monkeypatch):
    monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(b'{"a":1}\n\n{"a":2}'))
    assert list(hub.iter_jsonl("o/r", "f.jsonl")) == [{"a": 1}, {"a": 2}]


def test_retry_then_fail_fast_on_404(monkeypatch):
    n = {"c": 0}

    def flaky(req, timeout):
        n["c"] += 1
        if n["c"] < 3:
            raise urllib.error.HTTPError(req.full_url, 503, "busy", {}, None)
        return FakeResp(b"ok")

    monkeypatch.setattr(hub, "_urlopen", flaky)
    assert hub._open_with_retry("https://huggingface.co/x").read() == b"ok" and n["c"] == 3

    def gone(req, timeout):
        raise urllib.error.HTTPError(req.full_url, 404, "nope", {}, None)

    monkeypatch.setattr(hub, "_urlopen", gone)
    with pytest.raises(hub.HubError, match="404"):
        hub._open_with_retry("https://huggingface.co/x")


def test_list_repo_files_follows_link_and_caches_immutable(monkeypatch, tmp_path):
    sha = "a" * 40
    pages = {
        hub.tree_url("o/r", sha): FakeResp(
            json.dumps([{"path": "a.parquet", "type": "file"}]).encode(),
            headers={"Link": '<https://huggingface.co/page2>; rel="next"'},
        ),
        "https://huggingface.co/page2": FakeResp(
            json.dumps([{"path": "d", "type": "directory"}, {"path": "b.jsonl", "type": "file"}]).encode()
        ),
    }
    calls = []

    def fake(req, timeout):
        calls.append(req.full_url)
        return pages[req.full_url]

    monkeypatch.setattr(hub, "_urlopen", fake)
    assert hub.list_repo_file_paths("o/r", sha, cache_dir=tmp_path) == ["a.parquet", "b.jsonl"]
    n = len(calls)
    hub.list_repo_files("o/r", sha, cache_dir=tmp_path)
    assert len(calls) == n  # served from cache
    # moving ref is never cached
    monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(b"[]"))
    assert hub.list_repo_files("o/r", "main", cache_dir=tmp_path) == []


def _page(offset, n, total, truncated=False):
    rows = [
        {"row_idx": offset + i, "row": {"text": f"t{offset + i}"}, "truncated_cells": ["text"] if truncated else []}
        for i in range(n)
    ]
    return json.dumps({"rows": rows, "num_rows_total": total}).encode()


def test_iter_rows_paginates_caches_and_respects_max(monkeypatch, tmp_path):
    calls = []

    def fake(req, timeout):
        calls.append(req.full_url)
        off = int(req.full_url.split("offset=")[1].split("&")[0])
        return FakeResp(_page(off, min(100, 250 - off), 250))

    monkeypatch.setattr(hub, "_urlopen", fake)
    rows = list(hub.iter_rows("o/r", "default", "train", cache_dir=tmp_path))
    assert [r["text"] for r in rows] == [f"t{i}" for i in range(250)]
    assert len(calls) == 3
    list(hub.iter_rows("o/r", "default", "train", cache_dir=tmp_path))
    assert len(calls) == 3  # second pass fully cached
    assert len(list(hub.iter_rows("o/r", max_rows=150, cache_dir=tmp_path / "c2"))) == 150


def test_iter_rows_raises_on_truncated_cells(monkeypatch, tmp_path):
    monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(_page(0, 2, 2, truncated=True)))
    with pytest.raises(hub.HubError, match="truncated"):
        list(hub.iter_rows("o/r", cache_dir=tmp_path))


def test_rows_length_bound():
    with pytest.raises(ValueError):
        hub.fetch_rows_page("o/r", "default", "train", 0, length=101)


def test_corrupt_cache_page_is_refetched(monkeypatch, tmp_path):
    url = hub.rows_url("o/r", "default", "train", 0, 100)
    cf = hub._cache_file(tmp_path, "rows", url)
    cf.parent.mkdir(parents=True)
    cf.write_text("{not json")
    monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(_page(0, 1, 1)))
    assert list(hub.iter_rows("o/r", cache_dir=tmp_path)) == [{"text": "t0"}]


def test_pretrain_stream_goes_through_hub(monkeypatch, tmp_path):
    from align.data import pretrain_corpus as pc

    monkeypatch.setattr(hub, "DEFAULT_CACHE_DIR", str(tmp_path))
    monkeypatch.setattr(hub, "_urlopen", lambda r, t: FakeResp(_page(0, 3, 3)))
    assert list(pc.iter_streamed_texts("o/r")) == ["t0", "t1", "t2"]
    with pytest.raises(KeyError, match="nope"):
        list(pc.iter_streamed_texts("o/r", text_key="nope"))