from __future__ import annotations

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 (
    BatchMetadata,
    BatchMetadataAdapter,
    DataDiagnostics,
    DataDiagnosticsConfig,
    MappingBatchMetadataAdapter,
    TrainingSession,
)
from tests.service_test_utils import unique_project_name

class _TupleBatchAdapter(BatchMetadataAdapter):
    def extract_batch_metadata(self, batch) -> BatchMetadata:
        sample_ids, class_counts, sequence_lengths, source = batch
        return BatchMetadata(
            sample_ids=sample_ids,
            class_counts=class_counts,
            sequence_lengths=sequence_lengths,
            metadata={"source": source},
        )


class DataDiagnosticsTests(unittest.TestCase):
    def test_batch_diagnostics_are_logged_with_mapping_adapter_and_hard_batches_are_surfaced(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("data-diagnostics-demo"),
                storage_dir=storage_dir,
            )
            diagnostics = DataDiagnostics(
                session,
                config=DataDiagnosticsConfig(
                    batch_size=3,
                    adapter=MappingBatchMetadataAdapter(
                        sample_id_key="sample_ids",
                        label_key="labels",
                        sequence_length_key="lengths",
                        metadata_keys=("loader",),
                    ),
                ),
            )

            batches = (
                {
                    "sample_ids": ("sample-00a", "sample-00b"),
                    "labels": ("cat", "cat"),
                    "lengths": (24, 28),
                    "loader": "train",
                },
                {
                    "sample_ids": ("sample-01a", "sample-01b"),
                    "labels": ("cat", "dog"),
                    "lengths": (26, 29),
                    "loader": "train",
                },
                {
                    "sample_ids": ("sample-02a", "sample-02b"),
                    "labels": ("dog", "dog"),
                    "lengths": (25, 31),
                    "loader": "train",
                },
                {
                    "sample_ids": ("sample-03a", "sample-03b"),
                    "labels": ("dog", "rare"),
                    "lengths": (27, 33),
                    "loader": "train",
                },
                {
                    "sample_ids": ("sample-04a", "sample-04b"),
                    "labels": ("cat", "rare"),
                    "lengths": (22, 35),
                    "loader": "train",
                },
                {
                    "sample_ids": ("sample-05a", "sample-05b"),
                    "labels": ("rare", "rare"),
                    "lengths": (480, 512),
                    "loader": "train",
                },
            )
            losses = (0.92, 0.89, 0.94, 1.02, 1.08, 4.8)

            for step, (loss, batch) in enumerate(zip(losses, batches, strict=True)):
                diagnostics.log_batch(loss, step=step, epoch=0, batch=batch)

            records = diagnostics.read()
            sample_trace = diagnostics.read(sample_id="sample-05a")
            exposure_summary = diagnostics.summarize_exposure()
            hard_batches = diagnostics.find_hard_batches(limit=3)
            summary = diagnostics.summarize(hard_batch_limit=2)

            self.assertEqual(len(records), 6)
            self.assertEqual([record.step for record in records], [0, 1, 2, 3, 4, 5])
            self.assertEqual(records[-1].sample_ids, ("sample-05a", "sample-05b"))
            self.assertEqual(records[-1].class_counts, {"rare": 2})
            self.assertEqual(records[-1].metadata, {"loader": "train"})
            self.assertEqual(len(sample_trace), 1)
            self.assertEqual(sample_trace[0].step, 5)

            self.assertEqual(exposure_summary.total_batches, 6)
            self.assertEqual(exposure_summary.total_examples, 12)
            self.assertEqual(exposure_summary.unique_sample_id_count, 12)
            self.assertEqual(
                {item.label: item.count for item in exposure_summary.class_exposures},
                {"cat": 4, "dog": 4, "rare": 4},
            )

            self.assertEqual(hard_batches[0].batch.step, 5)
            self.assertIn("loss_outlier", hard_batches[0].reasons)
            self.assertIn("high_loss", hard_batches[0].reasons)
            self.assertIn("long_sequence_batch", hard_batches[0].reasons)
            self.assertEqual(summary.hard_batch_count, 1)
            self.assertEqual(summary.hard_batches[0].batch.step, 5)
            self.assertEqual(summary.sequence_length_summary.max_value, 512)
            self.assertEqual(len(session._store.read_batch_diagnostics(session.run_id)), 6)

    def test_custom_batch_metadata_adapters_can_drive_logging_without_mapping_batches(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("adapter-demo"),
                storage_dir=storage_dir,
            )
            diagnostics = DataDiagnostics(
                session,
                config=DataDiagnosticsConfig(
                    batch_size=10,
                    adapter=_TupleBatchAdapter(),
                    min_batches_for_outlier_detection=2,
                ),
            )

            diagnostics.log_batch(
                1.0,
                step=0,
                batch=(
                    ("a-0", "a-1"),
                    {"cat": 2},
                    (10, 12),
                    "sampler-a",
                ),
            )
            diagnostics.log_batch(
                1.3,
                step=1,
                batch=(
                    ("b-0",),
                    {"dog": 1},
                    (18,),
                    "sampler-b",
                ),
            )

            records = diagnostics.read()
            summary = diagnostics.summarize(hard_batch_limit=None)

            self.assertEqual([record.batch_size for record in records], [2, 1])
            self.assertEqual(records[0].metadata, {"source": "sampler-a"})
            self.assertEqual(records[1].class_counts, {"dog": 1})
            self.assertEqual(summary.exposure_summary.total_examples, 3)
            self.assertEqual(
                {item.label: item.count for item in summary.exposure_summary.class_exposures},
                {"cat": 2, "dog": 1},
            )
            self.assertEqual(summary.loss_summary.count, 2)

    def test_invalid_batch_diagnostics_payloads_are_rejected(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            storage_dir = Path(temp_dir)
            session = TrainingSession.start_run(
                project=unique_project_name("invalid-data-diagnostics"),
                storage_dir=storage_dir,
            )
            diagnostics = DataDiagnostics(session)

            with self.assertRaises(ValueError):
                diagnostics.log_batch(float("nan"), step=0)
            with self.assertRaises(ValueError):
                diagnostics.log_batch(1.0)
            with self.assertRaises(ValueError):
                diagnostics.log_batch(1.0, step=0, batch={"labels": ("cat",)})
            with self.assertRaises(ValueError):
                diagnostics.log_batch(1.0, step=0, sample_ids=("ok", ""), batch_size=2)
            with self.assertRaises(ValueError):
                DataDiagnosticsConfig(hard_batch_percentile=1.0)


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