import importlib.util
import sys
import types
import asyncio
from pathlib import Path

import torch


def _register_real_transfer_codec(package_name, utils_pkg):
    """Load the real utils/transfer_codec.py into the synthetic test package."""
    codec_path = Path(__file__).resolve().parents[1] / "utils" / "transfer_codec.py"
    codec_spec = importlib.util.spec_from_file_location(f"{package_name}.utils.transfer_codec", codec_path)
    codec_module = importlib.util.module_from_spec(codec_spec)
    sys.modules[f"{package_name}.utils.transfer_codec"] = codec_module
    codec_spec.loader.exec_module(codec_module)
    utils_pkg.transfer_codec = codec_module
    return codec_module


def _load_collector_module():
    module_path = Path(__file__).resolve().parents[1] / "nodes" / "collector.py"
    package_name = "dist_collector_list_testpkg"

    for mod_name in list(sys.modules):
        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
            del sys.modules[mod_name]

    root_pkg = types.ModuleType(package_name)
    root_pkg.__path__ = []
    sys.modules[package_name] = root_pkg

    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
    nodes_pkg.__path__ = []
    sys.modules[f"{package_name}.nodes"] = nodes_pkg

    utils_pkg = types.ModuleType(f"{package_name}.utils")
    utils_pkg.__path__ = []
    sys.modules[f"{package_name}.utils"] = utils_pkg
    _register_real_transfer_codec(package_name, utils_pkg)

    class _Routes:
        def post(self, _path):
            return lambda fn: fn

        def get(self, _path):
            return lambda fn: fn

    prompt_server = types.SimpleNamespace(
        routes=_Routes(),
        distributed_jobs_lock=None,
        distributed_pending_jobs={},
    )
    server_module = types.ModuleType("server")
    server_module.PromptServer = types.SimpleNamespace(instance=prompt_server)
    sys.modules["server"] = server_module

    comfy_module = types.ModuleType("comfy")
    model_management = types.ModuleType("comfy.model_management")

    class InterruptProcessingException(Exception):
        pass

    model_management.InterruptProcessingException = InterruptProcessingException
    model_management.throw_exception_if_processing_interrupted = lambda: None
    comfy_module.model_management = model_management

    comfy_utils = types.ModuleType("comfy.utils")

    class ProgressBar:
        def __init__(self, _total):
            self.total = _total
            self.updates = []

        def update(self, value):
            self.updates.append(value)

    comfy_utils.ProgressBar = ProgressBar
    comfy_module.utils = comfy_utils
    sys.modules["comfy"] = comfy_module
    sys.modules["comfy.model_management"] = model_management
    sys.modules["comfy.utils"] = comfy_utils

    aiohttp_module = types.ModuleType("aiohttp")
    aiohttp_module.ClientTimeout = lambda total: types.SimpleNamespace(total=total)
    sys.modules["aiohttp"] = aiohttp_module

    logging_module = types.ModuleType(f"{package_name}.utils.logging")
    logging_module.debug_log = lambda *_args, **_kwargs: None
    logging_module.log = lambda *_args, **_kwargs: None
    sys.modules[f"{package_name}.utils.logging"] = logging_module

    config_module = types.ModuleType(f"{package_name}.utils.config")
    config_module.get_worker_timeout_seconds = lambda: 0.1
    config_module.load_config = lambda: {"workers": []}
    config_module.is_master_delegate_only = lambda: False
    sys.modules[f"{package_name}.utils.config"] = config_module

    constants_module = types.ModuleType(f"{package_name}.utils.constants")
    constants_module.HEARTBEAT_INTERVAL = 1.0
    constants_module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024 * 1024
    constants_module.TRANSFER_MAX_REQUEST_BYTES = 96 * 1024 * 1024
    sys.modules[f"{package_name}.utils.constants"] = constants_module

    image_module = types.ModuleType(f"{package_name}.utils.image")
    def _ensure_contiguous(tensor):
        return tensor.contiguous() if hasattr(tensor, "contiguous") else tensor

    image_module.ensure_contiguous = _ensure_contiguous
    image_module.tensor_to_pil = lambda *_args, **_kwargs: None
    image_module.pil_to_tensor = lambda value: value
    sys.modules[f"{package_name}.utils.image"] = image_module

    network_module = types.ModuleType(f"{package_name}.utils.network")
    network_module.build_worker_url = lambda worker: "http://worker"
    network_module.get_client_session = lambda: None
    network_module.probe_worker = lambda *_args, **_kwargs: None
    sys.modules[f"{package_name}.utils.network"] = network_module

    audio_payload_module = types.ModuleType(f"{package_name}.utils.audio_payload")
    audio_payload_module.encode_audio_payload = lambda audio: audio
    sys.modules[f"{package_name}.utils.audio_payload"] = audio_payload_module

    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module

    spec = importlib.util.spec_from_file_location(f"{package_name}.nodes.collector", module_path)
    module = importlib.util.module_from_spec(spec)
    assert spec is not None and spec.loader is not None
    spec.loader.exec_module(module)
    return module


