"""Drawing the captures nobody can read.

A waterfall is checked here the way a picture has to be: by reading the
pixels back and asking whether the thing that was put in shows up where it
should, rather than by looking at the file and calling it a picture.
"""
import json
import wave
from pathlib import Path

import numpy as np
import pytest

from bandsaunter import waterfall as wf
from bandsaunter.images import PNG_SIGNATURE

FS = 16000


def tone(hz: float, seconds: float = 4.0, rate: int = FS, level: float = 0.5,
         seed: int = 0, keyed: bool = False) -> np.ndarray:
    rng = np.random.default_rng(seed)
    t = np.arange(int(seconds * rate)) / rate
    x = np.sin(2 * np.pi * hz * t) * level
    if keyed:
        x *= (np.sin(2 * np.pi * 1.5 * t) > 0)
    return x + 0.01 * rng.standard_normal(t.size)


def spectrum_body(canvas: np.ndarray) -> np.ndarray:
    """Just the drawn spectrum, with the margins and axes cut away."""
    return canvas[wf.TOP:canvas.shape[0] - wf.BOTTOM,
                  wf.LEFT:wf.LEFT + wf.WIDTH]


def brightest_column(canvas: np.ndarray) -> int:
    """Which column of the drawn body is loudest, in body coordinates."""
    body = spectrum_body(canvas)
    return int(np.argmax(body.astype(np.float64).sum(axis=(0, 2))))


# -- what it draws -----------------------------------------------------------

def test_a_tone_lands_at_its_own_frequency():
    """The whole point: where the picture is bright is where the signal was."""
    canvas, drawn = wf.render_waterfall(tone(2000.0), FS)
    hz = brightest_column(canvas) / wf.WIDTH * (FS / 2.0)
    assert abs(hz - 2000.0) < 120.0, f"the tone was drawn at {hz:.0f} Hz"


def test_two_tones_are_two_stripes():
    audio = tone(1000.0) + tone(3000.0, seed=1)
    canvas, _ = wf.render_waterfall(audio, FS)
    profile = spectrum_body(canvas).astype(np.float64).sum(axis=(0, 2))
    at = lambda hz: int(hz / (FS / 2.0) * wf.WIDTH)      # noqa: E731
    assert profile[at(1000)] > 3 * profile[at(2000)]
    assert profile[at(3000)] > 3 * profile[at(2000)]


def test_keying_shows_as_gaps_down_the_stripe():
    """Time runs down the picture, so a keyed carrier is a dashed line."""
    canvas, _ = wf.render_waterfall(tone(1500.0, seconds=6.0, keyed=True), FS)
    body = spectrum_body(canvas)
    column = body[:, brightest_column(canvas)].astype(np.float64).sum(axis=1)
    on = column > (column.max() + column.min()) / 2
    changes = int(np.count_nonzero(np.diff(on.astype(np.int8))))
    assert changes >= 8, f"the keying drew as {changes} edges, not a dashed line"


def test_the_picture_has_room_for_its_axes():
    canvas, drawn = wf.render_waterfall(tone(1200.0), FS)
    assert canvas.shape == (drawn.height, drawn.width, 3)
    assert drawn.width > wf.WIDTH and drawn.height > drawn.rows
    assert drawn.seconds == pytest.approx(4.0, abs=0.05)


def test_time_and_frequency_scales_are_actually_drawn():
    """Ink outside the body, or the axes are empty margins."""
    canvas, _ = wf.render_waterfall(tone(1200.0), FS)
    left_margin = canvas[:, :wf.LEFT]
    top_margin = canvas[:wf.TOP, :]
    assert (left_margin != np.array(wf.BACKGROUND, np.uint8)).any()
    assert (top_margin != np.array(wf.BACKGROUND, np.uint8)).any()


def test_a_caption_is_drawn_under_the_picture():
    plain, _ = wf.render_waterfall(tone(1200.0), FS)
    with_text, _ = wf.render_waterfall(tone(1200.0), FS, caption="146.520 MHZ")
    bottom = slice(plain.shape[0] - wf.BOTTOM, plain.shape[0])
    assert not np.array_equal(plain[bottom], with_text[bottom])


def test_an_empty_capture_is_refused_rather_than_drawn_blank():
    with pytest.raises(ValueError):
        wf.render_waterfall(np.zeros(0), FS)


