from __future__ import annotations

import base64
import hashlib

from crypto_spot_bot.training_coordination import TrainingCoordinator


def test_training_coordinator_claims_and_completes_job(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)

    requested = coordinator.request_retrain({"source": "android"})
    job_id = requested["job"]["id"]
    heartbeat = coordinator.heartbeat({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"})
    claimed = coordinator.claim({"worker_id": "win-1", "name": "DESKTOP-TMFDL0H"})

    assert requested["queued"] is True
    assert heartbeat["status"]["agent_online"] is True
    assert claimed["claimed"] is True
    assert claimed["job"]["id"] == job_id
    assert coordinator.status()["active_job"]["status"] == "running"

    completed = coordinator.complete(job_id, {"success": True, "message": "ok"})

    assert completed["job"]["status"] == "completed"
    assert coordinator.status()["active_job"] is None


def test_training_coordinator_accepts_chunked_artifact_upload(tmp_path) -> None:
    coordinator = TrainingCoordinator(tmp_path)
    job = coordinator.request_retrain({"source": "test"})["job"]
    payload = b'{"type":"pytorch_recurrent_forecaster","symbols":{}}\n'
    sha256 = hashlib.sha256(payload).hexdigest()
    first = payload[:20]
    second = payload[20:]

    part_1 = coordinator.save_artifact_chunk(
        job["id"],
        {
            "name": "lstm_forecaster.json",
            "index": 0,
            "total": 2,
            "sha256": sha256,
            "data_base64": base64.b64encode(first).decode("ascii"),
        },
    )
    part_2 = coordinator.save_artifact_chunk(
        job["id"],
        {
            "name": "lstm_forecaster.json",
            "index": 1,
            "total": 2,
            "sha256": sha256,
            "data_base64": base64.b64encode(second).decode("ascii"),
        },
    )

    assert part_1["complete"] is False
    assert part_2["complete"] is True
    assert (tmp_path / "lstm_forecaster.json").read_bytes() == payload
    assert coordinator.status()["latest_job"]["artifacts"][0]["sha256"] == sha256