"""Generating the corpus straight into the transfer.

The properties under test are the ones the old two-stage build got for free and
this one has to earn: that a shard is never staged locally, that nothing larger
than one record is ever held, that an interrupted run resumes without forgetting
what it learned, and that the vetting contract did not move when the pipeline
around it did.
"""

import json
import os

import numpy as np
import pytest

import dataset_build
from conftest import solar_disc, write_fits
from suvi import corpus, dataset, paths

CAD = paths.CADENCE
BASE = 1715385600          # 2024-05-11 00:00 UTC
DAY = "2024-05-11"

#: Scale putting a peak-1.0 synthetic disc's IMG_MEAN inside each band's
#: physical bounds (header_v1 rightly rejects a 304A frame as dim as a 94A one).
BAND_SCALE = {94: 0.5, 131: 0.5, 171: 1.0, 195: 2.0, 284: 2.0, 304: 3.0}


# ------------------------------------------------------------------ the archive


class FakeRow(dict):
    """sqlite3.Row-alike: the builder reads columns by name."""


class FakeConn:
    """Stands in for the index, counting reads so 'was this re-read' is testable."""

    def __init__(self, rows):
        self.rows = rows            # [(satellite, wavelength, t_start, relpath)]
        self.queries = 0

    def execute(self, sql, params):
        self.queries += 1
        satellite, times = params[0], set(params[1:])
        return [FakeRow(wavelength=w, t_start=t, path=p)
                for (s, w, t, p) in self.rows if s == satellite and t in times]


def make_archive(root, times, satellites=(16, 18), broken=()):
    """Write a synthetic archive and return the rows an index would hold."""
    rows = []
    for satellite in satellites:
        for t in times:
            for wavelength in paths.WAVELENGTHS:
                name = f"g{satellite}_{wavelength}_{t}.fits"
                image = (np.zeros((1280, 1280), dtype=np.float32)
                         if (satellite, t) in broken
                         else solar_disc(size=1280, radius=386, peak=1.0))
                write_fits(os.path.join(root, name),
                           image * BAND_SCALE[wavelength],
                           headers={"WAVELNTH": wavelength})
                rows.append((satellite, wavelength, t, name))
    return rows


# ------------------------------------------------------------------ destinations


def test_a_local_destination_writes_only_the_shard(tmp_path):
    """Nothing is staged: the destination is the only thing that appears."""
    root = tmp_path / "archive"
    out = tmp_path / "dest"
    root.mkdir()
    times = [BASE + i * CAD for i in range(3)]
    rows = make_archive(str(root), times, satellites=(16,))
    conn = FakeConn(rows)
    destination = dataset_build.Destination(str(out))

    size, kept, dropped, unusable = dataset_build.build_shard(
        conn, str(root), DAY, 16, times, destination, workers=1)

    assert kept == 3 and not dropped and not unusable and size > 0
    assert os.listdir(out) == [dataset.shard_name(DAY, 16)]
    with dataset.Shard(str(out / dataset.shard_name(DAY, 16))) as shard:
        assert shard.times() == times
        assert all(shard.usable(t) for t in times)


def test_a_failed_stream_leaves_nothing_at_the_real_name(tmp_path):
    """A shard is committed by its rename.  A run that dies mid-stream must leave
    no file the resume scan would mistake for finished work."""
    out = tmp_path / "dest"
    destination = dataset_build.Destination(str(out))
    with pytest.raises(RuntimeError):
        with destination.open("x.shard") as handle:
            handle.write(b"partial")
            raise RuntimeError("stream died")
    assert not os.path.exists(out / "x.shard")
    assert not os.path.exists(out / "x.shard.partial")


def test_existing_names_are_what_resume_is_keyed_on(tmp_path):
    out = tmp_path / "dest"
    out.mkdir()
    (out / "a.shard").write_bytes(b"")
    destination = dataset_build.Destination(str(out))
    assert destination.existing() == {"a.shard"}
    assert dataset_build.Destination(str(tmp_path / "absent")).existing() == set()


