"""Tests for the TykoCallback (HuggingFace Trainer integration)."""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any
from unittest.mock import MagicMock, patch

import pytest
from pytest_httpx import HTTPXMock

from tyko import Run, TykoClient


@dataclass
class MockTrainerState:
    """Minimal mock of transformers.TrainerState."""

    global_step: int = 100
    is_world_process_zero: bool = True
    is_local_process_zero: bool = True
    log_history: list[dict[str, float]] = field(default_factory=list)


@pytest.fixture(autouse=True)
def default_api_key():
    with patch.dict("os.environ", {"TYKO_API_KEY": "test-api-key"}):
        yield


def _create_run(httpx_mock: HTTPXMock) -> Run:
    """Create a Run with mocked HTTP responses.

    Note: start_run() only does a POST. Status PATCHes happen on
    __enter__ or first log() call, so tests that call log() must
    add their own PATCH mock.
    """
    httpx_mock.add_response(
        method="POST",
        url="https://api.tyko-labs.com/runs",
        json={
            "id": "run-cb-test",
            "project": {"name": "test-project"},
            "experiment": {"name": "default"},
            "name": "callback-run",
            "sequential": 1,
        },
    )

    client = TykoClient()
    return client.start_run(project="test-project")


def _get_entries_values(httpx_mock: HTTPXMock) -> list[dict[str, Any]]:
    """Extract metric values from entries POST requests."""
    results = []
    for req in httpx_mock.get_requests():
        if req.method == "POST" and "/entries" in str(req.url):
            body = json.loads(req.content)
            for entry in body["entries"]:
                results.append(entry["values"])
    return results


def test_callback_logs_metrics(httpx_mock: HTTPXMock) -> None:
    """on_log forwards metrics to run.log() with train/ prefix."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)

    # Allow status update + entries POST
    httpx_mock.add_response(method="PATCH", url="https://api.tyko-labs.com/runs/run-cb-test", json={})
    httpx_mock.add_response(method="POST", url="https://api.tyko-labs.com/runs/run-cb-test/entries", json={})

    callback = TykoCallback(run, log_training_args=False)
    callback.on_log(MagicMock(), MockTrainerState(), MagicMock(), logs={"loss": 0.5, "learning_rate": 2e-5})

    values_list = _get_entries_values(httpx_mock)
    assert len(values_list) == 1
    assert values_list[0]["train/loss"] == 0.5
    assert values_list[0]["train/learning_rate"] == 2e-5


def test_callback_rewrites_eval_prefix(httpx_mock: HTTPXMock) -> None:
    """on_log rewrites eval_ prefix to eval/ for dashboard grouping."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)
    httpx_mock.add_response(method="PATCH", url="https://api.tyko-labs.com/runs/run-cb-test", json={})
    httpx_mock.add_response(method="POST", url="https://api.tyko-labs.com/runs/run-cb-test/entries", json={})

    callback = TykoCallback(run, log_training_args=False)
    callback.on_log(MagicMock(), MockTrainerState(), MagicMock(), logs={"eval_loss": 0.3, "eval_accuracy": 0.85})

    values_list = _get_entries_values(httpx_mock)
    assert len(values_list) == 1
    assert "eval/loss" in values_list[0]
    assert "eval/accuracy" in values_list[0]
    assert "eval_loss" not in values_list[0]


