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,
    EventLogger,
    EventRecord,
    MetricRecord,
    MetricsLogger,
    ParameterDistributionRecord,
    ParameterSummaryRecord,
    TrainingSession,
    dump_record,
    validate_record,
)
from schemas import SchemaValidationError
from storage import LocalSessionStore
from tests.service_test_utils import unique_project_name


class SchemaCompatibilityTests(unittest.TestCase):
    def test_unnamespaced_metric_records_match_metrics_logger_contract(self) -> None:
        record = MetricRecord(
            id="metric-1",
            run_id="run-1",
            sequence_id=0,
            namespace="",
            metric_name="loss",
            metric_key="loss",
            value=0.25,
            step=1,
            epoch=None,
            timestamp="2026-05-07T00:00:00Z",
        )

        validate_record(record)
        self.assertEqual(dump_record(record)["namespace"], "")

    def test_parameter_summary_validation_matches_api_contract(self) -> None:
        record = ParameterSummaryRecord(
            id="parameter-summary-1",
            run_id="run-1",
            sequence_id=0,
            step=1,
            epoch=None,
            timestamp="2026-05-07T00:00:00Z",
            module_name="encoder.linear",
            module_type="Linear",
            parameter_name="weight",
            parameter_key="encoder.linear.weight",
            parameter_role="weight",
            element_count=8,
            tensor_shape=(2, 4),
            mean=0.1,
            std=0.2,
            zero_ratio=0.0,
            finite_ratio=1.0,
            update_to_parameter_norm_ratio=0.25,
            metadata={"source": "unit-test"},
        )

        validate_record(record)
        self.assertEqual(dump_record(record)["tensor_shape"], [2, 4])

        invalid_ratio = ParameterSummaryRecord(
            **{
                **record.to_dict(),
                "tensor_shape": record.tensor_shape,
                "update_to_parameter_norm_ratio": 1.01,
            }
        )
        with self.assertRaises(SchemaValidationError):
            validate_record(invalid_ratio)

        raw_tensor_metadata = ParameterSummaryRecord(
            **{
                **record.to_dict(),
                "tensor_shape": record.tensor_shape,
                "metadata": {"nested": {"tensor": [1.0, 2.0]}},
            }
        )
        with self.assertRaises(SchemaValidationError):
            validate_record(raw_tensor_metadata)

    def test_parameter_distribution_validation_matches_api_contract(self) -> None:
        record = ParameterDistributionRecord(
            id="parameter-distribution-1",
            run_id="run-1",
            sequence_id=0,
            step=1,
            epoch=None,
            timestamp="2026-05-07T00:00:00Z",
            module_name="encoder.linear",
            module_type="Linear",
            parameter_name="weight",
            parameter_key="encoder.linear.weight",
            parameter_role="weight",
            distribution_type="value",
            bin_edges=(-1.0, 0.0, 1.0),
            counts=(3, 5),
            sample_count=8,
            total_count=8,
            metadata={"source": "unit-test"},
        )

        validate_record(record)
        self.assertEqual(dump_record(record)["bin_edges"], [-1.0, 0.0, 1.0])

        bad_edges = ParameterDistributionRecord(
            **{
                **record.to_dict(),
                "bin_edges": (-1.0, 0.0, 0.0),
                "counts": record.counts,
            }
        )
        with self.assertRaises(SchemaValidationError):
            validate_record(bad_edges)

        bad_counts = ParameterDistributionRecord(
            **{
                **record.to_dict(),
                "bin_edges": record.bin_edges,
                "counts": (3, -1),
            }
        )
        with self.assertRaises(SchemaValidationError):
            validate_record(bad_counts)

    def test_invalid_persisted_event_records_are_rejected_before_write(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("schema-validation"),
                storage_dir=storage_dir,
            )
            invalid_event = EventRecord(
                id="event-1",
                run_id=session.run_id,
                sequence_id=1,
                event_type="checkpoint_saved",
                severity="urgent",  # type: ignore[arg-type]
                timestamp="2026-04-14T00:00:00Z",
                step=0,
                payload={"path": "checkpoints/step-0.pt"},
            )

            with self.assertRaises(SchemaValidationError):
                session._store.write_events(session.run_id, (invalid_event,))

    def test_current_store_writes_schema_versions_to_service_records_and_snapshots(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("schema-versioning"),
                config={"epochs": 2},
                storage_dir=storage_dir,
            )
            metrics = MetricsLogger(session, batch_size=8)
            events = EventLogger(session, batch_size=8)

            metrics.log("train/loss", 0.42, step=1, epoch=0)
            metrics.flush()
            events.log("checkpoint_saved", step=1, payload={"path": "checkpoints/step-1.pt"})
            events.flush()

            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"))
            event_history = self._read_jsonl(run_path / "events" / "events.ndjson")

            self.assertEqual(run_snapshot["project"]["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(run_snapshot["run"]["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(config_payload["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertTrue(all(item["schema_version"] == CURRENT_SCHEMA_VERSION for item in event_history))

            reloaded = session.reload()
            persisted_metrics = metrics.read(metric_key="train/loss")
            persisted_events = events.read(event_type="checkpoint_saved")

            self.assertEqual(reloaded.project.schema_version, CURRENT_SCHEMA_VERSION)
            self.assertEqual(reloaded.run.schema_version, CURRENT_SCHEMA_VERSION)
            self.assertEqual(persisted_metrics[0].schema_version, CURRENT_SCHEMA_VERSION)
            self.assertEqual(persisted_events[0].schema_version, CURRENT_SCHEMA_VERSION)

    def test_load_run_round_trips_current_schema_records_from_the_service(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("schema-load-run"),
                config={"epochs": 3, "optimizer": "adamw"},
                storage_dir=storage_dir,
            )
            session.add_log("Recovered from live service state.", identifier="Info")

            store = LocalSessionStore(storage_dir)
            context = store.load_run(session.run_id)

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

            self.assertEqual(context.project.schema_version, CURRENT_SCHEMA_VERSION)
            self.assertEqual(context.run.schema_version, CURRENT_SCHEMA_VERSION)
            self.assertEqual(context.status, "running")
            self.assertEqual(run_snapshot["run"]["status"], "active")
            self.assertEqual(config_payload["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(environment_payload["schema_version"], CURRENT_SCHEMA_VERSION)
            self.assertEqual(environment_payload["run_id"], session.run_id)
            events = session._store.read_events(session.run_id, event_type="user_log")
            self.assertEqual(tuple(event.payload["message"] for event in events), ("Recovered from live service state.",))

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