"""Tests for adaptive expand selector, data planner, and revision chain."""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
import torch.nn as nn

from model_osmosis import OsmosisConfig, OsmosisPipeline
from model_osmosis.optimizer.expand_engine import (
    finalize_expand_profile,
    select_expand_profile,
)
from model_osmosis.services.expand_data import infer_expand_task, plan_expand_data
from model_osmosis.services.data_registry import DataSourceProfile
from model_osmosis.services.data_resolver import resolve_data_source
from model_osmosis.services.expand_revisions import (
    active_revision,
    append_revision,
    list_revisions,
    rollback_revision,
)


@pytest.fixture(autouse=True)
def _local_state_homes(monkeypatch, tmp_path):
    monkeypatch.setenv("OSMOSIS_OPTIMIZER_HOME", str(tmp_path / "optimizer"))
    monkeypatch.setenv("OSMOSIS_EXPAND_HOME", str(tmp_path / "expand"))
    monkeypatch.setenv("OSMOSIS_EXPAND_DATA_HOME", str(tmp_path / "expand_data"))


def test_expand_selector_is_deterministic():
    context = {
        "task": "classification",
        "baseline_params": 128,
        "target_params": 256,
        "budget_minutes": 5,
        "budget_ram_gb": 1.5,
        "model_type": "Sequential",
    }
    a = select_expand_profile(context=context, optimizer_mode="auto", seed=42, record_selection=False)
    b = select_expand_profile(context=context, optimizer_mode="auto", seed=42, record_selection=False)
    assert a.profile_id == b.profile_id
    assert a.selector_version == b.selector_version
    assert a.pack_version == b.pack_version
    assert a.score == pytest.approx(b.score)


def test_expand_selector_quality_rollback():
    context = {
        "task": "text",
        "baseline_params": 1024,
        "target_params": 2048,
        "budget_minutes": 5,
        "budget_ram_gb": 1.5,
        "model_type": "Transformer",
    }
    selected = select_expand_profile(context=context, optimizer_mode="auto", seed=7, record_selection=False)
    finalized = finalize_expand_profile(
        selected,
        quality_ok=False,
        fallback_reason="quality_regression",
        quality_score=0.1,
    )
    assert finalized.canary_status == "rollback"
    assert finalized.rollback_reason == "quality_regression"


def test_expand_data_planner_synthetic():
    data, report = plan_expand_data(task="classification", dataset=None, seed=123, feature_dim=6)
    assert isinstance(data, np.ndarray)
    assert data.shape == (256, 6)
    assert report["dataset_source"] == "synthetic"
    assert report["sample_shape"] == [256, 6]
    assert report["license_policy"] == "any"
    assert report["task_batch_kind"] == "tabular"


def test_expand_data_planner_task_aware_text_synthetic():
    data, report = plan_expand_data(task="text", dataset=None, seed=321, feature_dim=6)
    assert isinstance(data, dict)
    assert data["kind"] == "text"
    input_ids = np.asarray(data["input_ids"])
    assert input_ids.ndim == 2
    assert input_ids.shape[0] > 0
    assert report["task_batch_kind"] == "text"
    assert report["task_batch_type"] == "dict"
    assert len(report["task_batch_shape"]) == 2


def test_expand_data_planner_task_aware_vision_synthetic():
    data, report = plan_expand_data(task="vision", dataset=None, seed=99, feature_dim=4)
    assert isinstance(data, dict)
    assert data["kind"] == "vision"
    images = np.asarray(data["images"])
    assert images.ndim == 4
    assert images.shape[1] == 3
    assert report["task_batch_kind"] == "vision"
    assert report["task_batch_type"] == "dict"


def test_expand_data_planner_user_path(tmp_path):
    dataset_path = tmp_path / "sample.npy"
    np.save(dataset_path, np.ones((10, 4), dtype=np.float32))
    data, report = plan_expand_data(
        task="classification",
        dataset=str(dataset_path),
        seed=1,
        feature_dim=4,
    )
    assert data.shape == (10, 4)
    assert report["dataset_source"] == "user_path"
    assert report["cache_hit"] is True


def test_data_resolver_is_deterministic():
    registry = [
        DataSourceProfile(
            dataset_id="candidate_a",
            task="classification",
            uri="file:///tmp/a.npy",
            hub="local",
            revision="v1",
            license="mit",
            priority=10,
            min_quality=0.2,
        ),
        DataSourceProfile(
            dataset_id="candidate_b",
            task="classification",
            uri="file:///tmp/b.npy",
            hub="local",
            revision="v1",
            license="mit",
            priority=20,
            min_quality=0.4,
        ),
    ]
    r1 = resolve_data_source(
        task="classification",
        model_ref="pt:model.pt",
        seed=42,
        registry=registry,
        allowed_licenses=["mit"],
    )
    r2 = resolve_data_source(
        task="classification",
        model_ref="pt:model.pt",
        seed=42,
        registry=registry,
        allowed_licenses=["mit"],
    )
    assert r1.selected is not None
    assert r2.selected is not None
    assert r1.selected.dataset_id == r2.selected.dataset_id
    assert r1.selection_reason == "best_score"


