"""Tests for DataStore artifact versioning and lineage tracking."""
import pickle
import shutil
from pathlib import Path

import pytest

from gsql_track.data import DataStore

TEST_DIR = Path("/tmp/gsql_test_datastore")


@pytest.fixture(autouse=True)
def clean_dir():
    if TEST_DIR.exists():
        shutil.rmtree(TEST_DIR)
    TEST_DIR.mkdir(parents=True, exist_ok=True)
    yield
    shutil.rmtree(TEST_DIR, ignore_errors=True)


@pytest.fixture
def store():
    return DataStore(db_path=TEST_DIR / "test_data.db")


def _make_pkl(name, data=None):
    """Create a pickle file with some data."""
    p = TEST_DIR / name
    with open(p, 'wb') as f:
        pickle.dump(data or {"hello": "world"}, f, protocol=4)
    return str(p)


class TestDataStore:

    def test_register_and_get(self, store):
        path = _make_pkl("raw.pkl")
        aid = store.register(path, "imdb_raw")

        # Get by id
        art = store.get(aid)
        assert art['name'] == "imdb_raw"
        assert art['version'] == 1
        assert art['size_bytes'] > 0

        # Get by name:version
        art2 = store.get("imdb_raw:1")
        assert art2['id'] == aid

    def test_auto_version_increment(self, store):
        p1 = _make_pkl("v1.pkl", {"v": 1})
        p2 = _make_pkl("v2.pkl", {"v": 2})
        p3 = _make_pkl("v3.pkl", {"v": 3})

        id1 = store.register(p1, "dataset")
        id2 = store.register(p2, "dataset")
        id3 = store.register(p3, "dataset")

        assert store.get(id1)['version'] == 1
        assert store.get(id2)['version'] == 2
        assert store.get(id3)['version'] == 3
        assert id1 != id2 != id3

    def test_lineage_chain(self, store):
        raw = _make_pkl("raw.pkl", "raw_data")
        tok = _make_pkl("tok.pkl", "tokenized")
        aug = _make_pkl("aug.pkl", "augmented")

        raw_id = store.register(raw, "raw")
        tok_id = store.register(tok, "tokenized", parent=raw_id)
        aug_id = store.register(aug, "augmented", parent=tok_id)

        chain = store.lineage(aug_id)
        assert len(chain) == 3
        assert chain[0]['name'] == "raw"
        assert chain[1]['name'] == "tokenized"
        assert chain[2]['name'] == "augmented"

    def test_children(self, store):
        raw = _make_pkl("raw.pkl")
        child1 = _make_pkl("c1.pkl", {"c": 1})
        child2 = _make_pkl("c2.pkl", {"c": 2})

        raw_id = store.register(raw, "raw")
        store.register(child1, "tok30k", parent=raw_id)
        store.register(child2, "tok50k", parent=raw_id)

        kids = store.children(raw_id)
        assert len(kids) == 2
        names = {k['name'] for k in kids}
        assert names == {"tok30k", "tok50k"}

    def test_versions_list(self, store):
        for i in range(4):
            p = _make_pkl(f"v{i}.pkl", {"i": i})
            store.register(p, "mydata")

        vers = store.versions("mydata")
        assert len(vers) == 4
        # Newest first
        assert vers[0]['version'] == 4
        assert vers[-1]['version'] == 1

    def test_load_artifact(self, store):
        data = {"features": [1, 2, 3], "labels": [0, 1, 0]}
        path = _make_pkl("train.pkl", data)
        aid = store.register(path, "train_data")

        loaded = store.load(aid)
        assert loaded == data

    def test_link_run(self, store):
        path = _make_pkl("data.pkl")
        aid = store.register(path, "dataset")

        store.link_run(aid, "run-abc-123", role='input')
        store.link_run(aid, "run-abc-123", role='output')

        # Verify via direct query
        import sqlite3
        conn = sqlite3.connect(str(TEST_DIR / "test_data.db"))
        rows = conn.execute(
            "SELECT * FROM artifact_runs WHERE artifact_id=?", (aid,)
        ).fetchall()
        conn.close()
        assert len(rows) == 2

    def test_ls_with_tag(self, store):
        p1 = _make_pkl("a.pkl", "a")
        p2 = _make_pkl("b.pkl", "b")
        p3 = _make_pkl("c.pkl", "c")

        id1 = store.register(p1, "data_a", tags=["production"])
        id2 = store.register(p2, "data_b", tags=["staging"])
        id3 = store.register(p3, "data_c")

        # Filter by tag
        prod = store.ls(tag="production")
        assert len(prod) == 1
        assert prod[0]['name'] == "data_a"

        # List all
        all_arts = store.ls()
        assert len(all_arts) == 3

        # Add tag and re-query
        store.tag(id3, "production")
        prod2 = store.ls(tag="production")
        assert len(prod2) == 2

    def test_delete(self, store):
        path = _make_pkl("to_delete.pkl")
        aid = store.register(path, "ephemeral")
        store.link_run(aid, "run-1")

        store.delete(aid)
        with pytest.raises(KeyError):
            store.get(aid)

    def test_latest(self, store):
        for i in range(3):
            p = _make_pkl(f"v{i}.pkl", {"i": i})
            store.register(p, "series")

        lat = store.latest("series")
        assert lat['version'] == 3

    def test_ls_pattern(self, store):
        _make_pkl("a.pkl")
        _make_pkl("b.pkl")
        store.register(_make_pkl("a.pkl"), "imdb_train")
        store.register(_make_pkl("b.pkl"), "imdb_test")
        store.register(_make_pkl("c.pkl", "c"), "yelp_train")

        imdb = store.ls(pattern="imdb_*")
        assert len(imdb) == 2

    def test_lineage_via_name_version(self, store):
        raw = _make_pkl("raw.pkl")
        tok = _make_pkl("tok.pkl")

        store.register(raw, "raw")
        store.register(tok, "tok", parent="raw:1")

        chain = store.lineage("tok:1")
        assert len(chain) == 2
        assert chain[0]['name'] == "raw"

    def test_register_with_steps(self, store):
        path = _make_pkl("processed.pkl")
        steps = [{"type": "tokenize", "params": {"vocab": 30000}}]
        aid = store.register(path, "processed", steps=steps)

        art = store.get(aid)
        assert art['step_configs'] is not None
        import json
        assert json.loads(art['step_configs']) == steps