def test_collector_opts_into_comfyui_list_inputs():
    collector = _load_collector_module().DistributedCollectorNode

    assert collector.INPUT_IS_LIST is True


def test_collector_exposes_images_as_optional_input():
    input_types = _load_collector_module().DistributedCollectorNode.INPUT_TYPES()

    assert "images" not in input_types["required"]
    assert input_types["optional"]["images"] == ("IMAGE",)


def test_audio_only_pass_through_returns_no_images_and_preserves_audio():
    collector = _load_collector_module().DistributedCollectorNode()
    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}

    images, returned_audio = collector.run(images=None, audio=[audio])

    assert images is None
    assert returned_audio is audio


def test_collector_rejects_missing_images_and_audio():
    collector = _load_collector_module().DistributedCollectorNode()

    try:
        collector.run(images=None, audio=None)
    except ValueError as exc:
        assert "image or audio" in str(exc).lower()
    else:
        raise AssertionError("Expected collector to reject a run with no media input")


def test_delegate_only_master_allows_no_local_media_input():
    collector = _load_collector_module().DistributedCollectorNode()

    images, audio = collector.run(
        images=None,
        audio=None,
        multi_job_id=["delegate-audio-job"],
        delegate_only=[True],
        enabled_worker_ids=["[]"],
    )

    assert images is None
    assert tuple(audio["waveform"].shape) == (1, 2, 1)


def test_pass_through_collapses_comfyui_image_list_to_batch_and_unwraps_hidden_inputs():
    collector = _load_collector_module().DistributedCollectorNode()
    first = torch.zeros(1, 2, 2, 3)
    second = torch.ones(1, 2, 2, 3)

    images, audio = collector.run(
        images=[first, second],
        load_balance=[False],
        audio=[None],
        multi_job_id=[""],
        is_worker=[False],
        master_url=[""],
        enabled_worker_ids=["[]"],
        worker_batch_size=[1],
        worker_id=[""],
        pass_through=[False],
        delegate_only=[False],
    )

    assert tuple(images.shape) == (2, 2, 2, 3)
    assert torch.equal(images[0:1], first)
    assert torch.equal(images[1:2], second)
    assert tuple(audio["waveform"].shape) == (1, 2, 1)


