from __future__ import annotations

from typing import Any

import expops_sklearn.runner as sklearn_runner
import pytest


class _LossCurveModel:
    def __init__(self) -> None:
        self.loss_curve_ = [1.0, 0.7, 0.4]


class _LossCurveAdapter:
    def __init__(self, **params: Any) -> None:
        self.params = dict(params)
        self._fitted_model = None

    def fit(self, X: Any, y: Any = None) -> "_LossCurveAdapter":  # noqa: ARG002
        self._fitted_model = _LossCurveModel()
        return self

    def predict(self, X: Any) -> list[int]:
        return [0 for _ in range(len(X))]


class _PartialFitModel:
    _estimator_type = "classifier"

    def __init__(self, **params: Any) -> None:
        self.params = dict(params)
        self.loss_ = 1.0
        self.partial_fit_calls = 0

    def partial_fit(self, X: Any, y: Any, classes: Any = None) -> "_PartialFitModel":  # noqa: ARG002
        self.partial_fit_calls += 1
        self.loss_ = max(0.0, self.loss_ - 0.2)
        return self

    def predict(self, X: Any) -> list[int]:
        return [0 for _ in range(len(X))]


class _PartialFitAdapter:
    _sklearn_class = _PartialFitModel

    def __init__(self, **params: Any) -> None:
        self.params = dict(params)
        self._fitted_model = None

    def fit(self, X: Any, y: Any = None) -> "_PartialFitAdapter":  # noqa: ARG002
        raise AssertionError("fit() should not be used when iterative_mode=partial_fit")

    def predict(self, X: Any) -> list[int]:
        return self._fitted_model.predict(X)


def test_run_estimator_component_partial_fit_iterative_logging(monkeypatch) -> None:
    metric_calls: list[tuple[str, float, int | None]] = []

    def _capture_metric(name: str, value: Any, step: int | None = None) -> None:
        metric_calls.append((name, float(value), step))

    monkeypatch.setattr(sklearn_runner, "log_metric", _capture_metric)

    result = sklearn_runner._run_estimator_component(
        adapter_cls=_PartialFitAdapter,
        available_inputs={"X": [[0.0], [1.0], [2.0], [3.0]], "y": [0, 0, 1, 1]},
        input_aliases={},
        params={
            "operation": "fit",
            "max_iter": 2,
        },
        iterative_cfg={
            "enabled": True,
            "metric": "train_loss",
            "eval_split": "train",
        },
    )

    model = result.get("model")
    assert isinstance(model, _PartialFitModel)
    assert model.partial_fit_calls == 2
    assert [name for name, _, _ in metric_calls] == ["train_loss", "train_loss"]
    assert [step for _, _, step in metric_calls] == [1, 2]
    assert [value for _, value, _ in metric_calls] == pytest.approx([0.8, 0.6])


def test_run_estimator_component_iterative_logging_requires_partial_fit(monkeypatch) -> None:
    metric_calls: list[tuple[str, float, int | None]] = []

    def _capture_metric(name: str, value: Any, step: int | None = None) -> None:  # noqa: ARG001
        metric_calls.append((name, float(value), step))

    monkeypatch.setattr(sklearn_runner, "log_metric", _capture_metric)

    with pytest.raises(ValueError, match="partial_fit"):
        sklearn_runner._run_estimator_component(
            adapter_cls=_LossCurveAdapter,
            available_inputs={"X": [[0.0], [1.0]], "y": [0, 1]},
            input_aliases={},
            params={
                "operation": "fit",
                "max_iter": 3,
            },
            iterative_cfg={
                "enabled": True,
                "metric": "train_loss",
                "eval_split": "train",
            },
        )


def test_run_sklearn_component_uses_operation_suffix(monkeypatch) -> None:
    from expops.core.runtime.step_system import StepContext, set_current_process_context

    class _PredictOnlyModel:
        def __init__(self) -> None:
            self.called_with: Any = None

        def predict(self, X: Any) -> list[int]:
            self.called_with = X
            return [1 for _ in range(len(X))]

    def _get_adapter(name: str):
        assert name == "MockEstimator"

        class _Adapter:
            def __init__(self, **params: Any) -> None:  # noqa: ARG002
                self.params = dict(params)

        return _Adapter

    monkeypatch.setattr(sklearn_runner, "get_adapter", _get_adapter)

    ctx = StepContext(
        project_id="p",
        run_id="r",
        global_config={
            "pipeline": {
                "processes": [
                    {
                        "name": "predict_component",
                        "component": "sklearn.MockEstimator.predict",
                        "input_transform": {"model": "model", "X": "X"},
                    }
                ]
            }
        },
    )

    model = _PredictOnlyModel()
    set_current_process_context("predict_component")
    try:
        result = sklearn_runner.run_sklearn_component(context=ctx, model=model, X=[[0.0], [1.0]])
    finally:
        set_current_process_context(None)

    assert "predictions" in result
    assert result["predictions"] == [1, 1]


def test_run_sklearn_component_rejects_parameters_operation(monkeypatch) -> None:
    from expops.core.runtime.step_system import StepContext, set_current_process_context

    def _get_adapter(name: str):
        class _Adapter:
            def __init__(self, **params: Any) -> None:  # noqa: ARG002
                self.params = dict(params)

            def fit(self, X: Any, y: Any = None) -> "_Adapter":  # noqa: ARG002
                self._fitted_model = self
                return self

        return _Adapter

    monkeypatch.setattr(sklearn_runner, "get_adapter", _get_adapter)

    ctx = StepContext(
        project_id="p",
        run_id="r",
        global_config={
            "pipeline": {
                "processes": [
                    {
                        "name": "component_train",
                        "component": "sklearn.LogisticRegression.fit",
                        "input_transform": {"X": "X", "y": "y"},
                        "parameters": {"operation": "fit"},
                    }
                ]
            }
        },
    )

    set_current_process_context("component_train")
    try:
        with pytest.raises(ValueError, match="parameters.operation"):
            sklearn_runner.run_sklearn_component(context=ctx, X=[[0.0], [1.0]], y=[0, 1])
    finally:
        set_current_process_context(None)
 
