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


class PhaseAnalyzerTests(unittest.TestCase):
    def test_phase_analyzer_segments_and_persists_major_training_regimes(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = self._build_sample_run(storage_dir)

            analyzer = PhaseAnalyzer(session)
            result = analyzer.analyze()

            self.assertEqual(
                [phase.phase_name for phase in result.phases],
                [
                    "initialization_adaptation",
                    "rapid_learning",
                    "stabilization",
                    "stagnation",
                    "overfitting_onset",
                    "divergence",
                ],
            )
            self.assertEqual(
                [(phase.step_start, phase.step_end) for phase in result.phases],
                [(0, 1), (2, 4), (5, 7), (8, 9), (10, 11), (12, 13)],
            )
            self.assertTrue(all(phase.evidence for phase in result.phases))

            overfitting_phase = analyzer.read_phases(step=10)[0]
            middle_phases = analyzer.read_phases(step_min=5, step_max=9)
            phases_payload = json.loads(
                (Path(session.run_path) / "analysis" / "phases.json").read_text(encoding="utf-8")
            )

            self.assertEqual(overfitting_phase.phase_name, "overfitting_onset")
            self.assertEqual(
                [phase.phase_name for phase in middle_phases],
                ["stabilization", "stagnation"],
            )
            self.assertEqual(
                [item["phase_name"] for item in phases_payload["phases"]],
                [phase.phase_name for phase in result.phases],
            )
            self.assertEqual(phases_payload["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertTrue(
                all(item["schema_version"] == CURRENT_SCHEMA_VERSION for item in phases_payload["phases"])
            )
            self.assertEqual(
                phases_payload["phases"][4]["metadata"]["entered_from_phase"],
                "stagnation",
            )

    def test_phase_segmentation_is_deterministic_for_identical_runs(self) -> None:
        with tempfile.TemporaryDirectory() as first_dir, tempfile.TemporaryDirectory() as second_dir:
            first_session = self._build_sample_run(Path(first_dir))
            second_session = self._build_sample_run(Path(second_dir))

            first_result = PhaseAnalyzer(first_session).analyze()
            second_result = PhaseAnalyzer(second_session).analyze()

            self.assertEqual(
                self._normalize_phases(first_result.phases),
                self._normalize_phases(second_result.phases),
            )

    @staticmethod
    def _build_sample_run(storage_dir: Path) -> TrainingSession:
        session = TrainingSession.start_run(
            project=unique_project_name("phase-demo"),
            storage_dir=storage_dir,
        )
        metrics = MetricsLogger(session, batch_size=32)

        train_loss = (
            1.50,
            1.42,
            1.10,
            0.78,
            0.58,
            0.50,
            0.45,
            0.42,
            0.418,
            0.416,
            0.404,
            0.392,
            0.52,
            0.71,
        )
        val_loss = (
            1.55,
            1.46,
            1.18,
            0.90,
            0.72,
            0.64,
            0.60,
            0.58,
            0.579,
            0.58,
            0.60,
            0.63,
            0.82,
            1.08,
        )

        for step, value in enumerate(train_loss):
            metrics.log("train/loss", value, step=step, epoch=0)
        for step, value in enumerate(val_loss):
            metrics.log("val/loss", value, step=step, epoch=0)
        metrics.flush()
        return session

    @staticmethod
    def _normalize_phases(phases) -> list[dict[str, object]]:
        normalized_payload: list[dict[str, object]] = []
        for phase in phases:
            payload = phase.to_dict()
            payload.pop("id", None)
            payload.pop("run_id", None)
            payload.pop("timestamp", None)
            normalized_payload.append(payload)
        return normalized_payload


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