def test_put_and_get_round_trip_small_files(tmp_path):
    destination = dataset_build.Destination(str(tmp_path / "dest"))
    assert destination.get("vetting.json") is None
    destination.put("vetting.json", b'[["2024-05-11", 16, 1000]]')
    assert json.loads(destination.get("vetting.json")) == [["2024-05-11", 16, 1000]]


def test_a_remote_destination_streams_through_ssh_without_staging():
    """The remote path builds a command that pipes into cat and commits by rename;
    no local file appears in it at any point."""
    destination = dataset_build.Destination("host:/var/shards")
    assert destination.host == "host" and destination.root == "/var/shards"
    writer = dataset_build._RemoteWrite(destination, "a.shard")
    assert "cat > " in writer.command
    assert "a.shard.partial" in writer.command
    assert "mv " in writer.command


# ------------------------------------------------------------------- generation


def test_the_shard_is_never_held_whole(tmp_path):
    """The memory property.  Records are pulled one at a time from the pool, so a
    600-slot shard is written with one ~4 MB payload live, not 1 GB."""
    root = tmp_path / "archive"
    root.mkdir()
    times = [BASE + i * CAD for i in range(4)]
    rows = make_archive(str(root), times, satellites=(16,))
    conn = FakeConn(rows)

    live = []
    real_write = dataset.write_shard_stream

    def watching(handle, day, satellite, wavelengths, records, level=3):
        def counted():
            for item in records:
                live.append(1)
                yield item
                live.pop()
        return real_write(handle, day, satellite, wavelengths, counted(), level)

    dataset.write_shard_stream = watching
    try:
        dataset_build.build_shard(conn, str(root), DAY, 16, times,
                                  dataset_build.Destination(str(tmp_path / "d")),
                                  workers=1)
    finally:
        dataset.write_shard_stream = real_write
    assert live == []          # every record released before the next was taken


def test_a_vetoed_slot_is_stored_as_a_sentinel_and_reported(tmp_path):
    """An all-zero frame parses fine and would enter the shard as 'truth' unvetted.
    It must land as a sentinel *and* be named, because the segment it sits in has
    to be split afterwards."""
    root = tmp_path / "archive"
    out = tmp_path / "dest"
    root.mkdir()
    times = [BASE + i * CAD for i in range(3)]
    rows = make_archive(str(root), times, satellites=(16,),
                        broken={(16, times[1])})
    conn = FakeConn(rows)

    size, kept, dropped, unusable = dataset_build.build_shard(
        conn, str(root), DAY, 16, times,
        dataset_build.Destination(str(out)), workers=1)

    assert kept == 2
    assert unusable == {times[1]}
    assert sum(dropped.values()) == 1
    with dataset.Shard(str(out / dataset.shard_name(DAY, 16))) as shard:
        assert times[1] in shard              # we looked
        assert not shard.usable(times[1])     # and there was nothing usable
        assert shard.raw(times[1]) is None


def test_a_slot_missing_from_the_index_becomes_a_sentinel(tmp_path):
    """A wanted slot the archive does not have must still occupy its place, or the
    shard's slots stop lining up with the segment that asked for them."""
    root = tmp_path / "archive"
    out = tmp_path / "dest"
    root.mkdir()
    times = [BASE + i * CAD for i in range(3)]
    rows = [r for r in make_archive(str(root), times, satellites=(16,))
            if r[2] != times[2]]
    conn = FakeConn(rows)

    _, kept, _, unusable = dataset_build.build_shard(
        conn, str(root), DAY, 16, times,
        dataset_build.Destination(str(out)), workers=1)
    assert kept == 2 and unusable == {times[2]}
    with dataset.Shard(str(out / dataset.shard_name(DAY, 16))) as shard:
        assert shard.times() == times