# -- IQ says radio frequency, audio says audio -------------------------------

def test_raw_iq_is_drawn_around_the_tuned_frequency():
    """Where a scan kept the IQ, the picture is the spectrum an operator
    would have been watching -- both sides of the carrier, not one."""
    rate = 48000.0
    t = np.arange(int(2.0 * rate)) / rate
    rng = np.random.default_rng(4)
    iq = (np.exp(2j * np.pi * -8000.0 * t)
          + 0.01 * (rng.standard_normal(t.size)
                    + 1j * rng.standard_normal(t.size))).astype(np.complex64)
    canvas, drawn = wf.render_waterfall(iq, rate, complex_input=True,
                                        centre_hz=146.52e6)
    assert drawn.source == "iq"
    assert drawn.centre_hz == pytest.approx(146.52e6)
    hz = (brightest_column(canvas) / wf.WIDTH - 0.5) * rate
    assert abs(hz + 8000.0) < 500.0, f"drawn at {hz:.0f} Hz from centre"


def test_the_caption_says_which_picture_it_is():
    """After an FM detector the frequency axis is audio, not radio, and a
    picture that did not say so would be a lie told in a convincing font."""
    assert "DEMODULATED AUDIO" in wf.caption_for(146.52e6, "nfm", 4.0, "", "audio")
    assert "RF SPECTRUM" in wf.caption_for(146.52e6, "nfm", 4.0, "", "iq")
    said = wf.caption_for(146.52e6, "nfm", 12.5, "OOK / ASK data burst", "audio")
    assert "146.520000 MHZ" in said and "NFM" in said and "12.5 S" in said
    assert "OOK / ASK DATA BURST" in said


# -- the file ---------------------------------------------------------------

def test_it_writes_a_real_png(tmp_path):
    out = tmp_path / "capture_waterfall.png"
    drawn = wf.write_waterfall(out, tone(1200.0), FS, caption="TEST")
    assert out.read_bytes()[:8] == PNG_SIGNATURE
    assert drawn.path == str(out)
    assert "demodulated audio" in drawn.summary()


def test_the_picture_is_named_for_the_recording(tmp_path):
    assert wf.waterfall_path(tmp_path / "0146.520000MHz--x-nfm.wav").name == \
        "0146.520000MHz--x-nfm_waterfall.png"


def test_a_recording_with_no_samples_draws_nothing(tmp_path):
    assert wf.draw_for_recording(tmp_path / "x.wav", audio=None, rate=0) is None
    assert wf.draw_for_recording(tmp_path / "x.wav",
                                 audio=np.zeros(4), rate=FS) is None


def test_the_iq_is_preferred_where_a_scan_kept_it(tmp_path):
    """The audio is what is left after a detector threw most of it away."""
    rate = 48000.0
    t = np.arange(int(1.0 * rate)) / rate
    rng = np.random.default_rng(5)
    iq = (np.exp(2j * np.pi * 5000.0 * t)
          + 0.01 * (rng.standard_normal(t.size)
                    + 1j * rng.standard_normal(t.size))).astype(np.complex64)
    raw = tmp_path / "cap.cf32"
    iq.tofile(raw)
    drawn = wf.draw_for_recording(tmp_path / "cap.wav", audio=tone(1200.0),
                                  rate=FS, iq_path=str(raw), iq_rate=rate,
                                  frequency=146.52e6)
    assert drawn is not None and drawn.source == "iq"
    assert drawn.seconds == pytest.approx(1.0, abs=0.02)


def test_unreadable_iq_falls_back_to_the_audio(tmp_path):
    broken = tmp_path / "cap.cf32"
    broken.write_bytes(b"")
    drawn = wf.draw_for_recording(tmp_path / "cap.wav", audio=tone(1200.0),
                                  rate=FS, iq_path=str(broken), iq_rate=48000.0)
    assert drawn is not None and drawn.source == "audio"


# -- the text -----------------------------------------------------------------

def test_an_unknown_character_is_a_space_not_a_failure():
    """A label is a convenience; a picture refused because of one odd
    character in a caption would be a poor trade."""
    canvas = np.zeros((20, 200, 3), dtype=np.uint8)
    wf.draw_text(canvas, 2, 2, "AéB")
    assert canvas.any()