def test_callback_skips_total_flos(httpx_mock: HTTPXMock) -> None:
    """on_log filters out total_flos (cumulative counter, not a metric)."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)
    httpx_mock.add_response(method="PATCH", url="https://api.tyko-labs.com/runs/run-cb-test", json={})
    httpx_mock.add_response(method="POST", url="https://api.tyko-labs.com/runs/run-cb-test/entries", json={})

    callback = TykoCallback(run, log_training_args=False)
    callback.on_log(MagicMock(), MockTrainerState(), MagicMock(), logs={"loss": 0.5, "total_flos": 123456789.0})

    values_list = _get_entries_values(httpx_mock)
    assert len(values_list) == 1
    assert "train/loss" in values_list[0]
    assert "total_flos" not in values_list[0]
    assert "train/total_flos" not in values_list[0]


def test_callback_filters_non_numeric(httpx_mock: HTTPXMock) -> None:
    """on_log skips non-numeric values (run.log only accepts float | int)."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)
    httpx_mock.add_response(method="PATCH", url="https://api.tyko-labs.com/runs/run-cb-test", json={})
    httpx_mock.add_response(method="POST", url="https://api.tyko-labs.com/runs/run-cb-test/entries", json={})

    callback = TykoCallback(run, log_training_args=False)
    callback.on_log(MagicMock(), MockTrainerState(), MagicMock(), logs={"loss": 0.5, "some_string": "hello"})

    values_list = _get_entries_values(httpx_mock)
    assert len(values_list) == 1
    assert "train/loss" in values_list[0]
    assert "train/some_string" not in values_list[0]


def test_callback_skips_on_non_zero_rank(httpx_mock: HTTPXMock) -> None:
    """on_log does nothing when not on world process zero (distributed)."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)

    callback = TykoCallback(run, log_training_args=False)
    callback.on_log(
        MagicMock(),
        MockTrainerState(is_world_process_zero=False),
        MagicMock(),
        logs={"loss": 0.5},
    )

    values_list = _get_entries_values(httpx_mock)
    assert len(values_list) == 0


def test_callback_handles_none_logs(httpx_mock: HTTPXMock) -> None:
    """on_log handles logs=None gracefully."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)

    callback = TykoCallback(run, log_training_args=False)
    callback.on_log(MagicMock(), MockTrainerState(), MagicMock(), logs=None)

    values_list = _get_entries_values(httpx_mock)
    assert len(values_list) == 0


def test_callback_logs_training_args(httpx_mock: HTTPXMock) -> None:
    """on_train_begin logs curated TrainingArguments as params."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)
    # PATCH for params update
    httpx_mock.add_response(method="PATCH", url="https://api.tyko-labs.com/runs/run-cb-test", json={})

    callback = TykoCallback(run, log_training_args=True)

    mock_args = MagicMock()
    mock_args.to_dict.return_value = {
        "learning_rate": 2e-5,
        "num_train_epochs": 3.0,
        "per_device_train_batch_size": 16,
        "seed": 42,
        "output_dir": "/tmp/output",
        "logging_dir": "/tmp/logs",
        "dataloader_pin_memory": True,
    }

    callback.on_train_begin(mock_args, MockTrainerState(), MagicMock())

    # Find the params PATCH request
    params_patches = []
    for req in httpx_mock.get_requests():
        if req.method == "PATCH":
            body = json.loads(req.content)
            if "params" in body:
                params_patches.append(body["params"])

    assert len(params_patches) == 1
    params = params_patches[0]
    assert params["learning_rate"] == 2e-5
    assert params["num_train_epochs"] == 3.0
    assert params["seed"] == 42
    # Internal config should be filtered out
    assert "output_dir" not in params
    assert "logging_dir" not in params
    assert "dataloader_pin_memory" not in params


def test_callback_skips_training_args_when_disabled(httpx_mock: HTTPXMock) -> None:
    """on_train_begin does not log params when log_training_args=False."""
    from tyko.integrations import TykoCallback

    run = _create_run(httpx_mock)

    callback = TykoCallback(run, log_training_args=False)
    callback.on_train_begin(MagicMock(), MockTrainerState(), MagicMock())

    # Only the initial POST (create run) and PATCH (status) should exist
    params_patches = []
    for req in httpx_mock.get_requests():
        if req.method == "PATCH":
            body = json.loads(req.content)
            if "params" in body:
                params_patches.append(body)

    assert len(params_patches) == 0


def test_lazy_import_from_tyko() -> None:
    """TykoCallback is accessible via `from tyko import TykoCallback`."""
    from tyko import TykoCallback as TykoCallbackFromRoot
    from tyko.integrations import TykoCallback

    assert TykoCallbackFromRoot is TykoCallback
