from __future__ import annotations

import json
import sys
import tempfile
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SRC = ROOT / "src"

if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

from quark import ARTIFACT_STORAGE_LAYOUT, ArtifactStore, TrainingSession
from tests.service_test_utils import unique_project_name


class ArtifactStoreTests(unittest.TestCase):
    def test_histograms_are_persisted_to_deterministic_paths_and_indexed(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("artifacts-demo"),
                storage_dir=storage_dir,
            )
            store = ArtifactStore(session)

            record = store.log_histogram(
                "activation",
                bins=[-1.0, 0.0, 1.0],
                counts=[3, 7],
                step=12,
                epoch=1,
                namespace="layers/encoder",
                metadata={"module": "encoder.0"},
            )

            discovered = store.read(
                artifact_type="histogram",
                artifact_key="layers/encoder/activation",
            )
            loaded_payload = store.load(record)
            layout_payload = json.loads(
                (Path(session.run_path) / "artifacts" / "layout.json").read_text(encoding="utf-8")
            )

            sequence_label = f"{record.sequence_id:08d}"
            self.assertGreaterEqual(record.sequence_id, 0)
            self.assertEqual(
                record.payload_relative_path,
                (
                    "artifacts/histograms/layers/encoder/activation/"
                    f"step-00000012__epoch-00000001__seq-{sequence_label}.json"
                ),
            )
            self.assertEqual(
                record.metadata_relative_path,
                (
                    "artifacts/histograms/layers/encoder/activation/"
                    f"step-00000012__epoch-00000001__seq-{sequence_label}.meta.json"
                ),
            )
            self.assertTrue(Path(record.payload_path).exists())
            self.assertTrue(Path(record.metadata_path).exists())
            self.assertEqual(loaded_payload, {"bins": [-1.0, 0.0, 1.0], "counts": [3.0, 7.0]})
            self.assertEqual(len(discovered), 1)
            self.assertEqual(discovered[0].artifact_key, "layers/encoder/activation")
            self.assertEqual(discovered[0].metadata, {"module": "encoder.0"})
            self.assertEqual(discovered[0].payload_relative_path, record.payload_relative_path)
            self.assertEqual(layout_payload["type_directories"], ARTIFACT_STORAGE_LAYOUT["type_directories"])

    def test_snapshot_metadata_is_discoverable_and_retention_hooks_are_invoked(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("snapshot-demo"),
                storage_dir=storage_dir,
            )
            retained_artifacts: list[tuple[str, str]] = []

            def capture_retention_hook(store: ArtifactStore, record) -> None:
                retained_artifacts.append((store.run_id, record.artifact_key))

            store = ArtifactStore(session, retention_hooks=[capture_retention_hook])
            record = store.log_snapshot(
                "warmup",
                {
                    "checkpoint_path": "checkpoints/step-40.pt",
                    "checkpoint_size_bytes": 8192,
                },
                snapshot_type="checkpoint_metadata",
                step=40,
                namespace="trainer",
                metadata={
                    "phase": "warmup",
                    "kind": "checkpoint",
                },
            )

            discovered = store.read(
                artifact_type="checkpoint_metadata",
                namespace="trainer",
                step_min=40,
            )
            loaded_payload = store.load(record.id)
            sidecar_payload = json.loads(Path(record.metadata_path).read_text(encoding="utf-8"))

            self.assertEqual(retained_artifacts, [(session.run_id, "trainer/warmup")])
            self.assertEqual(len(discovered), 1)
            self.assertEqual(discovered[0].artifact_name, "warmup")
            self.assertEqual(discovered[0].metadata, {"phase": "warmup", "kind": "checkpoint"})
            self.assertEqual(
                loaded_payload,
                {
                    "checkpoint_path": "checkpoints/step-40.pt",
                    "checkpoint_size_bytes": 8192,
                },
            )
            self.assertEqual(sidecar_payload["artifact_type"], "checkpoint_metadata")
            self.assertEqual(sidecar_payload["metadata"]["phase"], "warmup")

    def test_invalid_artifact_inputs_are_rejected(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("invalid-artifacts"),
                storage_dir=storage_dir,
            )
            store = ArtifactStore(session)

            with self.assertRaises(ValueError):
                store.log_histogram("activation", bins=[0, 1], counts=[1, 2], step=0)
            with self.assertRaises(ValueError):
                store.log_snapshot("bad", {}, snapshot_type="histogram")
            with self.assertRaises(TypeError):
                store.store("histogram", "weights", object(), payload_format="bytes")
            with self.assertRaises(TypeError):
                ArtifactStore(session, retention_hooks=[object()])  # type: ignore[list-item]


if __name__ == "__main__":
    unittest.main(verbosity=2)