def test_a_label_that_runs_off_the_edge_is_clipped_not_wrapped():
    canvas = np.zeros((20, 40, 3), dtype=np.uint8)
    wf.draw_text(canvas, 30, 2, "1234567890")
    assert not canvas[10:].any()      # nothing wrapped onto a later row


# -- what a scan does with it ------------------------------------------------

def _scan(tmp_path, said: str | None, **over):
    """One capture of one voice transmission, with a stubbed recogniser.

    ``said`` is what the recogniser comes back with; None means no recogniser
    is installed at all, which is the commonest case in the wild.
    """
    from bandsaunter import transcribe as tr
    from bandsaunter.config import ScanConfig
    from bandsaunter.ranges import parse_range_list
    from bandsaunter.scanner import Scanner
    from bandsaunter.simulator import SimulatedDevice, VirtualTransmitter as V

    cfg = ScanConfig(ranges=parse_range_list("146.4M-146.6M"),
                     output_dir=str(tmp_path), record_seconds=4.0,
                     hang_seconds=1.0, threshold_db=12, dwell_seconds=0.05,
                     max_cycles=1, revisit_seconds=0.2,
                     transcribe=said is not None,
                     transcribe_engine="fake", callsign_lookup=False, **over)
    scanner = Scanner(cfg, device=SimulatedDevice(
        transmitters=[V(146_520_000, "nfm", 0.4, 12_500, "v")]).open())
    scanner.prepare()
    scanner.run()
    return scanner


@pytest.fixture
def recogniser(monkeypatch):
    """A stub recogniser whose answer each test chooses."""
    from bandsaunter import transcribe as tr
    said = {"text": ""}

    monkeypatch.setitem(
        tr._DISPATCH, "fake",
        lambda audio, rate, model, lang: tr.Transcript(text=said["text"],
                                                       engine="fake"))
    monkeypatch.setattr(tr, "ENGINES", ("fake",) + tr.ENGINES)
    monkeypatch.setattr(tr, "_is_present", lambda name: name == "fake")
    monkeypatch.setattr("bandsaunter.scanner.available_engine", lambda: "fake")
    # Nothing in this file is about whether a clip has signal in it.
    monkeypatch.setattr(tr, "_has_signal", lambda audio, rate: True)
    return said


def drawings(tmp_path) -> list[Path]:
    return sorted(tmp_path.glob("*_waterfall.png"))


def test_a_capture_with_no_recogniser_is_drawn(tmp_path):
    """The commonest case in the wild: nothing installed, so nothing can be
    read, so everything gets a picture."""
    scanner = _scan(tmp_path, None)
    assert scanner.stats.recordings >= 1
    assert len(drawings(tmp_path)) == scanner.stats.recordings


def test_a_capture_the_recogniser_could_read_is_not_drawn(tmp_path,
                                                          recogniser):
    recogniser["text"] = "Net control, this is W1AW, standing by."
    scanner = _scan(tmp_path, recogniser["text"])
    assert scanner.stats.recordings >= 1
    assert drawings(tmp_path) == []


def test_a_capture_the_recogniser_could_not_read_is_drawn(tmp_path,
                                                          recogniser):
    """A recogniser handed something that is not speech comes back with a
    word or two of nothing in particular, and a picture is worth more."""
    recogniser["text"] = "You"
    scanner = _scan(tmp_path, recogniser["text"])
    assert scanner.stats.recordings >= 1
    assert len(drawings(tmp_path)) == scanner.stats.recordings


def test_a_capture_that_produced_no_words_at_all_is_drawn(tmp_path,
                                                          recogniser):
    """The recogniser writes no transcript file for a silent capture, so
    this is the case that had nothing at all to show for it."""
    recogniser["text"] = ""
    scanner = _scan(tmp_path, "")
    assert scanner.stats.recordings >= 1
    assert len(drawings(tmp_path)) == scanner.stats.recordings
    assert not list(tmp_path.glob("*_transcription.txt"))


def test_the_picture_is_recorded_in_the_sidecar(tmp_path):
    _scan(tmp_path, None)
    for meta in tmp_path.glob("*.json"):
        if meta.name.startswith("scan_log"):
            continue
        hit = json.loads(meta.read_text()).get("hit") or {}
        assert hit.get("waterfall_path"), meta.name
        assert Path(hit["waterfall_path"]).is_file()


