from datetime import datetime, timedelta
from pathlib import Path

import pytest
from numpy import dtype, float64, int64
from pandas import Int64Dtype

import vallenae as vae
from vallenae.io import HitRecord, MarkerRecord, ParametricRecord, SetType, StatusRecord

STEEL_PLATE_DIR = Path(__file__).resolve().parent / "../examples/steel_plate"
PRIDB_FILE_PATH = STEEL_PLATE_DIR / "sample.pridb"

LABELS_EXPECTED = [
    "",
    "10:52 Resume",
    "2019-09-20 10:54:52",
    "TimeZone: +02:00 (W. Europe Standard Time)",
    "10:56 Suspend",
]

# Caution: units are µV and µs (except time)
HITS_EXPECTED = [
    HitRecord(
        set_id=10,
        time=3.992771,
        channel=3,
        status=4,
        param_id=4,
        threshold=100.469451,
        amplitude=46538.6675710947,
        rise_time=49.5,
        duration=19671.6,
        energy=27995100.3052212,
        signal_strength=29064.1915258135,
        rms=5.05799020605523,
        counts=2180,
        trai=1,
        cascade_hits=4,
        cascade_counts=2193,
        cascade_energy=27995448.6329966,
        cascade_signal_strength=29131.710036357,
    ),
    HitRecord(
        set_id=11,
        time=3.9927747,
        channel=2,
        status=4,
        param_id=3,
        threshold=100.469451,
        amplitude=59621.0079186672,
        rise_time=192.1,
        duration=20575.7,
        energy=22762787.1460552,
        signal_strength=24734.02914512,
        rms=5.02806126992473,
        counts=2047,
        trai=2,
        cascade_hits=4,
        cascade_counts=2054,
        cascade_energy=22762918.7161395,
        cascade_signal_strength=24761.7395309888,
    ),
    HitRecord(
        set_id=12,
        time=3.9928129,
        channel=4,
        status=4,
        param_id=5,
        threshold=100.469451,
        amplitude=34118.5122422787,
        rise_time=354.6,
        duration=19130.7,
        energy=12867002.8866017,
        signal_strength=19956.3935214563,
        rms=4.8085824049677,
        counts=1854,
        trai=3,
        cascade_hits=4,
        cascade_counts=1858,
        cascade_energy=12867202.6824819,
        cascade_signal_strength=19997.9738662233,
    ),
    HitRecord(
        set_id=13,
        time=3.9928143,
        channel=1,
        status=4,
        param_id=2,
        threshold=100.469451,
        amplitude=29114.8291235365,
        rise_time=160.3,
        duration=19266.1,
        energy=12652748.5220513,
        signal_strength=20545.4817939854,
        rms=4.90834552540271,
        counts=1985,
        trai=4,
        cascade_hits=5,
        cascade_counts=1988,
        cascade_energy=12652874.448889,
        cascade_signal_strength=20573.5164221492,
    ),
]

PARAMETRIC_EXPECTED = [
    ParametricRecord(set_id=5, time=0.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=6, time=1.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=7, time=2.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=8, time=3.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=9, time=3.99, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=14, time=4.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=15, time=5.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=16, time=6.0, status=12, param_id=1, pctd=0, pcta=0),
    ParametricRecord(set_id=17, time=6.45, status=0, param_id=1, pctd=0, pcta=0),
]


@pytest.fixture(name="sample_pridb")
def fixture_sample_pridb() -> vae.io.PriDatabase:
    pridb = vae.io.PriDatabase(PRIDB_FILE_PATH)
    yield pridb
    pridb.close()


@pytest.fixture(name="fresh_pridb")
def fixture_fresh_pridb(tmp_path) -> vae.io.PriDatabase:
    filename = tmp_path / "test.pridb"
    with vae.io.PriDatabase(filename, mode="rwc") as pridb:
        con = pridb.connection()
        con.execute(
            """
            INSERT INTO ae_params (
                ID, SetupID, Chan, ADC_µV, ADC_TE, ADC_SS
            ) VALUES (
                1, 1, 1, 1, 1, 1
            )
            """
        )
        yield pridb


