"""Tests for upstream job sync and translation."""

from __future__ import annotations

import json
from unittest.mock import AsyncMock

import pytest

from radix_edge.upstream.translate import cloud_job_to_edge_fields
from radix_edge.upstream.job_sync import _poll_and_insert_jobs, _report_completions


class TestJobTranslation:
    def test_basic_translation(self):
        cloud_job = {
            "job_id": "job-cloud123",
            "cluster_id": "cluster-abc",
            "tenant_id": "tenant-1",
            "job_kind": "training",
            "image": "myimg:latest",
            "command": ["python", "train.py"],
            "args": ["--epochs", "10"],
            "env": {"LR": "0.01"},
            "orchestration": {"num_gpus": 2, "timeout_seconds": 7200},
            "params": {"model": "resnet"},
        }
        fields = cloud_job_to_edge_fields(cloud_job)

        assert fields["user_id"] == "tenant-1"
        assert fields["image"] == "myimg:latest"
        assert fields["command"] == ["python", "train.py"]
        assert fields["args"] == ["--epochs", "10"]
        assert fields["env"] == {"LR": "0.01"}
        assert fields["num_gpus"] == 2
        assert fields["timeout_seconds"] == 7200
        assert fields["origin"] == "upstream"
        assert fields["upstream_job_id"] == "job-cloud123"
        assert fields["job_type"] == "training"

    def test_defaults_when_missing(self):
        cloud_job = {"job_id": "job-min", "image": "img:1"}
        fields = cloud_job_to_edge_fields(cloud_job)

        assert fields["num_gpus"] == 1
        assert fields["timeout_seconds"] == 3600
        assert fields["priority"] == 5
        assert fields["origin"] == "upstream"
        assert fields["user_id"] == "upstream"
        assert fields["command"] == []
        assert fields["args"] == []
        assert fields["env"] == {}

    def test_benchmark_job(self):
        cloud_job = {
            "job_id": "job-bench",
            "job_kind": "resnet50_benchmark",
            "image": "iangreen74/radix-resnet50-benchmark:latest",
            "command": ["python", "train_resnet50.py", "--epochs", "1"],
            "orchestration": {"num_gpus": 1},
            "params": {"epochs": 1, "batch_size": 128},
            "tenant_id": "demo-tenant",
        }
        fields = cloud_job_to_edge_fields(cloud_job)
        assert fields["job_type"] == "resnet50_benchmark"
        assert fields["metadata"]["params"]["epochs"] == 1

    def test_name_format(self):
        cloud_job = {"job_id": "job-x", "job_kind": "inference", "image": "img:1"}
        fields = cloud_job_to_edge_fields(cloud_job)
        assert fields["name"] == "upstream:inference:job-x"


class TestPollAndInsertJobs:
    @pytest.mark.asyncio
    async def test_inserts_polled_job(self, db):
        mock_client = AsyncMock()
        mock_client.poll_jobs.return_value = [{
            "job_id": "job-upstream1",
            "job_kind": "training",
            "image": "img:1",
            "tenant_id": "t1",
            "command": ["python", "train.py"],
            "orchestration": {"num_gpus": 1},
        }]

        count = await _poll_and_insert_jobs(mock_client)
        assert count == 1

        row = db.execute(
            "SELECT * FROM jobs WHERE upstream_job_id = 'job-upstream1'"
        ).fetchone()
        assert row is not None
        assert row["origin"] == "upstream"
        assert row["status"] == "queued"
        assert row["image"] == "img:1"
        assert row["user_id"] == "t1"

    @pytest.mark.asyncio
    async def test_deduplication(self, db):
        mock_client = AsyncMock()
        cloud_job = {
            "job_id": "job-dup1",
            "job_kind": "generic",
            "image": "img:1",
            "tenant_id": "t1",
        }
        mock_client.poll_jobs.return_value = [cloud_job]

        count1 = await _poll_and_insert_jobs(mock_client)
        assert count1 == 1

        # Second poll with same job — should be deduped
        mock_client.poll_jobs.return_value = [cloud_job]
        count2 = await _poll_and_insert_jobs(mock_client)
        assert count2 == 0

    @pytest.mark.asyncio
    async def test_empty_poll(self, db):
        mock_client = AsyncMock()
        mock_client.poll_jobs.return_value = []

        count = await _poll_and_insert_jobs(mock_client)
        assert count == 0

    @pytest.mark.asyncio
    async def test_poll_error_returns_zero(self, db):
        mock_client = AsyncMock()
        mock_client.poll_jobs.side_effect = Exception("Connection refused")

        count = await _poll_and_insert_jobs(mock_client)
        assert count == 0


class TestReportCompletions:
    @pytest.mark.asyncio
    async def test_reports_completed_job(self, db):
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, origin, upstream_job_id,
                exit_code, runtime_seconds, completed_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
            ("job-local1", "t1", "img:1", "succeeded", "upstream", "job-upstream1",
             0, 42.0, "2026-02-27T00:00:00Z"),
        )
        db.commit()

        mock_client = AsyncMock()
        mock_client.complete_job.return_value = {"status": "ok"}

        count = await _report_completions(mock_client)
        assert count == 1
        mock_client.complete_job.assert_called_once()

        call_kwargs = mock_client.complete_job.call_args.kwargs
        assert call_kwargs["job_id"] == "job-upstream1"
        assert call_kwargs["status"] == "succeeded"

    @pytest.mark.asyncio
    async def test_marks_as_reported(self, db):
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, origin, upstream_job_id,
                exit_code, completed_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
            ("job-local2", "t1", "img:1", "failed", "upstream", "job-upstream2",
             1, "2026-02-27T00:00:00Z"),
        )
        db.commit()

        mock_client = AsyncMock()
        mock_client.complete_job.return_value = {"status": "ok"}

        await _report_completions(mock_client)

        # Should be marked as reported
        row = db.execute("SELECT metadata FROM jobs WHERE job_id = 'job-local2'").fetchone()
        meta = json.loads(row["metadata"])
        assert meta["reported"] is True

        # Second call should not report again
        mock_client.reset_mock()
        count = await _report_completions(mock_client)
        assert count == 0
        mock_client.complete_job.assert_not_called()

    @pytest.mark.asyncio
    async def test_skips_local_jobs(self, db):
        db.execute(
            """INSERT INTO jobs (job_id, user_id, image, status, origin, exit_code, completed_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)""",
            ("job-local3", "t1", "img:1", "succeeded", "local", 0, "2026-02-27T00:00:00Z"),
        )
        db.commit()

        mock_client = AsyncMock()
        count = await _report_completions(mock_client)
        assert count == 0
        mock_client.complete_job.assert_not_called()

    @pytest.mark.asyncio
    async def test_completion_error_continues(self, db):
        for i in range(2):
            db.execute(
                """INSERT INTO jobs (job_id, user_id, image, status, origin, upstream_job_id,
                    exit_code, completed_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (f"job-local{i}", "t1", "img:1", "succeeded", "upstream", f"job-up{i}",
                 0, "2026-02-27T00:00:00Z"),
            )
        db.commit()

        mock_client = AsyncMock()
        # First call fails, second succeeds
        mock_client.complete_job.side_effect = [Exception("Network error"), {"status": "ok"}]

        count = await _report_completions(mock_client)
        assert count == 1  # Only the second succeeded


class TestGracefulDegradation:
    def test_upstream_disabled_by_default(self):
        from radix_edge.config import get_config
        config = get_config()
        assert config.upstream_enabled is False
        assert config.api_base == ""
        assert config.cluster_id == ""
