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 quark import AnalyzerEngine, AnalyzerEngineConfig, AnalyzerThresholds, MetricsLogger, TrainingSession
from models import LayerSummaryRecord
from storage import utc_now
from tests.service_test_utils import unique_project_name


class AnalyzerEngineTests(unittest.TestCase):
    def test_analyzer_emits_persisted_alerts_with_evidence_and_health_scores(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("analyzer-demo"),
                storage_dir=storage_dir,
            )
            metrics = MetricsLogger(session, batch_size=20)

            for step, value in enumerate((0.9, 0.7, 0.5)):
                metrics.log("train/loss", value, step=step, epoch=0)
            for step, value in enumerate((0.95, 1.0, 1.15)):
                metrics.log("val/loss", value, step=step, epoch=0)
            for step, value in enumerate((620.0, 700.0, 760.0)):
                metrics.log("system/runtime/step_time_ms", value, step=step, epoch=0)
            for step, value in enumerate((410.0, 450.0, 520.0)):
                metrics.log("system/runtime/data_time_ms", value, step=step, epoch=0)
            metrics.flush()

            self._write_layer_summaries(
                session,
                [
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=0,
                        step=0,
                        module_name="encoder",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=60.0,
                    ),
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=1,
                        step=0,
                        module_name="stem",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=0.00005,
                    ),
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=2,
                        step=0,
                        module_name="body",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=0.00004,
                    ),
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=3,
                        step=1,
                        module_name="head",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=0.00003,
                    ),
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=4,
                        step=2,
                        module_name="head",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=0.00002,
                    ),
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=5,
                        step=2,
                        module_name="relu",
                        module_type="ReLU",
                        summary_type="activation",
                        hook_direction="forward",
                        norm=0.01,
                        std=0.0005,
                        dead_unit_ratio=0.99,
                    ),
                ],
            )

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

            self.assertEqual(
                [alert.rule_id for alert in result.alerts],
                [
                    "exploding_gradients",
                    "vanishing_gradients",
                    "stagnant_layers",
                    "activation_collapse",
                    "overfitting_onset",
                    "runtime_bottleneck_risk",
                ],
            )
            self.assertTrue(all(alert.evidence for alert in result.alerts))
            self.assertEqual(
                [score.score_key for score in result.health_scores],
                ["optimization", "representation", "generalization", "runtime", "overall"],
            )

            optimization_score = analyzer.read_health_scores(score_key="optimization")[0]
            overall_score = analyzer.read_health_scores(score_key="overall")[0]
            runtime_alert = analyzer.read_alerts(rule_id="runtime_bottleneck_risk")[0]

            self.assertLess(optimization_score.score, 50.0)
            self.assertLess(overall_score.score, 100.0)
            self.assertEqual(runtime_alert.category, "runtime")
            self.assertEqual(runtime_alert.severity, "error")
            self.assertTrue(
                any(evidence.kind == "metric" for evidence in runtime_alert.evidence)
            )

    def test_rule_thresholds_are_configurable_and_previous_outputs_are_replaced(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("threshold-demo"),
                storage_dir=storage_dir,
            )
            self._write_layer_summaries(
                session,
                [
                    self._layer_summary(
                        run_id=session.run_id,
                        sequence_id=0,
                        step=0,
                        module_name="encoder",
                        module_type="Linear",
                        summary_type="gradient",
                        hook_direction="backward",
                        norm=15.0,
                    )
                ],
            )

            default_analyzer = AnalyzerEngine(session)
            default_result = default_analyzer.analyze()
            self.assertEqual([alert.rule_id for alert in default_result.alerts], ["exploding_gradients"])

            strict_analyzer = AnalyzerEngine(
                session,
                config=AnalyzerEngineConfig(
                    thresholds=AnalyzerThresholds(
                        exploding_gradient_warning_norm=20.0,
                        exploding_gradient_error_norm=30.0,
                        exploding_gradient_critical_norm=40.0,
                    )
                ),
            )
            strict_result = strict_analyzer.analyze()

            self.assertEqual(strict_result.alerts, ())
            self.assertEqual(strict_analyzer.read_alerts(), ())
            optimization_score = strict_analyzer.read_health_scores(score_key="optimization")[0]
            self.assertEqual(optimization_score.score, 100.0)

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

    @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,
        std: float | None = 0.1,
        dead_unit_ratio: float | None = None,
    ) -> 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=std,
            norm=norm,
            variance=0.0,
            dead_unit_ratio=dead_unit_ratio,
            saturation_ratio=None,
            metadata={},
        )


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