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 EventLogger, TrainingSession
from tests.service_test_utils import unique_project_name


class EventLoggerTests(unittest.TestCase):
    def test_structured_events_are_queryable_by_type_and_step_range(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("events-demo"),
                storage_dir=storage_dir,
            )
            logger = EventLogger(session, batch_size=4)

            logger.log(
                "checkpoint_saved",
                step=3,
                epoch=0,
                payload={
                    "checkpoint_name": "step-3",
                    "path": "checkpoints/step-3.pt",
                },
            )
            logger.log(
                "scheduler_step",
                step=4,
                epoch=0,
                payload={
                    "scheduler": "cosine",
                    "next_lr": 0.0005,
                },
            )
            session.add_log("Manual inspection after warmup.", identifier="Info")
            logger.add_log(
                "Gradient spike observed.",
                identifier="Warning",
                step=5,
                payload={"author": "researcher"},
            )
            session.end_run(status="completed")

            all_events = logger.read()
            log_events = logger.read(event_type="user_log")
            scheduled_events = logger.read(
                event_types=("checkpoint_saved", "scheduler_step"),
            )
            step_window = logger.read(step_min=4, step_max=5)

            self.assertEqual(
                [event.event_type for event in all_events],
                [
                    "run_started",
                    "checkpoint_saved",
                    "scheduler_step",
                    "user_log",
                    "user_log",
                    "run_ended",
                ],
            )
            self.assertEqual([event.sequence_id for event in all_events], [1, 2, 3, 4, 5, 6])
            self.assertEqual([event.event_type for event in scheduled_events], ["checkpoint_saved", "scheduler_step"])
            self.assertEqual([event.step for event in step_window], [4, 5])
            self.assertEqual(len(log_events), 2)
            self.assertEqual(log_events[0].payload["message"], "Manual inspection after warmup.")
            self.assertEqual(log_events[0].payload["identifier"], "Info")
            self.assertEqual(log_events[1].payload["message"], "Gradient spike observed.")
            self.assertEqual(log_events[1].payload["identifier"], "Warning")
            self.assertEqual(log_events[1].payload["author"], "researcher")
            self.assertEqual(log_events[0].severity, "info")
            self.assertEqual(log_events[1].severity, "warning")
            self.assertEqual(all_events[-1].status, "completed")

            event_history = self._read_jsonl(Path(session.run_path) / "events" / "events.ndjson")
            self.assertEqual(len(event_history), 6)
            ordered_history = sorted(event_history, key=lambda item: item["sequence_id"])
            self.assertEqual(
                [item["event_type"] for item in ordered_history],
                [
                    "run_started",
                    "checkpoint_saved",
                    "scheduler_step",
                    "user_log",
                    "user_log",
                    "run_ended",
                ],
            )
            self.assertEqual(len(session._store.read_events(session.run_id)), 6)

    def test_invalid_event_payloads_are_rejected(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("invalid-events"),
                storage_dir=storage_dir,
            )
            logger = EventLogger(session)

            with self.assertRaises(ValueError):
                logger.log("", step=0)
            with self.assertRaises(ValueError):
                logger.log("Bad Event", step=0)
            with self.assertRaises(TypeError):
                logger.log("checkpoint_saved", step=0, payload={"tags": {"bad"}})
            with self.assertRaises(ValueError):
                logger.add_log("   ")
            with self.assertRaises(ValueError):
                logger.add_log("Loss became nan.", identifier="Critical")
            with self.assertRaises(ValueError):
                logger.add_log("Loss became nan.", payload={"message": "duplicate"})
            with self.assertRaises(ValueError):
                logger.log("anomaly_detected", step=0, severity="urgent", payload={"name": "nan_loss"})

    @staticmethod
    def _read_jsonl(path: Path) -> list[dict[str, object]]:
        return [
            json.loads(line)
            for line in path.read_text(encoding="utf-8").splitlines()
            if line.strip()
        ]


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