def test_worker_list_input_sends_one_completion_sequence_with_last_only_on_final_item():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    first = torch.zeros(1, 2, 2, 3)
    second = torch.ones(1, 2, 2, 3)
    posted_payloads = []

    class _FakeImage:
        def save(self, fp, format=None, compress_level=None):
            fp.write(b"png-bytes")

    class _FakeResponse:
        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb):
            return False

        def raise_for_status(self):
            return None

    class _FakeSession:
        def post(self, url, json, timeout):
            posted_payloads.append(json)
            return _FakeResponse()

    async def _fake_get_client_session():
        return _FakeSession()

    module.tensor_to_pil = lambda *_args, **_kwargs: _FakeImage()
    module.get_client_session = _fake_get_client_session
    module.encode_audio_payload = lambda _audio: None

    images, audio = collector.run(
        images=[first, second],
        load_balance=[False],
        audio=[None],
        multi_job_id=["job-list-1"],
        is_worker=[True],
        master_url=["http://master"],
        enabled_worker_ids=["[]"],
        worker_batch_size=[1],
        worker_id=["worker-a"],
        pass_through=[False],
        delegate_only=[False],
    )

    assert tuple(images.shape) == (2, 2, 2, 3)
    assert tuple(audio["waveform"].shape) == (1, 2, 1)
    assert len(posted_payloads) == 2
    assert [payload["batch_idx"] for payload in posted_payloads] == [0, 1]
    assert [payload["is_last"] for payload in posted_payloads] == [False, True]
    assert {payload["job_id"] for payload in posted_payloads} == {"job-list-1"}
    assert {payload["worker_id"] for payload in posted_payloads} == {"worker-a"}


def test_audio_only_worker_sends_one_completion_without_image():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    audio = {"waveform": torch.ones(1, 2, 4), "sample_rate": 48000}
    posted = []

    class _FakeResponse:
        async def __aenter__(self):
            return self

        async def __aexit__(self, exc_type, exc, tb):
            return False

        def raise_for_status(self):
            return None

    class _FakeSession:
        def post(self, url, json, timeout):
            posted.append((url, json, timeout.total))
            return _FakeResponse()

    async def _fake_get_client_session():
        return _FakeSession()

    module.get_client_session = _fake_get_client_session
    module.encode_audio_payload = lambda value: {"encoded": value is audio}

    images, returned_audio = collector.run(
        images=None,
        audio=[audio],
        multi_job_id=["audio-job"],
        is_worker=[True],
        master_url=["http://master"],
        worker_id=["worker-a"],
    )

    assert images is None
    assert returned_audio is audio
    assert len(posted) == 1
    assert posted[0][0] == "http://master/distributed/job_complete"
    assert posted[0][2] == 600
    assert posted[0][1] == {
        "job_id": "audio-job",
        "worker_id": "worker-a",
        "batch_idx": 0,
        "audio": {"encoded": True},
        "is_last": True,
    }


def test_audio_only_master_combines_local_and_worker_audio():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    master_audio = {"waveform": torch.ones(1, 2, 2), "sample_rate": 48000}
    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
    queue = asyncio.Queue()
    queue.put_nowait(
        {
            "worker_id": "worker-a",
            "image_index": 0,
            "tensor": None,
            "audio": worker_audio,
            "is_last": True,
        }
    )
    module.prompt_server.distributed_pending_jobs = {"audio-job": queue}

    images, combined_audio = asyncio.run(
        collector.execute(
            images=None,
            audio=master_audio,
            multi_job_id="audio-job",
            enabled_worker_ids='["worker-a"]',
        )
    )

    assert images is None
    assert combined_audio["sample_rate"] == 48000
    assert tuple(combined_audio["waveform"].shape) == (1, 2, 5)
    assert torch.equal(combined_audio["waveform"][..., :2], master_audio["waveform"])
    assert torch.equal(combined_audio["waveform"][..., 2:], worker_audio["waveform"])


def test_delegate_only_audio_collects_worker_audio_without_placeholder_image():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    worker_audio = {"waveform": torch.full((1, 2, 3), 2.0), "sample_rate": 48000}
    module.prompt_server.distributed_jobs_lock = asyncio.Lock()
    queue = asyncio.Queue()
    queue.put_nowait(
        {
            "worker_id": "worker-a",
            "image_index": 0,
            "tensor": None,
            "audio": worker_audio,
            "is_last": True,
        }
    )
    module.prompt_server.distributed_pending_jobs = {"delegate-audio-job": queue}

    images, combined_audio = asyncio.run(
        collector.execute(
            images=None,
            audio=None,
            multi_job_id="delegate-audio-job",
            enabled_worker_ids='["worker-a"]',
            delegate_only=True,
        )
    )

    assert images is None
    assert combined_audio["sample_rate"] == 48000
    assert torch.equal(combined_audio["waveform"], worker_audio["waveform"])


