from __future__ import annotations
from pathlib import Path

import numpy as np, h5py

from sfm.database import SeismoDb
from sfm.waveform_store import WaveformStore
from scripts.backfill_event_shape import backfill_shape
from minimateplus.models import Event, Timestamp, PeakValues

_FIX = Path(__file__).parent / "fixtures/histogram-extension-re/events-5-21-26/K558LL8B.7I0W"


def _event(waveform_key="0111abcd"):
    ev = Event(index=0)
    ev._waveform_key = bytes.fromhex(waveform_key)
    ev.timestamp = Timestamp(raw=b"", flag=0x10, year=2026, unknown_byte=0,
                             month=6, day=25, hour=8, minute=50, second=0)
    ev.record_type = "Waveform"
    ev.peak_values = PeakValues(tran=0.075, vert=0.220, long=0.045,
                                 peak_vector_sum=0.231, micl=0.01)
    return ev


def test_insert_stores_offset_from_record(tmp_path: Path):
    db = SeismoDb(tmp_path / "s.db")
    ev = _event()
    rec = {ev._waveform_key.hex(): {
        "filename": "F.CE0W", "filesize": 10,
        "shape_offset": 1, "shape_offset_axis": "Tran",
        "shape_offset_pre": 0.05, "shape_offset_spread": 0.001}}
    db.insert_events([ev], serial="BE1", waveform_records=rec)
    row = db.query_events(serial="BE1")[0]
    assert row["shape_offset"] == 1
    assert row["shape_offset_axis"] == "Tran"
    assert abs(row["shape_offset_pre"] - 0.05) < 1e-6
    assert abs(row["shape_offset_spread"] - 0.001) < 1e-6


def test_save_imported_bw_attaches_offset(tmp_path: Path):
    store = WaveformStore(tmp_path / "waveforms")
    ev, rec = store.save_imported_bw(_FIX.read_bytes(), source_path=_FIX, serial_hint="BE9558")
    assert rec["shape_offset"] in (0, 1)
    assert rec["shape_offset_axis"] in ("Tran", "Vert", "Long")
    assert "shape_offset_pre" in rec and "shape_offset_spread" in rec


def test_backfill_updates_offset(tmp_path: Path):
    db = SeismoDb(tmp_path / "s.db")
    store = WaveformStore(tmp_path / "waveforms")
    ev = Event(index=0); ev._waveform_key = bytes.fromhex("0111abcd")
    db.insert_events([ev], serial="BE1",
                     waveform_records={ev._waveform_key.hex(): {"filename": "F.CE0W", "filesize": 10}})
    p = store.hdf5_path_for("BE1", "F.CE0W")
    with h5py.File(p, "w") as f:
        g = f.create_group("samples")
        g.create_dataset("Tran", data=np.full(300, 0.05, "float32"))
        g.create_dataset("Vert", data=np.zeros(300, "float32"))
        g.create_dataset("Long", data=np.zeros(300, "float32"))
        f.attrs["pretrig_samples"] = 50
    backfill_shape(db, store)
    row = db.query_events(serial="BE1")[0]
    assert row["shape_offset"] == 1
    assert row["shape_offset_axis"] == "Tran"