from __future__ import annotations

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


class MetricsLoggerTests(unittest.TestCase):
    def test_metrics_are_batched_and_read_back_in_run_order(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("metrics-demo"),
                storage_dir=storage_dir,
            )
            logger = MetricsLogger(session, batch_size=3)

            logger.log(
                "train/loss",
                1.25,
                step=2,
                epoch=0,
                description="Training loss for the current optimization step.",
                metadata={"unit": "cross_entropy"},
            )
            logger.log("optimizer/lr", 0.001, step=1, epoch=0)

            buffered_records = session._store.read_scalar_metrics(session.run_id)
            self.assertEqual(buffered_records, ())

            logger.log("val/accuracy", 0.61, epoch=0)
            records = logger.read()

            self.assertEqual(
                [record.metric_key for record in records],
                ["train/loss", "optimizer/lr", "val/accuracy"],
            )
            self.assertEqual([record.step for record in records], [2, 1, None])
            self.assertEqual(records[0].description, "Training loss for the current optimization step.")
            self.assertEqual(records[0].metadata, {"unit": "cross_entropy"})

            persisted_records = session._store.read_scalar_metrics(session.run_id)
            self.assertEqual([record.metric_key for record in persisted_records], ["train/loss", "optimizer/lr", "val/accuracy"])

    def test_training_loop_metrics_can_be_logged_and_filtered_by_namespace(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("loop-demo"),
                storage_dir=storage_dir,
            )
            logger = MetricsLogger(session, batch_size=10)

            global_step = 0
            for epoch in range(2):
                for loss_value in (1.0 - (epoch * 0.2), 0.8 - (epoch * 0.2)):
                    logger.log("loss", loss_value, step=global_step, epoch=epoch, namespace="train")
                    logger.log(
                        "step_time_ms",
                        12.5 + global_step,
                        step=global_step,
                        epoch=epoch,
                        namespace="system/runtime",
                    )
                    global_step += 1
                logger.log("loss", 0.7 - (epoch * 0.1), epoch=epoch, namespace="val")

            train_loss = logger.read(metric_key="train/loss")
            validation_loss = logger.read(namespace="val", metric_name="loss")
            runtime = logger.read(metric_key="system/runtime/step_time_ms")

            self.assertEqual(len(train_loss), 4)
            self.assertEqual([record.step for record in train_loss], [0, 1, 2, 3])
            self.assertEqual([record.epoch for record in validation_loss], [0, 1])
            self.assertEqual(
                [record.metric_key for record in runtime],
                ["system/runtime/step_time_ms"] * 4,
            )
            self.assertTrue(all(record.namespace == "system/runtime" for record in runtime))

    def test_invalid_metric_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-metrics"),
                storage_dir=storage_dir,
            )
            logger = MetricsLogger(session)

            with self.assertRaises(ValueError):
                logger.log("", 1.0, step=0)
            with self.assertRaises(ValueError):
                logger.log("bad metric", 1.0, step=0)
            with self.assertRaises(ValueError):
                logger.log("train/loss", float("nan"), step=0)
            with self.assertRaises(ValueError):
                logger.log("train/loss", 1.0)


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