import os

import numpy as np
import pytest

from suvi import availability, paths, segments

CAD = paths.CADENCE
BASE = 1715400000 // CAD * CAD


def make_model(pattern16, pattern18):
    # Slot patterns, applied to every band: segments need all six GOOD.
    slots = np.stack([np.asarray(pattern16), np.asarray(pattern18)])
    states = np.repeat(slots[:, None, :], len(availability.BANDS), axis=1)
    return availability.Availability(states=states, t0=BASE, cadence=CAD,
                                     satellites=(16, 18))


def test_segments_are_jointly_pristine_runs_of_at_least_the_minimum():
    g16 = np.full(100, availability.GOOD)
    g18 = np.full(100, availability.GOOD)
    g16[30] = availability.BAD                 # splits the run
    g18[70:80] = availability.MISSING          # second split, other satellite
    found = segments.from_availability(make_model(g16, g18), minimum_slots=10)
    assert found == [
        (BASE, BASE + 30 * CAD),                       # 30 slots
        (BASE + 31 * CAD, BASE + 70 * CAD),            # 39 slots
        (BASE + 80 * CAD, BASE + 100 * CAD),           # 20 slots
    ]
    # A tighter minimum drops the tail fragment.
    found = segments.from_availability(make_model(g16, g18), minimum_slots=35)
    assert found == [(BASE + 31 * CAD, BASE + 70 * CAD)]


def test_days_touched_spans_midnights():
    start = BASE  # 2024-05-11 04:00 UTC
    segment = (start, start + 350 * CAD)       # ~23.3 h -> crosses one midnight
    days = segments.days_touched([segment])
    assert days == ["2024-05-11", "2024-05-12"]


def test_verify_splits_where_extraction_disagrees():
    """disc_v1 can veto a slot the header profile called pristine; the segment
    must split there, and fragments below the model input span must go."""
    segment = (BASE, BASE + 100 * CAD)
    bad = {BASE + 50 * CAD}
    out = segments.verify([segment], lambda t: t not in bad, minimum_slots=30)
    assert out == [(BASE, BASE + 50 * CAD), (BASE + 51 * CAD, BASE + 100 * CAD)]
    out = segments.verify([segment], lambda t: t not in bad, minimum_slots=60)
    assert out == []                            # both fragments too short


def test_targets_keep_a_full_span_inside_the_segment():
    segment = (BASE, BASE + 100 * CAD)
    eligible = segments.targets([segment], span_slots=30)
    assert eligible[0] == BASE + 30 * CAD
    assert eligible[-1] == BASE + 69 * CAD
    assert len(eligible) == 40


def test_split_assigns_whole_segments_and_is_deterministic():
    runs = [(BASE + i * 1000 * CAD, BASE + (i * 1000 + 200) * CAD)
            for i in range(10)]
    first = segments.split(runs, seed=3, val_targets=100, test_targets=100,
                           span_slots=50)
    again = segments.split(runs, seed=3, val_targets=100, test_targets=100,
                           span_slots=50)
    assert first == again
    everything = first["train"] + first["val"] + first["test"]
    assert sorted(everything) == sorted(map(tuple, runs))
    for name in ("val", "test"):
        mass = sum((b - a) // CAD - 100 for a, b in first[name])
        assert mass >= 100
    assert first["train"], "everything landed in val/test"


def test_save_load_round_trip(tmp_path):
    runs = [(BASE, BASE + 200 * CAD), (BASE + 500 * CAD, BASE + 900 * CAD)]
    path = str(tmp_path / "segments.json")
    segments.save(path, runs)
    assert segments.load(path) == runs


def test_containing():
    runs = [(BASE, BASE + 100 * CAD)]
    assert segments.containing(runs, BASE + 50 * CAD) == runs[0]
    assert segments.containing(runs, BASE - CAD) is None


def test_ladder_for_matches_spans():
    from suvi import samples

    assert samples.ladder_for(64) == (1, 4, 16, 64)
    assert samples.ladder_for(128) == (1, 4, 16, 64, 128)
    assert samples.ladder_for(512) == (1, 4, 16, 64, 128, 256, 512)
    with pytest.raises(ValueError):
        samples.ladder_for(100)


def test_saved_segments_are_plain_json_ints(tmp_path):
    """numpy int64 is not JSON serializable, and the failure lands at the very
    end of a multi-hour corpus build."""
    model = make_model(np.full(600, availability.GOOD), np.full(600, availability.GOOD))
    found = segments.from_availability(model, minimum_slots=10)
    assert all(type(v) is int for pair in found for v in pair)
    path = str(tmp_path / "segments.json")
    segments.save(path, found)
    assert segments.load(path) == found
    assert not os.path.exists(path + ".tmp")