def test_generation_defaults_to_a_single_reader():
    """One outstanding read against a shared mount is a promise to the other
    consumers of it, not an implementation detail free to drift.

    Read concurrency is the one lever that would make this faster -- 48 readers
    reach ~207 MB/s against one reader's ~62 -- so the temptation to raise it is
    permanent and the default is worth pinning.
    """
    args = dataset_build.build_parser().parse_args(
        ["--availability", "a.npz", "--dest", "/tmp/x", "--budget-gb", "1"])
    assert args.workers == 1


# ------------------------------------------------------------------ verification


def test_a_vetoed_slot_splits_the_segment_that_contained_it():
    """The profile is header-only; disc_v1 rejects inside a nominally pristine run.
    Verification runs off what extraction observed, so nothing is read back."""
    segment = (BASE, BASE + 10 * CAD)
    vetoed = {(DAY, 16, BASE + 4 * CAD)}
    verified = dataset_build.verified_segments(
        {"train": [segment]}, vetoed, satellites=(16, 18), minimum_slots=3)
    assert [pair for _, pair in verified] == [
        (BASE, BASE + 4 * CAD), (BASE + 5 * CAD, BASE + 10 * CAD)]


def test_fragments_below_the_span_minimum_are_dropped():
    """A stub too short to hold one model input is not a segment."""
    segment = (BASE, BASE + 10 * CAD)
    vetoed = {(DAY, 16, BASE + 2 * CAD)}
    verified = dataset_build.verified_segments(
        {"train": [segment]}, vetoed, satellites=(16,), minimum_slots=5)
    assert [pair for _, pair in verified] == [(BASE + 3 * CAD, BASE + 10 * CAD)]


def test_a_veto_on_either_satellite_splits_the_segment():
    """Segments are jointly pristine: a slot is usable only if every satellite is."""
    segment = (BASE, BASE + 6 * CAD)
    verified = dataset_build.verified_segments(
        {"train": [segment]}, {(DAY, 18, BASE + 3 * CAD)},
        satellites=(16, 18), minimum_slots=2)
    assert [pair for _, pair in verified] == [
        (BASE, BASE + 3 * CAD), (BASE + 4 * CAD, BASE + 6 * CAD)]


# --------------------------------------------------------------------- end to end


def test_a_whole_corpus_is_generated_into_the_destination(tmp_path, monkeypatch):
    """The pieces composed: profile -> slices -> shards -> manifest, with the
    destination the only thing written and no library anywhere."""
    from suvi import availability as availability_module, db

    root = tmp_path / "archive"
    out = tmp_path / "dest"
    root.mkdir()

    span_min = 2                       # segments need 2*2+1 = 5 slots
    slots = 24
    times = [BASE + i * CAD for i in range(slots)]
    rows = make_archive(str(root), times, satellites=(16, 18))
    conn = FakeConn(rows)

    states = np.full((2, len(availability_module.BANDS), slots), availability_module.GOOD,
                     dtype=np.uint8)
    profile = availability_module.Availability(
        states=states, t0=BASE, cadence=CAD, satellites=(16, 18))
    profile.save(str(tmp_path / "availability.npz"))

    monkeypatch.setattr(db, "connect", lambda *a, **k: conn)
    monkeypatch.setattr(dataset_build.db, "connect", lambda *a, **k: conn)
    monkeypatch.setattr(dataset_build.paths, "data_root", lambda: str(root))
    monkeypatch.setattr(dataset_build, "holdout_days", lambda *a, **k: [])
    monkeypatch.chdir(tmp_path / "archive")     # cwd must stay untouched

    before = set(os.listdir(root))
    assert dataset_build.main([
        "--availability", str(tmp_path / "availability.npz"),
        "--dest", str(out), "--budget-gb", "1",
        "--span-min", str(span_min), "--val-segments", "0",
        "--test-segments", "0", "--workers", "1",
    ]) == 0

    # Nothing was staged: the archive directory is exactly as it was.
    assert set(os.listdir(root)) == before

    produced = sorted(os.listdir(out))
    assert "manifest.json" in produced and "segments.json" in produced
    manifest = json.loads((out / "manifest.json").read_text())
    assert manifest["satellites"] == [16, 18]
    assert manifest["span_min"] == span_min
    assert manifest["digest"]
    assert manifest["segments"]["train"]

    for satellite in (16, 18):
        name = dataset.shard_name(DAY, satellite)
        assert name in produced
        with dataset.Shard(str(out / name)) as shard:
            assert shard.times()                       # real slots, readable
            assert all(shard.usable(t) for t in shard.times())
            assert shard.raw(shard.times()[0]).shape == (6, 640, 640)