def test_switching_it_off_draws_nothing(tmp_path):
    scanner = _scan(tmp_path, None, waterfall=False)
    assert scanner.stats.recordings >= 1
    assert drawings(tmp_path) == []


def test_the_bar_for_readable_can_be_moved(tmp_path, recogniser):
    """Five characters is a default, not a law."""
    recogniser["text"] = "Roger"
    scanner = _scan(tmp_path, recogniser["text"], waterfall_min_chars=40)
    assert scanner.stats.recordings >= 1
    assert len(drawings(tmp_path)) == scanner.stats.recordings


# -- and for a directory already recorded ------------------------------------

def _recording(directory: Path, name: str, audio: np.ndarray,
               rate: int = FS, hit: dict | None = None,
               transcript: str = "") -> Path:
    wav = directory / f"{name}.wav"
    with wave.open(str(wav), "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(rate)
        w.writeframes((np.clip(audio, -1, 1) * 32000).astype("<i2").tobytes())
    if hit is not None:
        (directory / f"{name}.json").write_text(json.dumps({"hit": hit}))
    if transcript:
        (directory / f"{name}_transcription.txt").write_text(transcript + "\n")
    return wav


def _waterfall_command(*args) -> int:
    from bandsaunter.cli import main
    return main(["waterfall", *args])


def test_the_command_draws_only_what_cannot_be_read(tmp_path, monkeypatch):
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0146.520000MHz--a-nfm", tone(1200.0),
               hit={"category": "voice", "frequency": 146.52e6, "mode": "nfm"},
               transcript="net control this is W1AW standing by")
    _recording(tmp_path, "0146.940000MHz--b-ook", tone(2400.0, keyed=True),
               hit={"category": "digital", "frequency": 146.94e6,
                    "mode": "nfm", "classification": "OOK / ASK data burst"})
    assert _waterfall_command(str(tmp_path)) == 0
    drawn = [p.name for p in drawings(tmp_path)]
    assert drawn == ["0146.940000MHz--b-ook_waterfall.png"]


def test_short_words_do_not_count_as_having_read_it(tmp_path, monkeypatch):
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0146.520000MHz--a-nfm", tone(1200.0),
               hit={"category": "voice", "frequency": 146.52e6}, transcript="You")
    assert _waterfall_command(str(tmp_path)) == 0
    assert len(drawings(tmp_path)) == 1


def test_all_draws_the_readable_ones_too(tmp_path, monkeypatch):
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0146.520000MHz--a-nfm", tone(1200.0),
               hit={"category": "voice", "frequency": 146.52e6},
               transcript="a whole sentence of perfectly good speech")
    assert _waterfall_command(str(tmp_path), "--all") == 0
    assert len(drawings(tmp_path)) == 1


def test_it_does_not_redraw_what_it_already_drew(tmp_path, monkeypatch):
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    wav = _recording(tmp_path, "0146.940000MHz--b-ook", tone(2400.0),
                     hit={"category": "digital", "frequency": 146.94e6})
    assert _waterfall_command(str(tmp_path)) == 0
    first = wf.waterfall_path(wav).stat().st_mtime_ns
    assert _waterfall_command(str(tmp_path)) == 0
    assert wf.waterfall_path(wav).stat().st_mtime_ns == first
    assert _waterfall_command(str(tmp_path), "--redraw") == 0
    assert wf.waterfall_path(wav).stat().st_mtime_ns != first


def test_a_recording_with_no_sidecar_is_drawn(tmp_path, monkeypatch):
    """Nothing is known about it, which is the strongest reason to draw it."""
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0146.940000MHz--b-nfm", tone(2400.0))
    assert _waterfall_command(str(tmp_path)) == 0
    assert len(drawings(tmp_path)) == 1


def test_the_command_notes_the_picture_in_the_sidecar(tmp_path, monkeypatch):
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0146.940000MHz--b-ook", tone(2400.0),
               hit={"category": "digital", "frequency": 146.94e6})
    assert _waterfall_command(str(tmp_path)) == 0
    hit = json.loads((tmp_path / "0146.940000MHz--b-ook.json").read_text())["hit"]
    assert Path(hit["waterfall_path"]).is_file()