def test_init():
    pridb = vae.io.PriDatabase(PRIDB_FILE_PATH)
    pridb.close()


def test_create(tmp_path):
    filename = tmp_path / "empty.pridb"
    vae.io.PriDatabase.create(filename)
    with vae.io.PriDatabase(filename) as pridb:
        assert pridb.tables() == {
            "acq_setup",
            "ae_data",
            "ae_fieldinfo",
            "ae_globalinfo",
            "ae_markers",
            "ae_params",
            "data_integrity",
        }


def test_channel(sample_pridb):
    assert sample_pridb.channel() == {1, 2, 3, 4}


def test_to_datetime(fresh_pridb):
    with pytest.raises(RuntimeError):
        fresh_pridb.to_datetime(0.0)

    fresh_pridb.write_marker(
        MarkerRecord(time=0.0, set_type=SetType.DATETIME, number=1, data="2025-01-02 12:13:14")
    )
    dt1 = datetime(year=2025, month=1, day=2, hour=12, minute=13, second=14)

    fresh_pridb.write_marker(
        MarkerRecord(time=10.0, set_type=SetType.DATETIME, number=1, data="2025-01-02 15:16:17")
    )
    dt2 = datetime(year=2025, month=1, day=2, hour=15, minute=16, second=17)

    assert fresh_pridb.to_datetime(0.0) == dt1
    assert fresh_pridb.to_datetime(1.1) == dt1 + timedelta(seconds=1.1)
    assert fresh_pridb.to_datetime(10.0) == dt2
    assert fresh_pridb.to_datetime(11.1) == dt2 + timedelta(seconds=1.1)


def test_iread_markers(sample_pridb):
    iterable = sample_pridb.iread_markers()
    labels = [marker.data for marker in iterable]

    assert len(iterable) == len(LABELS_EXPECTED)
    assert len(labels) == len(LABELS_EXPECTED)
    assert labels == LABELS_EXPECTED


def test_iread_markers_query_filter(sample_pridb):
    markers = list(sample_pridb.iread_markers(query_filter="Data LIKE '%TimeZone%'"))
    assert len(markers) == 1


def test_read_markers(sample_pridb):
    markers = sample_pridb.read_markers()

    assert len(markers) == len(LABELS_EXPECTED)
    assert markers.index.name == "set_id"
    assert markers.index.dtype == int64
    assert dict(markers.dtypes) == {
        "time": float64,
        "set_type": Int64Dtype(),
        "number": Int64Dtype(),
        "data": dtype("O"),
    }

    labels = list(markers["data"])
    assert labels == LABELS_EXPECTED


def test_iread_hits(sample_pridb):
    hits = list(sample_pridb.iread_hits())

    assert len(hits) == len(HITS_EXPECTED)

    for hit, hit_expected in zip(hits, HITS_EXPECTED):
        assert hit.set_id == hit_expected.set_id
        assert hit.time == pytest.approx(hit_expected.time)
        assert hit.channel == hit_expected.channel
        assert hit.status == hit_expected.status
        assert hit.param_id == hit_expected.param_id
        assert hit.threshold == pytest.approx(hit_expected.threshold / 1e6)
        assert hit.amplitude == pytest.approx(hit_expected.amplitude / 1e6)
        assert hit.rise_time == pytest.approx(hit_expected.rise_time / 1e6)
        assert hit.duration == pytest.approx(hit_expected.duration / 1e6)
        assert hit.energy == pytest.approx(hit_expected.energy)
        assert hit.signal_strength == pytest.approx(hit_expected.signal_strength)
        assert hit.rms == pytest.approx(hit_expected.rms / 1e6)
        assert hit.counts == hit_expected.counts
        assert hit.trai == hit_expected.trai
        assert hit.cascade_hits == hit_expected.cascade_hits
        assert hit.cascade_counts == hit_expected.cascade_counts
        assert hit.cascade_energy == pytest.approx(hit_expected.cascade_energy)
        assert hit.cascade_signal_strength == pytest.approx(hit_expected.cascade_signal_strength)


def test_iread_hits_query_filter(sample_pridb):
    hits = list(sample_pridb.iread_hits(query_filter="SetID >= 12"))
    assert len(hits) == 2