def test_a_resumed_build_skips_finished_shards_and_keeps_their_vetoes(tmp_path):
    """Nothing is local, so the destination is the whole memory of a run.  A
    resume that forgot the vetoes would emit a manifest claiming segments longer
    than the data behind them."""
    out = tmp_path / "dest"
    destination = dataset_build.Destination(str(out))
    destination.put("vetting.json", json.dumps([[DAY, 16, BASE + 4 * CAD]]).encode())

    carried = {tuple(entry) for entry in json.loads(destination.get("vetting.json"))}
    assert carried == {(DAY, 16, BASE + 4 * CAD)}
    verified = dataset_build.verified_segments(
        {"train": [(BASE, BASE + 10 * CAD)]}, carried,
        satellites=(16,), minimum_slots=3)
    assert [pair for _, pair in verified] == [
        (BASE, BASE + 4 * CAD), (BASE + 5 * CAD, BASE + 10 * CAD)]


def test_roles_survive_verification():
    """A split segment stays in the split it was chosen for -- a val fragment that
    drifted into train would unblind the model."""
    verified = dataset_build.verified_segments(
        {"train": [(BASE, BASE + 4 * CAD)],
         "val": [(BASE + 100 * CAD, BASE + 104 * CAD)]},
        vetoed=set(), satellites=(16,), minimum_slots=2)
    assert {role for role, _ in verified} == {"train", "val"}
    for role, pair in verified:
        assert (role == "val") == (pair[0] >= BASE + 100 * CAD)


def test_the_corpus_carries_its_availability_profile(tmp_path, monkeypatch):
    """train.py looks for availability.npz beside the shards and exits without it,
    so a corpus that does not ship it cannot be trained on at all."""
    import dataset_build

    profile = tmp_path / "availability.npz"
    profile.write_bytes(b"not really an npz, but bytes are bytes")
    destination = dataset_build.Destination(str(tmp_path / "dest"))
    destination.put("availability.npz", profile.read_bytes())
    assert (tmp_path / "dest" / "availability.npz").read_bytes() == profile.read_bytes()


def test_an_unreachable_destination_raises_rather_than_reading_as_empty(monkeypatch):
    """A failed listing once came back empty, which reads as 'nothing has been
    built yet' -- that silently restarts a finished 25-hour build from zero."""
    import subprocess as sp
    import dataset_build

    class Failed:
        returncode = 255
        stdout = b""
        stderr = b"ssh: connect to host: No route to host"

    monkeypatch.setattr(dataset_build.subprocess, "run", lambda *a, **k: Failed())
    monkeypatch.setattr(dataset_build.gpubox, "ssh_argv", lambda host=None: ["ssh", host])
    destination = dataset_build.Destination("nowhere:/tmp/x")
    with pytest.raises(OSError, match="cannot list"):
        destination.existing()


def test_remote_streams_use_the_authenticated_ssh_invocation(monkeypatch):
    """The host authenticates by password; a plain `ssh -o BatchMode=yes` cannot
    connect at all, and failed as a broken pipe on the first shard."""
    import dataset_build

    monkeypatch.setattr(dataset_build.gpubox, "_password", lambda: "secret")
    argv = dataset_build.gpubox.ssh_argv("user@host")
    assert argv[0] == "sshpass" and "secret" in argv
    assert argv[-1] == "user@host"
    assert "BatchMode=yes" not in argv


