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


class SessionManagerTests(unittest.TestCase):
    def test_start_run_persists_metadata_and_lifecycle_files(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)

            session = start_run(
                project=unique_project_name("image-classifier"),
                config={"epochs": 3, "optimizer": "adamw"},
                tags=["baseline", "resnet18"],
                storage_dir=storage_dir,
            )

            run_path = Path(session.run_path)
            run_snapshot = json.loads((run_path / "meta" / "run.json").read_text(encoding="utf-8"))
            config_payload = json.loads((run_path / "meta" / "config.json").read_text(encoding="utf-8"))
            environment_payload = json.loads(
                (run_path / "meta" / "environment.json").read_text(encoding="utf-8")
            )
            lifecycle_lines = self._read_jsonl(run_path / "events" / "events.ndjson")

            self.assertEqual(session.status, "running")
            self.assertEqual(config_payload["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(config_payload["run_id"], session.run_id)
            self.assertEqual(config_payload["config"]["epochs"], 3)
            self.assertEqual(environment_payload["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(environment_payload["run_id"], session.run_id)
            self.assertEqual(run_snapshot["run"]["status"], "active")
            self.assertEqual(run_snapshot["project"]["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(run_snapshot["run"]["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(tuple(run_snapshot["run"]["tags"]), ("baseline", "resnet18"))
            self.assertEqual(lifecycle_lines[0]["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(lifecycle_lines[0]["event_type"], "run_started")
            self.assertEqual(lifecycle_lines[0]["status"], "active")
            self.assertIn("python", environment_payload["environment"])
            self.assertTrue(Path(session.run_path).exists())
            self.assertEqual(Path(session.run_path).parent.parent.parent, storage_dir / ".quark")
            self.assertEqual(run_snapshot["run"]["parent_run_id"], None)
            self.assertEqual(run_snapshot["run"]["resume_count"], 0)
            self.assertTrue(session.context.db_path)

    def test_parent_child_lineage_is_persisted_for_new_runs(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            project_name = unique_project_name("lineage-demo")
            parent = TrainingSession.start_run(project=project_name, storage_dir=storage_dir)
            parent.end_run(status="completed")

            child = TrainingSession.start_run(
                project=project_name,
                parent_run_id=parent.run_id,
                storage_dir=storage_dir,
            )

            child_snapshot = json.loads(
                (Path(child.run_path) / "meta" / "run.json").read_text(encoding="utf-8")
            )

            self.assertEqual(child_snapshot["run"]["parent_run_id"], parent.run_id)
            self.assertEqual(child_snapshot["run"]["lineage_root_run_id"], parent.run_id)
            self.assertEqual(child.context.run.parent_run_id, parent.run_id)

    def test_resume_run_updates_resume_count_and_final_metadata(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            project_name = unique_project_name("resume-demo")
            first_session = TrainingSession.start_run(
                project=project_name,
                config={"seed": 7},
                tags=["initial"],
                storage_dir=storage_dir,
            )
            first_session.add_log("Checkpoint saved before shutdown.", identifier="Info")
            first_session.end_run(status="interrupted")

            resumed = TrainingSession.resume_run(run_id=first_session.run_id, storage_dir=storage_dir)
            resumed.add_tag("resumed")
            resumed.add_log("Recovered training from local metadata.", identifier="Info")
            resumed.end_run(status="completed")

            run_path = Path(resumed.run_path)
            run_snapshot = json.loads((run_path / "meta" / "run.json").read_text(encoding="utf-8"))
            event_lines = self._read_jsonl(run_path / "events" / "events.ndjson")
            lifecycle_lines = [item for item in event_lines if item["event_type"] != "user_log"]
            log_lines = [item for item in event_lines if item["event_type"] == "user_log"]
            lifecycle_types = [item["event_type"] for item in lifecycle_lines]
            lifecycle_statuses = [item["status"] for item in lifecycle_lines]

            self.assertEqual(resumed.status, "completed")
            self.assertEqual(run_snapshot["run"]["resume_count"], 1)
            self.assertEqual(tuple(run_snapshot["run"]["tags"]), ("initial", "resumed"))
            self.assertIsNotNone(run_snapshot["run"]["latest_resume_at"])
            self.assertIsNotNone(run_snapshot["run"]["ended_at"])
            self.assertEqual(len(log_lines), 2)
            self.assertEqual(
                [item["payload"]["message"] for item in log_lines],
                [
                    "Checkpoint saved before shutdown.",
                    "Recovered training from local metadata.",
                ],
            )
            self.assertEqual([item["payload"]["identifier"] for item in log_lines], ["Info", "Info"])
            self.assertEqual(
                lifecycle_types,
                ["run_started", "run_ended", "run_resumed", "run_ended"],
            )
            self.assertEqual(
                lifecycle_statuses,
                ["active", "interrupted", "active", "completed"],
            )

    def test_context_manager_logs_runtime_errors_before_marking_run_failed(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session: TrainingSession | None = None

            with self.assertRaisesRegex(RuntimeError, "loss exploded"):
                with TrainingSession.start_run(
                    project=unique_project_name("runtime-error-demo"),
                    storage_dir=storage_dir,
                ) as active_session:
                    session = active_session
                    raise RuntimeError("loss exploded")

            self.assertIsNotNone(session)
            assert session is not None
            event_lines = self._read_jsonl(Path(session.run_path) / "events" / "events.ndjson")
            runtime_logs = [item for item in event_lines if item["event_type"] == "runtime_error"]

            self.assertEqual(session.status, "failed")
            self.assertEqual(len(runtime_logs), 1)
            self.assertEqual(runtime_logs[0]["severity"], "error")
            self.assertEqual(runtime_logs[0]["status"], "failed")
            self.assertEqual(runtime_logs[0]["payload"]["identifier"], "Error")
            self.assertIn("loss exploded", runtime_logs[0]["payload"]["message"])
            self.assertEqual(runtime_logs[0]["payload"]["exception_type"], "RuntimeError")

    def test_context_manager_logs_keyboard_interrupts_before_marking_run_interrupted(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session: TrainingSession | None = None

            with self.assertRaises(KeyboardInterrupt):
                with TrainingSession.start_run(
                    project=unique_project_name("runtime-interruption-demo"),
                    storage_dir=storage_dir,
                ) as active_session:
                    session = active_session
                    raise KeyboardInterrupt("manual stop")

            self.assertIsNotNone(session)
            assert session is not None
            event_lines = self._read_jsonl(Path(session.run_path) / "events" / "events.ndjson")
            interruption_logs = [item for item in event_lines if item["event_type"] == "runtime_interruption"]

            self.assertEqual(session.status, "interrupted")
            self.assertEqual(len(interruption_logs), 1)
            self.assertEqual(interruption_logs[0]["severity"], "warning")
            self.assertEqual(interruption_logs[0]["status"], "interrupted")
            self.assertEqual(interruption_logs[0]["payload"]["identifier"], "Warning")
            self.assertIn("manual stop", interruption_logs[0]["payload"]["message"])
            self.assertEqual(interruption_logs[0]["payload"]["exception_type"], "KeyboardInterrupt")

    @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)