# --------------------------------------------------------------------------- #
# Compressed worker -> master transfer
# --------------------------------------------------------------------------- #

class _RecordingResponse:
    def __init__(self, status=200, text=""):
        self.status = status
        self._text = text

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        return False

    async def text(self):
        return self._text

    def raise_for_status(self):
        if self.status >= 400:
            raise RuntimeError(f"HTTP {self.status}")


class _RecordingSession:
    def __init__(self, statuses=None, text="Invalid JSON payload"):
        self.calls = []
        self._statuses = list(statuses or [])
        self._text = text

    def post(self, url, timeout, json=None, data=None):
        self.calls.append({"url": url, "json": json, "data": data, "timeout": timeout.total})
        status = self._statuses.pop(0) if self._statuses else 200
        return _RecordingResponse(status=status, text=self._text)


def _install_fake_transport(module, session):
    async def _fake_get_client_session():
        return session

    module.get_client_session = _fake_get_client_session
    # Record the multipart content instead of building a real aiohttp.FormData.
    module.DistributedCollectorNode._build_multipart = staticmethod(
        lambda envelope, parts: {"envelope": envelope, "parts": list(parts)}
    )


def _run_worker(collector, images, audio=None, **kwargs):
    return collector.run(
        images=[images],
        audio=[audio],
        multi_job_id=["job-c"],
        is_worker=[True],
        master_url=["http://master"],
        worker_id=["worker-a"],
        **{key: [value] for key, value in kwargs.items()},
    )


def test_collector_exposes_transfer_widgets_with_legacy_default():
    module = _load_collector_module()
    optional = module.DistributedCollectorNode.INPUT_TYPES()["optional"]

    formats, options = optional["transfer_format"]
    assert "legacy_png" in formats and "h264" in formats and "ffv1" in formats
    assert options["default"] == "legacy_png"
    assert optional["transfer_quality"][0] == "INT"


def test_worker_sends_video_batch_as_one_multipart_chunk_with_audio_on_last():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    module.encode_audio_payload = lambda value: None if value is None else {"encoded": True}
    images = torch.rand(6, 32, 48, 3)

    _run_worker(
        collector,
        images,
        audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
        transfer_format="ffv1",
        transfer_quality=90,
    )

    assert len(session.calls) == 1
    call = session.calls[0]
    assert call["url"] == "http://master/distributed/job_complete"
    assert call["json"] is None
    envelope = call["data"]["envelope"]
    assert envelope["format"] == "ffv1"
    assert (envelope["start_index"], envelope["count"]) == (0, 6)
    assert (envelope["width"], envelope["height"]) == (48, 32)
    assert envelope["is_last"] is True
    assert envelope["audio"] == {"encoded": True}
    assert envelope["job_id"] == "job-c" and envelope["worker_id"] == "worker-a"
    assert len(call["data"]["parts"]) == 1


def test_worker_marks_only_final_chunk_as_last():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    module.TRANSFER_MAX_CHUNK_BYTES = 1  # force one frame per request

    _run_worker(collector, torch.rand(3, 8, 8, 3), transfer_format="png")

    envelopes = [call["data"]["envelope"] for call in session.calls]
    assert [(e["start_index"], e["count"], e["is_last"]) for e in envelopes] == [
        (0, 1, False),
        (1, 1, False),
        (2, 1, True),
    ]
    assert "audio" not in envelopes[-1]


def test_worker_falls_back_to_legacy_png_when_master_rejects_multipart():
    from PIL import Image
    import numpy as np

    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession(statuses=[400])
    _install_fake_transport(module, session)
    module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
        (255 * tensor[idx].cpu().numpy()).astype(np.uint8)
    )

    _run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="h264")

    assert session.calls[0]["data"] is not None  # compressed attempt
    legacy = [call["json"] for call in session.calls[1:]]
    assert [payload["batch_idx"] for payload in legacy] == [0, 1]
    assert all(payload["image"].startswith("data:image/png;base64,") for payload in legacy)
    assert [payload["is_last"] for payload in legacy] == [False, True]