class _Args:
    def __init__(self, retries=2, retry_wait=1, workers=1):
        self.retries, self.retry_wait, self.workers = retries, retry_wait, workers


def test_a_dropped_destination_is_waited_out_and_the_shard_retried(monkeypatch):
    """The host is on WiFi and dropped twice in three days, each time killing a
    build partway through a shard.  Resume was already correct; what was missing
    was anyone to restart it."""
    import dataset_build

    calls = {"build": 0}

    def flaky(conn, root, day, satellite, times, destination, workers):
        calls["build"] += 1
        if calls["build"] == 1:
            raise BrokenPipeError(32, "Broken pipe")
        return (123, 4, {}, [])

    monkeypatch.setattr(dataset_build, "build_shard", flaky)
    monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)

    class Back:
        def existing(self):
            return set()

    result = dataset_build.build_with_retry(
        None, "/root", "2024-01-01", 16, [1], Back(), _Args())
    assert result == (123, 4, {}, [])
    assert calls["build"] == 2          # failed once, retried once


def test_retries_are_bounded_and_the_error_surfaces(monkeypatch):
    """A destination that never returns must eventually fail loudly, not spin."""
    import dataset_build

    def always_fails(*a, **k):
        raise BrokenPipeError(32, "Broken pipe")

    monkeypatch.setattr(dataset_build, "build_shard", always_fails)
    monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)

    class Back:
        def existing(self):
            return set()

    with pytest.raises(BrokenPipeError):
        dataset_build.build_with_retry(
            None, "/root", "2024-01-01", 16, [1], Back(), _Args(retries=2))


def test_an_archive_read_failure_is_not_retried(monkeypatch):
    """Only the destination is retried.  Spinning on a mount that cannot serve a
    file would hide the problem instead of reporting it."""
    import dataset_build

    def bad_archive(*a, **k):
        raise ValueError("corrupt FITS")

    monkeypatch.setattr(dataset_build, "build_shard", bad_archive)
    with pytest.raises(ValueError, match="corrupt FITS"):
        dataset_build.build_with_retry(
            None, "/root", "2024-01-01", 16, [1], None, _Args())


def test_a_build_queued_against_an_absent_host_waits_rather_than_crashing(monkeypatch):
    """The corpus is queued against a host that may be away when queueing happens.
    Crashing at startup means the work waits for a person; waiting means it starts
    itself the moment the host returns."""
    import dataset_build

    state = {"calls": 0}

    class LateHost:
        def existing(self):
            state["calls"] += 1
            if state["calls"] < 3:
                raise OSError("No route to host")
            return {"already.shard"}

    monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)
    got = dataset_build.wait_for_destination(LateHost(), seconds=600, poll=0)
    assert got == {"already.shard"}
    assert state["calls"] == 3


def test_waiting_for_a_destination_is_bounded(monkeypatch):
    """A host that never returns must surface the error, not wait forever."""
    import dataset_build

    class NeverHost:
        def existing(self):
            raise OSError("No route to host")

    times = iter([0, 1, 10_000, 10_001, 10_002])
    monkeypatch.setattr(dataset_build.time, "sleep", lambda s: None)
    monkeypatch.setattr(dataset_build.time, "time", lambda: next(times))
    with pytest.raises(OSError, match="No route to host"):
        dataset_build.wait_for_destination(NeverHost(), seconds=60, poll=0)