# -- Morse is never words ----------------------------------------------------

def test_a_transcript_of_morse_tones_is_not_words():
    """Whisper renders a CW ident as one digit per tone: sixteen characters
    of nothing, which clears the bar and means nothing."""
    hit = {"category": "voice", "morse_complete": "KSQ330"}
    assert not wf.is_readable(hit, "2-2-2-3-3-5-2-7-0-5-9-7-0-8-1-0", 5)


def test_words_over_the_top_of_an_ident_are_still_words():
    assert wf.is_readable({"category": "voice"},
                          "net control this is W1AW standing by", 5)


def test_the_partial_morse_text_counts_as_much_as_the_complete_one():
    """A keyed carrier whose timing only half resolved is still a keyed
    carrier, and the transcript of it is still digits."""
    hit = {"category": "voice", "morse_text": "K SQ 3 30"}
    assert not wf.is_readable(hit, "3-3-0-5-9-7-0-8", 5)


def test_a_capture_that_identified_itself_in_morse_is_drawn(tmp_path,
                                                            recogniser,
                                                            monkeypatch):
    """The whole point: an FM carrier with a CW ident on the end of it gets
    a picture, however long a string of digits the recogniser made of it."""
    from bandsaunter.scanner import Scanner

    class _Ident:
        is_morse = True
        text = "KSQ330"
        complete_text = "KSQ330"
        wpm = 20.0
        confidence = 0.9

    monkeypatch.setattr(Scanner, "_morse_from_recording",
                        lambda self, rec, morse: _Ident())
    recogniser["text"] = "2-2-2-3-3-5-2-7-0-5-9-7-0-8-1-0"
    scanner = _scan(tmp_path, recogniser["text"])
    assert scanner.stats.recordings >= 1
    assert len(drawings(tmp_path)) == scanner.stats.recordings
    for meta in tmp_path.glob("*.json"):
        if meta.name.startswith("scan_log"):
            continue
        assert json.loads(meta.read_text())["hit"]["morse_complete"] == "KSQ330"


def test_the_command_draws_an_ident_it_finds_in_the_sidecar(tmp_path,
                                                            monkeypatch):
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0154.369068MHz--a-nfm", tone(1200.0),
               hit={"category": "voice", "frequency": 154.369068e6,
                    "mode": "nfm", "morse_complete": "KSQ330"},
               transcript="2-2-2-3-3-5-2-7-0-5-9-7-0-8-1-0")
    assert _waterfall_command(str(tmp_path)) == 0
    assert len(drawings(tmp_path)) == 1


def test_check_morse_listens_to_what_the_sidecar_calls_readable(tmp_path,
                                                                monkeypatch):
    """A sidecar written before the CW decoder could hear an ident over an FM
    carrier calls a repeater readable.  Listening again finds it."""
    from morse_gen import morse_audio

    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    ident = morse_audio("DE KSQ330 KSQ330", wpm=20, fs=FS, snr_db=25)
    _recording(tmp_path, "0154.369068MHz--a-nfm", ident,
               hit={"category": "voice", "frequency": 154.369068e6,
                    "mode": "nfm", "morse_text": "", "morse_complete": ""},
               transcript="2-2-2-3-3-5-2-7-0-5-9-7-0-8-1-0")
    assert _waterfall_command(str(tmp_path)) == 0
    assert drawings(tmp_path) == []            # the sidecar was believed
    assert _waterfall_command(str(tmp_path), "--check-morse") == 0
    assert len(drawings(tmp_path)) == 1
    hit = json.loads((tmp_path / "0154.369068MHz--a-nfm.json").read_text())["hit"]
    assert "KSQ330" in hit["morse_complete"]


def test_check_morse_leaves_a_conversation_alone(tmp_path, monkeypatch):
    """It has to be able to say no, or it is just --all with a long wait."""
    monkeypatch.setenv("XDG_CONFIG_HOME", str(tmp_path / "cfg"))
    _recording(tmp_path, "0146.520000MHz--a-nfm", tone(1200.0),
               hit={"category": "voice", "frequency": 146.52e6},
               transcript="net control this is W1AW standing by")
    assert _waterfall_command(str(tmp_path), "--check-morse") == 0
    assert drawings(tmp_path) == []