def test_worker_uses_legacy_path_for_unknown_transfer_format():
    from PIL import Image
    import numpy as np

    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    module.tensor_to_pil = lambda tensor, idx: Image.fromarray(
        (255 * tensor[idx].cpu().numpy()).astype(np.uint8)
    )

    _run_worker(collector, torch.rand(1, 8, 8, 3), transfer_format="gif")

    assert len(session.calls) == 1
    assert session.calls[0]["json"]["image"].startswith("data:image/png;base64,")


def test_worker_falls_back_when_video_encoder_missing():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    codec = module.transfer_codec
    original = codec.video_codec_available
    codec.video_codec_available = lambda fmt: fmt == "h264" and original("h264")
    try:
        _run_worker(collector, torch.rand(2, 16, 16, 3), transfer_format="av1")
    finally:
        codec.video_codec_available = original

    assert session.calls[0]["data"]["envelope"]["format"] == "h264"


def test_worker_does_not_fall_back_on_non_legacy_rejection():
    import pytest

    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession(statuses=[400], text="Invalid multipart payload: Request Entity Too Large")
    _install_fake_transport(module, session)

    with pytest.raises(RuntimeError, match="MAX_REQUEST_BYTES"):
        _run_worker(collector, torch.rand(2, 8, 8, 3), transfer_format="png")

    # One compressed attempt, no legacy resend of the batch.
    assert len(session.calls) == 1
    assert session.calls[0]["json"] is None


def test_worker_sends_nothing_when_a_frame_cannot_fit():
    import numpy as np
    import pytest

    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
    module.TRANSFER_MAX_REQUEST_BYTES = reserve + 32 * 1024
    module.TRANSFER_MAX_CHUNK_BYTES = 32 * 1024
    # Frame 0 is tiny; frame 2 is noise that cannot fit in 32 KiB as PNG.
    images = torch.zeros(3, 128, 128, 3)
    images[2] = torch.from_numpy(np.random.default_rng(1).random((128, 128, 3), dtype=np.float32))

    with pytest.raises(module.transfer_codec.FrameTooLargeError):
        _run_worker(collector, images, transfer_format="png")

    assert session.calls == []


def test_worker_sends_audio_separately_when_it_would_overflow_last_request():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    reserve = module.DistributedCollectorNode.TRANSFER_ENVELOPE_RESERVE_BYTES
    module.TRANSFER_MAX_REQUEST_BYTES = reserve + 64 * 1024
    module.TRANSFER_MAX_CHUNK_BYTES = 64 * 1024
    module.encode_audio_payload = lambda value: None if value is None else {"data": "x" * (128 * 1024)}

    _run_worker(
        collector,
        torch.zeros(2, 16, 16, 3),
        audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
        transfer_format="png",
    )

    multipart = [call["data"]["envelope"] for call in session.calls if call["data"] is not None]
    assert [e["is_last"] for e in multipart] == [False]
    assert all("audio" not in e for e in multipart)
    final = session.calls[-1]["json"]
    assert final["is_last"] is True and final["audio"]["data"].startswith("x")
    assert "image" not in final


def test_worker_attaches_audio_to_last_chunk_when_it_fits():
    module = _load_collector_module()
    collector = module.DistributedCollectorNode()
    session = _RecordingSession()
    _install_fake_transport(module, session)
    module.encode_audio_payload = lambda value: None if value is None else {"data": "small"}

    _run_worker(
        collector,
        torch.zeros(2, 16, 16, 3),
        audio={"waveform": torch.ones(1, 2, 4), "sample_rate": 48000},
        transfer_format="png",
    )

    assert len(session.calls) == 1
    envelope = session.calls[0]["data"]["envelope"]
    assert envelope["is_last"] is True and envelope["audio"] == {"data": "small"}