"""Tests for log_table and log_series on the Run class."""

from unittest.mock import MagicMock

import pandas as pd
import polars as pl
import pytest

from tyko._types import Experiment, Project, Run

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_run() -> Run:
    """Create a Run with a mocked client (no HTTP needed)."""
    mock_client = MagicMock()
    mock_client.create_artifact_version.return_value = "v-abc"
    return Run(
        id="run-123",
        name="test-run",
        sequential=1,
        experiment=Experiment(name="default"),
        project=Project(name="test-project"),
        _client=mock_client,
        _client_uuid="test-uuid",
        _silent=True,
    )


# ---------------------------------------------------------------------------
# log_table — polars DataFrame
# ---------------------------------------------------------------------------


class TestLogTablePolars:
    def test_uploads_polars_dataframe(self) -> None:
        run = _make_run()
        df = pl.DataFrame({"loss": [0.5, 0.3, 0.1], "accuracy": [0.7, 0.8, 0.9]})

        version_id = run.log_table("metrics", df)

        assert version_id == "v-abc"
        run._client.create_artifact_version.assert_called_once_with(
            "run-123",
            "metrics",
            "table",
            filter_data=None,
        )
        run._client.upload_artifact_file.assert_called_once()
        args = run._client.upload_artifact_file.call_args[0]
        assert args[0] == "v-abc"
        assert args[1] == "data.parquet"
        assert isinstance(args[2], bytes)
        assert len(args[2]) > 0
        run._client.complete_artifact_version.assert_called_once_with("v-abc")

    def test_polars_parquet_is_readable(self) -> None:
        """Verify the uploaded bytes are valid Parquet."""
        run = _make_run()
        df = pl.DataFrame({"x": [1, 2, 3], "y": [4, 5, 6]})

        run.log_table("data", df)

        parquet_bytes = run._client.upload_artifact_file.call_args[0][2]
        # Read back and verify
        import io

        result = pl.read_parquet(io.BytesIO(parquet_bytes))
        assert result.shape == (3, 2)
        assert result["x"].to_list() == [1, 2, 3]


# ---------------------------------------------------------------------------
# log_table — pandas DataFrame
# ---------------------------------------------------------------------------


class TestLogTablePandas:
    def test_uploads_pandas_dataframe(self) -> None:
        run = _make_run()
        df = pd.DataFrame({"loss": [0.5, 0.3], "acc": [0.8, 0.9]})

        version_id = run.log_table("metrics", df)

        assert version_id == "v-abc"
        run._client.upload_artifact_file.assert_called_once()
        args = run._client.upload_artifact_file.call_args[0]
        assert args[1] == "data.parquet"
        assert len(args[2]) > 0

    def test_pandas_parquet_is_readable(self) -> None:
        """Verify pandas-produced Parquet is readable by polars."""
        run = _make_run()
        df = pd.DataFrame({"a": [10, 20], "b": [30, 40]})

        run.log_table("data", df)

        parquet_bytes = run._client.upload_artifact_file.call_args[0][2]
        import io

        result = pl.read_parquet(io.BytesIO(parquet_bytes))
        assert result.shape == (2, 2)
        assert result["a"].to_list() == [10, 20]


# ---------------------------------------------------------------------------
# log_table — dict of lists
# ---------------------------------------------------------------------------


class TestLogTableDictOfLists:
    def test_uploads_dict_of_lists(self) -> None:
        run = _make_run()
        data = {"y": [0.5, 0.3], "pred": [0.6, 0.4]}

        version_id = run.log_table("predictions", data)

        assert version_id == "v-abc"
        run._client.upload_artifact_file.assert_called_once()

    def test_dict_of_lists_is_readable(self) -> None:
        run = _make_run()
        data = {"col_a": [1, 2, 3], "col_b": ["x", "y", "z"]}

        run.log_table("mixed", data)

        parquet_bytes = run._client.upload_artifact_file.call_args[0][2]
        import io

        result = pl.read_parquet(io.BytesIO(parquet_bytes))
        assert result.shape == (3, 2)
        assert result["col_a"].to_list() == [1, 2, 3]
        assert result["col_b"].to_list() == ["x", "y", "z"]


