from __future__ import annotations

import io
import json
from types import SimpleNamespace

import numpy as np
import pytest

from dovla_cil.data.images import CILImageReader
from dovla_cil.generation.pipeline import generate_cil_dataset
from dovla_cil.tasks.library import built_in_toy_tasks
from dovla_cil.training.trainer import DoVLATrainer, TrainerConfig


def test_image_reader_resolves_collection_source_and_caches_archive(tmp_path) -> None:
    h5py = pytest.importorskip("h5py")
    image_module = pytest.importorskip("PIL.Image")
    source = tmp_path / "source"
    source.mkdir()
    expected = np.full((6, 8, 3), 117, dtype=np.uint8)
    buffer = io.BytesIO()
    image_module.fromarray(expected).save(buffer, format="JPEG", quality=95)
    encoded = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
    with h5py.File(source / "observations.h5", "w") as handle:
        dataset = handle.create_dataset(
            "initial_rgb_jpeg",
            shape=(1,),
            dtype=h5py.vlen_dtype(np.dtype("uint8")),
        )
        dataset[0] = encoded
    record = SimpleNamespace(
        record_id="record-0",
        observation_ref="observations.h5#initial_rgb_jpeg/0",
        next_observation_ref=None,
        metadata={"source_dataset": str(source)},
    )

    with CILImageReader(tmp_path / "collection") as reader:
        actual = reader.read(record)
        cached = reader.read(record)
        assert len(reader._handles) == 1

    assert actual.shape == expected.shape
    assert actual.dtype == np.uint8
    assert np.array_equal(actual, cached)


def test_rgb_trainer_smoke_writes_checkpoint(tmp_path) -> None:
    h5py = pytest.importorskip("h5py")
    image_module = pytest.importorskip("PIL.Image")
    torch = pytest.importorskip("torch")
    dataset_dir = tmp_path / "data"
    generate_cil_dataset(
        backend="toy",
        tasks=built_in_toy_tasks()[:2],
        out_dir=dataset_dir,
        num_states_per_task=1,
        k=2,
        seed=9,
        shard_size=8,
        inline_observations=True,
    )
    image = np.zeros((24, 24, 3), dtype=np.uint8)
    image[..., 0] = 180
    buffer = io.BytesIO()
    image_module.fromarray(image).save(buffer, format="JPEG", quality=90)
    encoded = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
    with h5py.File(dataset_dir / "observations.h5", "w") as handle:
        images = handle.create_dataset(
            "initial_rgb_jpeg",
            shape=(1,),
            dtype=h5py.vlen_dtype(np.dtype("uint8")),
        )
        images[0] = encoded
    for shard in (dataset_dir / "shards").glob("*.jsonl"):
        records = [json.loads(line) for line in shard.read_text().splitlines()]
        for record in records:
            record["observation_ref"] = "observations.h5#initial_rgb_jpeg/0"
        shard.write_text("".join(json.dumps(record) + "\n" for record in records))

    run_dir = tmp_path / "run"
    result = DoVLATrainer(
        TrainerConfig(
            dataset_dir=dataset_dir,
            output_dir=run_dir,
            epochs=1,
            batch_groups=1,
            records_per_group=2,
            pair_count_per_group=1,
            hidden_dim=24,
            action_horizon=2,
            effect_dim=8,
            observation_mode="rgb",
            device="cpu",
        )
    ).train()

    checkpoint = torch.load(run_dir / "best.pt", map_location="cpu", weights_only=False)
    assert result["best"]
    assert checkpoint["model_config"]["observation_mode"] == "rgb"