from __future__ import annotations

import sys
import tempfile
import unittest
import uuid
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 models import AlertRecord, EvidenceRecord, LayerSummaryRecord, PhaseRecord
from quark import (
    ComparisonEngine,
    ComparisonEngineConfig,
    MetricTargetSpec,
    MetricsLogger,
    TrainingSession,
)
from storage import utc_now
from tests.service_test_utils import unique_project_name


class ComparisonEngineTests(unittest.TestCase):
    def test_comparison_engine_returns_structured_diffs_for_metrics_alerts_phases_and_layers(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            project_name = unique_project_name("compare-demo")
            primary = TrainingSession.start_run(project=project_name, storage_dir=storage_dir)
            reference = TrainingSession.start_run(project=project_name, storage_dir=storage_dir)

            self._write_metrics(
                primary,
                {
                    "train/loss": (1.2, 0.9, 0.6),
                    "val/accuracy": (0.55, 0.72, 0.83),
                    "system/runtime/step_time_ms": (120.0, 110.0, 100.0),
                },
            )
            self._write_metrics(
                reference,
                {
                    "train/loss": (1.4, 1.1, 0.8),
                    "val/accuracy": (0.50, 0.65, 0.76),
                    "system/runtime/step_time_ms": (140.0, 132.0, 128.0),
                },
            )

            self._write_alerts(
                primary,
                [
                    self._alert(
                        run_id=primary.run_id,
                        sequence_id=0,
                        rule_id="exploding_gradients",
                        category="optimization",
                        severity="warning",
                    ),
                    self._alert(
                        run_id=primary.run_id,
                        sequence_id=1,
                        rule_id="runtime_bottleneck_risk",
                        category="runtime",
                        severity="error",
                    ),
                ],
            )
            self._write_alerts(
                reference,
                [
                    self._alert(
                        run_id=reference.run_id,
                        sequence_id=0,
                        rule_id="exploding_gradients",
                        category="optimization",
                        severity="info",
                    )
                ],
            )

            self._write_phases(
                primary,
                [
                    self._phase(
                        run_id=primary.run_id,
                        sequence_id=0,
                        phase_name="rapid_learning",
                        display_name="Rapid Learning",
                        step_start=0,
                        step_end=3,
                    ),
                    self._phase(
                        run_id=primary.run_id,
                        sequence_id=1,
                        phase_name="stabilization",
                        display_name="Stabilization",
                        step_start=4,
                        step_end=6,
                    ),
                ],
            )
            self._write_phases(
                reference,
                [
                    self._phase(
                        run_id=reference.run_id,
                        sequence_id=0,
                        phase_name="rapid_learning",
                        display_name="Rapid Learning",
                        step_start=0,
                        step_end=1,
                    ),
                    self._phase(
                        run_id=reference.run_id,
                        sequence_id=1,
                        phase_name="stabilization",
                        display_name="Stabilization",
                        step_start=2,
                        step_end=6,
                    ),
                ],
            )

            self._write_layer_summaries(
                primary,
                [
                    self._layer_summary(
                        run_id=primary.run_id,
                        sequence_id=0,
                        step=0,
                        module_name="encoder",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=1.0,
                    ),
                    self._layer_summary(
                        run_id=primary.run_id,
                        sequence_id=1,
                        step=1,
                        module_name="encoder",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=0.8,
                    ),
                ],
            )
            self._write_layer_summaries(
                reference,
                [
                    self._layer_summary(
                        run_id=reference.run_id,
                        sequence_id=0,
                        step=0,
                        module_name="encoder",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=1.4,
                    ),
                    self._layer_summary(
                        run_id=reference.run_id,
                        sequence_id=1,
                        step=1,
                        module_name="encoder",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=1.2,
                    ),
                ],
            )

            comparison = ComparisonEngine(primary).compare(reference)
            payload = comparison.to_dict()

            self.assertEqual(comparison.schema_version, 1)
            self.assertEqual(comparison.mode, "pairwise")
            self.assertFalse(comparison.reference.is_baseline)
            self.assertEqual(payload["mode"], "pairwise")
            self.assertIn("metric_diffs", payload)

            train_loss = self._metric_diff(comparison, "train/loss")
            rapid_learning = self._phase_diff(comparison, "rapid_learning")
            encoder_gradient = self._layer_diff(
                comparison,
                module_name="encoder",
                summary_type="gradient",
                hook_direction="backward",
            )

            self.assertAlmostEqual(train_loss.latest_value_diff.delta, -0.2)
            self.assertEqual(comparison.alert_diff.total_count_diff, 1)
            self.assertEqual(comparison.alert_diff.severity_count_diffs["error"], 1)
            self.assertEqual(rapid_learning.total_duration_diff, 2)
            self.assertAlmostEqual(encoder_gradient.norm_diff.delta, -0.4)
            self.assertTrue(
                any(rule.rule_id == "runtime_bottleneck_risk" for rule in comparison.alert_diff.rule_diffs)
            )

    def test_comparison_engine_supports_baseline_workflow_and_metric_targets(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            project_name = unique_project_name("compare-demo")
            candidate = TrainingSession.start_run(project=project_name, storage_dir=storage_dir)
            baseline = TrainingSession.start_run(
                project=project_name,
                storage_dir=storage_dir,
                tags=["baseline"],
            )

            self._write_metrics(
                candidate,
                {
                    "val/accuracy": (0.40, 0.60, 0.82),
                },
            )
            self._write_metrics(
                baseline,
                {
                    "val/accuracy": (0.45, 0.75, 0.78),
                },
            )

            comparison = ComparisonEngine(
                candidate,
                config=ComparisonEngineConfig(
                    metric_targets=(
                        MetricTargetSpec(
                            metric_key="val/accuracy",
                            target_value=0.80,
                            direction="at_least",
                        ),
                    )
                ),
            ).compare_to_baseline(baseline.run_id)

            metric_diff = self._metric_diff(comparison, "val/accuracy")

            self.assertEqual(comparison.mode, "baseline")
            self.assertTrue(comparison.reference.is_baseline)
            self.assertEqual(metric_diff.primary.time_to_target_step, 2)
            self.assertIsNone(metric_diff.reference.time_to_target_step)
            self.assertIsNotNone(metric_diff.primary.target)
            self.assertEqual(metric_diff.primary.target.target_value, 0.8)
            self.assertIsNone(metric_diff.time_to_target_step_diff.delta)
            self.assertEqual(comparison.alert_diff.total_count_diff, 0)
            self.assertEqual(comparison.phase_diffs, ())
            self.assertEqual(comparison.layer_summary_diffs, ())

    @staticmethod
    def _write_metrics(
        session: TrainingSession,
        series_by_key: dict[str, tuple[float, ...]],
    ) -> None:
        logger = MetricsLogger(session, batch_size=64)
        for metric_key, values in series_by_key.items():
            for step, value in enumerate(values):
                logger.log(metric_key, value, step=step, epoch=0)
        logger.flush()

    @staticmethod
    def _write_alerts(
        session: TrainingSession,
        alerts: list[AlertRecord],
    ) -> None:
        session._store.write_analysis_outputs(
            session.run_id,
            alerts=alerts,
            health_scores=(),
            replace_existing=True,
        )

    @staticmethod
    def _write_phases(
        session: TrainingSession,
        phases: list[PhaseRecord],
    ) -> None:
        session._store.write_phases(session.run_id, phases, replace_existing=True)

    @staticmethod
    def _write_layer_summaries(
        session: TrainingSession,
        summaries: list[LayerSummaryRecord],
    ) -> None:
        session._store.write_layer_summaries(session.run_id, summaries)

    @staticmethod
    def _alert(
        *,
        run_id: str,
        sequence_id: int,
        rule_id: str,
        category: str,
        severity: str,
    ) -> AlertRecord:
        return AlertRecord(
            id=uuid.uuid4().hex,
            run_id=run_id,
            sequence_id=sequence_id,
            rule_id=rule_id,
            category=category,
            severity=severity,
            title=rule_id.replace("_", " ").title(),
            summary=f"{rule_id} summary",
            fact=f"{rule_id} fact",
            inference=f"{rule_id} inference",
            timestamp=utc_now(),
            step=sequence_id,
            step_start=sequence_id,
            step_end=sequence_id,
            evidence=(
                EvidenceRecord(
                    kind="metric",
                    source="test-signal",
                    summary="Synthetic comparison evidence.",
                    observed_value=sequence_id,
                    step=sequence_id,
                    step_start=sequence_id,
                    step_end=sequence_id,
                ),
            ),
            metadata={},
        )

    @staticmethod
    def _phase(
        *,
        run_id: str,
        sequence_id: int,
        phase_name: str,
        display_name: str,
        step_start: int,
        step_end: int,
    ) -> PhaseRecord:
        return PhaseRecord(
            id=uuid.uuid4().hex,
            run_id=run_id,
            sequence_id=sequence_id,
            phase_name=phase_name,
            display_name=display_name,
            summary=f"{display_name} summary",
            step_start=step_start,
            step_end=step_end,
            timestamp=utc_now(),
            evidence=(),
            metadata={},
        )

    @staticmethod
    def _layer_summary(
        *,
        run_id: str,
        sequence_id: int,
        step: int,
        module_name: str,
        module_type: str,
        summary_type: str,
        hook_direction: str,
        norm: float,
    ) -> LayerSummaryRecord:
        return LayerSummaryRecord(
            id=uuid.uuid4().hex,
            run_id=run_id,
            sequence_id=sequence_id,
            step=step,
            epoch=0,
            module_name=module_name,
            module_type=module_type,
            summary_type=summary_type,
            hook_direction=hook_direction,
            call_index=0,
            element_count=16,
            timestamp=utc_now(),
            tensor_shape=(4, 4),
            mean=0.0,
            std=0.1,
            norm=norm,
            variance=0.01,
            dead_unit_ratio=0.0,
            saturation_ratio=None,
            metadata={},
        )

    @staticmethod
    def _metric_diff(comparison, metric_key: str):
        return next(item for item in comparison.metric_diffs if item.metric_key == metric_key)

    @staticmethod
    def _phase_diff(comparison, phase_name: str):
        return next(item for item in comparison.phase_diffs if item.phase_name == phase_name)

    @staticmethod
    def _layer_diff(
        comparison,
        *,
        module_name: str,
        summary_type: str,
        hook_direction: str,
    ):
        return next(
            item
            for item in comparison.layer_summary_diffs
            if item.module_name == module_name
            and item.summary_type == summary_type
            and item.hook_direction == hook_direction
        )


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