def test_read_hits(sample_pridb):
    hits = sample_pridb.read_hits()

    assert len(hits) == len(HITS_EXPECTED)
    assert hits.index.name == "set_id"
    assert hits.index.dtype == int64
    assert dict(hits.dtypes) == {
        "time": float64,
        "channel": Int64Dtype(),
        "status": Int64Dtype(),
        "param_id": Int64Dtype(),
        "threshold": float64,
        "amplitude": float64,
        "rise_time": float64,
        "cascade_counts": Int64Dtype(),
        "cascade_energy": float64,
        "cascade_hits": Int64Dtype(),
        "cascade_signal_strength": float64,
        "counts": Int64Dtype(),
        "duration": float64,
        "energy": float64,
        "rms": float64,
        "signal_strength": float64,
        "trai": Int64Dtype(),
    }


def test_iread_parametric(sample_pridb):
    records = list(sample_pridb.iread_parametric())
    assert len(records) == len(PARAMETRIC_EXPECTED)

    for record, record_expected in zip(records, PARAMETRIC_EXPECTED):
        assert record.set_id == record_expected.set_id
        assert record.time == pytest.approx(record_expected.time)
        assert record.status == record_expected.status
        assert record.param_id == record_expected.param_id
        assert record.pctd == record_expected.pctd
        assert record.pcta == record_expected.pcta
        assert record.pa0 == record_expected.pa0
        assert record.pa1 == record_expected.pa1
        assert record.pa2 == record_expected.pa2
        assert record.pa3 == record_expected.pa3
        assert record.pa4 == record_expected.pa4
        assert record.pa5 == record_expected.pa5
        assert record.pa6 == record_expected.pa6
        assert record.pa7 == record_expected.pa7


def test_iread_parametric_query_filter(sample_pridb):
    records = list(sample_pridb.iread_parametric(query_filter="SetID > 10"))
    assert len(records) == 4


def test_read_parametric(sample_pridb):
    df = sample_pridb.read_parametric()

    assert len(df) == len(PARAMETRIC_EXPECTED)
    assert df.index.name == "set_id"
    assert df.index.dtype == int64
    assert dict(df.dtypes) == {
        "time": float64,
        "status": Int64Dtype(),
        "param_id": Int64Dtype(),
        "pctd": Int64Dtype(),
        "pcta": Int64Dtype(),
    }


def test_read(sample_pridb):
    df = sample_pridb.read()

    assert df.index.name == "set_id"
    assert df.index.dtype == int64
    assert dict(df.dtypes) == {
        "set_type": int64,
        "time": float64,
        "channel": Int64Dtype(),
        "status": Int64Dtype(),
        "param_id": Int64Dtype(),
        "threshold": float64,
        "amplitude": float64,
        "rise_time": float64,
        "cascade_counts": Int64Dtype(),
        "cascade_energy": float64,
        "cascade_hits": Int64Dtype(),
        "cascade_signal_strength": float64,
        "counts": Int64Dtype(),
        "duration": float64,
        "energy": float64,
        "rms": float64,
        "signal_strength": float64,
        "trai": Int64Dtype(),
        "pctd": Int64Dtype(),
        "pcta": Int64Dtype(),
    }


def test_listen(sample_pridb):
    assert len(list(sample_pridb.listen())) == 0
    assert len(list(sample_pridb.listen(existing=True))) == 18

    def records_by_type(type_):
        return list(filter(lambda r: isinstance(r, type_), sample_pridb.listen(existing=True)))

    assert records_by_type(HitRecord) == list(sample_pridb.iread_hits())
    assert records_by_type(MarkerRecord) == list(sample_pridb.iread_markers())
    assert records_by_type(ParametricRecord) == list(sample_pridb.iread_parametric())
    assert records_by_type(StatusRecord) == list(sample_pridb.iread_status())


def test_write_readonly(sample_pridb):
    with pytest.raises(ValueError):
        sample_pridb.write_marker(
            MarkerRecord(set_id=1, time=0, set_type=SetType.LABEL, number=1, data="Test")
        )