def test_a_failed_shard_reaps_its_worker_pool(monkeypatch, tmp_path):
    """A link that drops mid-shard makes this run hundreds of times.  `with Pool`
    only terminates -- it never joins -- so every abandoned imap kept its buffered
    4.9 MB blocks alive and the builder grew 260 MB -> 16.6 GB over 80 attempts."""
    import dataset_build

    reaped = {"terminate": 0, "join": 0, "closed": 0}

    class FakePool:
        def __init__(self, workers):
            pass

        def imap(self, fn, jobs, chunksize=1):
            for job in jobs:
                yield job[2], None, None

        def terminate(self):
            reaped["terminate"] += 1

        def join(self):
            reaped["join"] += 1

    class Exploding:
        def open(self, name, sealed=False):
            raise BrokenPipeError(32, "Broken pipe")

    monkeypatch.setattr(dataset_build, "Pool", FakePool)
    monkeypatch.setattr(dataset_build, "slot_paths", lambda conn, sat, times: {})

    with pytest.raises(BrokenPipeError):
        dataset_build.build_shard(None, "/root", "2024-01-01", 16, [1, 2, 3],
                                  Exploding(), 1)
    assert reaped["terminate"] == 1, "pool must be terminated on failure"
    assert reaped["join"] == 1, "terminate without join leaks the result buffers"


def _truncated_and_whole_shard(tmp_path):
    """Bytes of a real v2 shard, and a prefix of it cut short mid-record."""
    import numpy as np
    from suvi import dataset
    block = np.full((6, 640, 640), 1234, dtype=np.uint16)
    path = tmp_path / "whole.shard"
    dataset.write_shard(str(path), "2024-01-01", 16, (94, 131, 171, 195, 284, 304),
                        {1700000000: block, 1700000240: block})
    whole = path.read_bytes()
    return whole[: len(whole) // 3], whole


def test_the_remote_commit_refuses_a_stream_cut_short(tmp_path):
    """The real failure: the local builder died, ssh forwarded a clean EOF, the
    remote `cat` exited 0, and `mv` committed a 33 MB fragment of a 1.4 GB day
    under its final name.  Runs the exact remote shell command through bash."""
    import subprocess
    import dataset_build

    short, whole = _truncated_and_whole_shard(tmp_path)
    dest = dataset_build.Destination(f"host:{tmp_path / 'remote'}")

    cut = dataset_build._RemoteWrite(dest, "day_g16.shard", sealed=True)
    result = subprocess.run(["bash", "-c", cut.command], input=short,
                            stderr=subprocess.PIPE)
    assert result.returncode != 0
    assert b"truncated" in result.stderr
    assert not (tmp_path / "remote" / "day_g16.shard").exists()
    assert (tmp_path / "remote" / "day_g16.shard.partial").exists()

    ok = dataset_build._RemoteWrite(dest, "day_g16.shard", sealed=True)
    result = subprocess.run(["bash", "-c", ok.command], input=whole)
    assert result.returncode == 0
    assert (tmp_path / "remote" / "day_g16.shard").read_bytes() == whole


def test_unsealed_remote_files_commit_without_a_trailer(tmp_path):
    """manifest.json and friends do not end in the shard magic and must still land."""
    import subprocess
    import dataset_build

    dest = dataset_build.Destination(f"host:{tmp_path / 'remote'}")
    put = dataset_build._RemoteWrite(dest, "manifest.json")
    assert subprocess.run(["bash", "-c", put.command], input=b'{"a": 1}').returncode == 0
    assert (tmp_path / "remote" / "manifest.json").read_bytes() == b'{"a": 1}'


def test_a_local_sealed_write_cut_short_is_not_committed(tmp_path):
    import dataset_build

    short, whole = _truncated_and_whole_shard(tmp_path)
    dest = dataset_build.Destination(str(tmp_path / "local"))
    with pytest.raises(OSError, match="truncated"):
        with dest.open("day_g16.shard", sealed=True) as handle:
            handle.write(short)
    assert not (tmp_path / "local" / "day_g16.shard").exists()
    assert not (tmp_path / "local" / "day_g16.shard.partial").exists()

    with dest.open("day_g16.shard", sealed=True) as handle:
        handle.write(whole)
    assert (tmp_path / "local" / "day_g16.shard").read_bytes() == whole