def test_expand_data_planner_auto_registry_fetch(tmp_path):
    dataset_path = tmp_path / "auto.npy"
    np.save(dataset_path, np.ones((12, 3), dtype=np.float32))
    registry = [
        DataSourceProfile(
            dataset_id="local_auto",
            task="classification",
            uri=dataset_path.resolve().as_uri(),
            hub="local",
            revision="v1",
            license="mit",
            priority=1,
            min_quality=0.9,
        )
    ]
    data, report = plan_expand_data(
        task="classification",
        dataset=None,
        seed=1,
        feature_dim=3,
        model_ref="pt:model.pt",
        data_mode="auto",
        registry=registry,
        allow_licenses=["mit"],
    )
    assert data.shape == (12, 3)
    assert report["dataset_source"] == "auto_hub"
    assert report["dataset_id"] == "local_auto"
    assert report["source_hub"] == "local"

    _, report2 = plan_expand_data(
        task="classification",
        dataset=None,
        seed=1,
        feature_dim=3,
        model_ref="pt:model.pt",
        data_mode="auto",
        registry=registry,
        allow_licenses=["mit"],
        offline=True,
    )
    assert report2["cache_hit"] is True


def test_infer_expand_task():
    class DummyModel:
        pass

    task = infer_expand_task(task="auto", model_ref="hf:HuggingFaceTB/SmolLM2-135M-Instruct", model_obj=DummyModel())
    assert task == "text"


def test_expand_revision_chain_and_rollback(tmp_path):
    first = append_revision(
        name="demo_model",
        artifact_path=str(Path("models") / "demo_model.r0001.osmoimg"),
        accepted=True,
        baseline_metric=1.0,
        post_expand_metric=1.2,
        quality_delta=0.2,
        report={"status": "accepted"},
    )
    second = append_revision(
        name="demo_model",
        artifact_path=str(Path("models") / "demo_model.r0002.osmoimg"),
        accepted=False,
        baseline_metric=1.2,
        post_expand_metric=1.0,
        quality_delta=-0.2,
        report={"status": "rejected"},
    )
    assert first["revision_id"].endswith("r0001")
    assert second["revision_id"].endswith("r0002")
    assert active_revision(name="demo_model")["revision_id"] == first["revision_id"]

    rolled = rollback_revision(second["revision_id"])
    assert rolled["revision_id"] == second["revision_id"]
    assert active_revision(name="demo_model")["revision_id"] == second["revision_id"]
    assert len(list_revisions(name="demo_model")) == 2


def test_expand_strict_mode_fails_without_structural_backend():
    model = nn.Linear(6, 4)
    pipeline = OsmosisPipeline(OsmosisConfig(verbose=False, seed=42))
    sample = np.random.default_rng(1).standard_normal((32, 6)).astype("float32")

    _, report = pipeline.expand(
        model,
        sample_data=sample,
        task="classification",
        target_params="x1.2",
        budget_minutes=1,
        seed=42,
        fallback_mode="off",
        strict_mode=True,
        route="native_torch",
        runtime_source="native",
    )
    assert report["strict_mode"] is True
    assert report["structural_applied"] is False
    assert report["accepted"] is False
    assert report["hard_failure"] is True


def test_expand_distill_fallback_marks_degraded():
    model = nn.Linear(6, 4)
    pipeline = OsmosisPipeline(OsmosisConfig(verbose=False, seed=42))
    sample = np.random.default_rng(2).standard_normal((16, 6)).astype("float32")

    _, report = pipeline.expand(
        model,
        sample_data=sample,
        task="classification",
        target_params="x1.2",
        budget_minutes=1,
        seed=42,
        fallback_mode="distill",
        strict_mode=True,
        route="native_torch",
        runtime_source="native",
    )
    assert report["fallback_mode"] == "distill"
    assert report["degraded"] is True
    assert "self_check_status" in report


def test_expand_report_core_fields_are_deterministic():
    model = nn.Sequential(nn.Linear(6, 8), nn.Linear(8, 4))
    pipeline = OsmosisPipeline(OsmosisConfig(verbose=False, seed=42))
    sample = np.random.default_rng(3).standard_normal((24, 6)).astype("float32")

    _, report_a = pipeline.expand(
        model,
        sample_data=sample,
        task="classification",
        target_params="x1.2",
        budget_minutes=1,
        seed=42,
        fallback_mode="off",
        strict_mode=True,
        route="native_torch",
        runtime_source="native",
    )
    _, report_b = pipeline.expand(
        model,
        sample_data=sample,
        task="classification",
        target_params="x1.2",
        budget_minutes=1,
        seed=42,
        fallback_mode="off",
        strict_mode=True,
        route="native_torch",
        runtime_source="native",
    )
    keys = [
        "task",
        "route",
        "runtime_source",
        "strict_mode",
        "fallback_mode",
        "profile_id",
        "pack_version",
        "selector_version",
        "status",
        "degraded",
    ]
    for key in keys:
        assert report_a.get(key) == report_b.get(key)