def test_write_hit(fresh_pridb):
    new_hit = HitRecord(
        time=5,
        channel=1,
        param_id=1,
        threshold=111,
        amplitude=22222,
        rise_time=111,
        duration=11111,
        energy=12345678,
        signal_strength=20000,
        rms=5,
        counts=1987,
    )

    assert fresh_pridb.write_hit(new_hit) == 1
    assert fresh_pridb.rows() == 1

    hit_read = next(iter(fresh_pridb.iread_hits()))
    assert hit_read.set_id == 1  # first data set
    assert hit_read.time == new_hit.time
    assert hit_read.channel == new_hit.channel
    assert hit_read.param_id == new_hit.param_id
    assert hit_read.threshold == new_hit.threshold
    assert hit_read.amplitude == new_hit.amplitude
    assert hit_read.rise_time == new_hit.rise_time
    assert hit_read.duration == new_hit.duration
    assert hit_read.energy == new_hit.energy
    assert hit_read.signal_strength == new_hit.signal_strength
    assert hit_read.rms == pytest.approx(new_hit.rms)
    assert hit_read.counts == new_hit.counts
    assert hit_read.trai == new_hit.trai
    assert hit_read.cascade_hits == new_hit.cascade_hits
    assert hit_read.cascade_counts == new_hit.cascade_counts
    assert hit_read.cascade_energy == new_hit.cascade_energy
    assert hit_read.cascade_signal_strength == new_hit.cascade_signal_strength


def test_write_marker(fresh_pridb):
    new_marker = MarkerRecord(
        time=11.11,
        set_type=SetType.LABEL,
        number=1,
        data="Test label",
    )

    assert fresh_pridb.write_marker(new_marker) == 1
    assert fresh_pridb.rows() == 1

    marker_read = next(iter(fresh_pridb.iread_markers()))
    assert marker_read.set_id == 1  # first data set
    assert marker_read.time == new_marker.time
    assert marker_read.set_type == new_marker.set_type
    assert marker_read.number == new_marker.number
    assert marker_read.data == new_marker.data


def test_write_status(fresh_pridb):
    new_status = StatusRecord(
        time=11.11,
        channel=1,
        param_id=1,
        threshold=100,
        energy=111,
        signal_strength=222,
        rms=5,
    )

    assert fresh_pridb.write_status(new_status) == 1
    assert fresh_pridb.rows() == 1

    status_read = next(iter(fresh_pridb.iread_status()))
    assert status_read.set_id == 1  # first data set
    assert status_read.time == new_status.time
    assert status_read.channel == new_status.channel
    assert status_read.param_id == new_status.param_id
    assert status_read.threshold == new_status.threshold
    assert status_read.energy == new_status.energy
    assert status_read.signal_strength == new_status.signal_strength
    assert status_read.rms == pytest.approx(new_status.rms)


def test_write_parametric(fresh_pridb):
    new_parametric = ParametricRecord(
        time=11.11,
        param_id=1,
        pctd=11,
        pcta=22,
    )

    assert fresh_pridb.write_parametric(new_parametric) == 1
    assert fresh_pridb.rows() == 1

    parametric_read = next(iter(fresh_pridb.iread_parametric()))
    assert parametric_read.set_id == 1  # first data set
    # assert parametric_read.status == new_parametric.status
    assert parametric_read.time == new_parametric.time
    assert parametric_read.param_id == new_parametric.param_id
    assert parametric_read.pctd == new_parametric.pctd
    assert parametric_read.pcta == new_parametric.pcta


def test_check_monotonic_time(fresh_pridb):
    def generate_marker(time: float):
        return MarkerRecord(
            time=time,
            set_type=SetType.LABEL,
            number=1,
            data="Test label",
        )

    fresh_pridb.write_marker(generate_marker(1.0))
    fresh_pridb.write_marker(generate_marker(1.0))
    with pytest.raises(ValueError):
        fresh_pridb.write_marker(generate_marker(0.9))


def test_read_empty_pridb(fresh_pridb):
    # dtype conversion might fail because of missing columns
    df = fresh_pridb.read()
    assert len(df) == 0
    assert list(df.columns) == []