# ---------------------------------------------------------------------------
# log_table — list of dicts
# ---------------------------------------------------------------------------


class TestLogTableListOfDicts:
    def test_uploads_list_of_dicts(self) -> None:
        run = _make_run()
        data = [{"y": 0.5, "pred": 0.6}, {"y": 0.3, "pred": 0.4}]

        version_id = run.log_table("predictions", data)

        assert version_id == "v-abc"
        run._client.upload_artifact_file.assert_called_once()

    def test_list_of_dicts_is_readable(self) -> None:
        run = _make_run()
        data = [{"epoch": 1, "loss": 0.5}, {"epoch": 2, "loss": 0.3}]

        run.log_table("training", data)

        parquet_bytes = run._client.upload_artifact_file.call_args[0][2]
        import io

        result = pl.read_parquet(io.BytesIO(parquet_bytes))
        assert result.shape == (2, 2)
        assert result["epoch"].to_list() == [1, 2]


# ---------------------------------------------------------------------------
# log_table — error cases
# ---------------------------------------------------------------------------


class TestLogTableErrors:
    def test_rejects_string(self) -> None:
        run = _make_run()
        with pytest.raises(TypeError, match="Expected a DataFrame"):
            run.log_table("bad", "not a table")

    def test_rejects_plain_list(self) -> None:
        run = _make_run()
        with pytest.raises(TypeError, match="Expected a DataFrame"):
            run.log_table("bad", [1, 2, 3])

    def test_rejects_empty_list(self) -> None:
        run = _make_run()
        with pytest.raises(TypeError, match="Expected a DataFrame"):
            run.log_table("bad", [])

    def test_fails_version_on_upload_error(self) -> None:
        run = _make_run()
        run._client.upload_artifact_file.side_effect = RuntimeError("upload failed")

        with pytest.raises(RuntimeError, match="upload failed"):
            run.log_table("data", {"a": [1]})

        run._client.fail_artifact_version.assert_called_once_with("v-abc")
        run._client.complete_artifact_version.assert_not_called()

    def test_metadata_passed_as_filter_data(self) -> None:
        run = _make_run()
        run.log_table("data", {"a": [1]}, metadata={"epoch": 5})

        run._client.create_artifact_version.assert_called_once_with(
            "run-123",
            "data",
            "table",
            filter_data={"epoch": 5},
        )


# ---------------------------------------------------------------------------
# log_series
# ---------------------------------------------------------------------------


class TestLogSeries:
    def test_polars_series(self) -> None:
        run = _make_run()
        series = pl.Series("loss", [0.5, 0.3, 0.1])

        version_id = run.log_series("training_loss", series)

        assert version_id == "v-abc"
        run._client.upload_artifact_file.assert_called_once()

    def test_pandas_series(self) -> None:
        run = _make_run()
        series = pd.Series([0.5, 0.3, 0.1], name="loss")

        version_id = run.log_series("training_loss", series)

        assert version_id == "v-abc"

    def test_rejects_non_series(self) -> None:
        run = _make_run()
        with pytest.raises(TypeError, match="Expected a Series-like object"):
            run.log_series("bad", [1, 2, 3])

    def test_series_parquet_is_readable(self) -> None:
        run = _make_run()
        series = pl.Series("accuracy", [0.7, 0.8, 0.9])

        run.log_series("acc", series)

        parquet_bytes = run._client.upload_artifact_file.call_args[0][2]
        import io

        result = pl.read_parquet(io.BytesIO(parquet_bytes))
        assert result.shape == (3, 1)
        assert result["accuracy"].to_list() == [0.7, 0.8, 0.9]
