"""Tests for file transfer: download and upload handlers."""

from __future__ import annotations

import base64
import hashlib
import os
import tempfile

import pytest

from qremotex.transfer.downloader import FileDownloader, MAX_FILE_SIZE_BYTES
from qremotex.transfer.uploader import FileUploader, CHUNK_SIZE_RAW


# ──────────────────────────────────────────── Downloader Tests ──


class TestFileDownloader:
    def setup_method(self):
        self.tmpdir = tempfile.mkdtemp()
        self.dl = FileDownloader(download_dir=self.tmpdir)

    def teardown_method(self):
        self.dl.cleanup_all()
        import shutil
        shutil.rmtree(self.tmpdir, ignore_errors=True)

    def _start_msg(self, file_id="f1", target="out/data.json", chunks=2, size=100):
        return {
            "file_id": file_id,
            "target_path": target,
            "total_chunks": chunks,
            "total_size_bytes": size,
        }

    def _chunk_msg(self, file_id="f1", index=0, data=b"hello"):
        return {
            "file_id": file_id,
            "chunk_index": index,
            "data_b64": base64.b64encode(data).decode(),
        }

    def _complete_msg(self, file_id="f1", checksum=""):
        return {"file_id": file_id, "checksum": checksum}

    def test_full_download_success(self):
        data = b"hello world test data"
        sha = hashlib.sha256(data).hexdigest()

        r = self.dl.handle_start(self._start_msg(size=len(data), chunks=1))
        assert r["ok"]

        r = self.dl.handle_chunk(self._chunk_msg(data=data))
        assert r["ok"]

        r = self.dl.handle_complete(self._complete_msg(checksum=f"sha256:{sha}"))
        assert r["ok"]
        assert r["size_bytes"] == len(data)

        # File should exist at target
        assert os.path.isfile(r["local_path"])
        with open(r["local_path"], "rb") as f:
            assert f.read() == data

    def test_start_missing_fields(self):
        r = self.dl.handle_start({"file_id": "", "target_path": ""})
        assert not r["ok"]
        assert "Missing" in r["error"]

    def test_path_traversal_blocked(self):
        r = self.dl.handle_start(self._start_msg(target="../../../etc/passwd"))
        assert not r["ok"]
        assert "traversal" in r["error"].lower()

    def test_chunk_no_active_download(self):
        r = self.dl.handle_chunk(self._chunk_msg(file_id="nonexistent"))
        assert not r["ok"]
        assert "No active download" in r["error"]

    def test_complete_no_active_download(self):
        r = self.dl.handle_complete(self._complete_msg(file_id="nonexistent"))
        assert not r["ok"]

    def test_checksum_mismatch(self):
        data = b"some data"
        self.dl.handle_start(self._start_msg(size=len(data), chunks=1))
        self.dl.handle_chunk(self._chunk_msg(data=data))
        r = self.dl.handle_complete(self._complete_msg(checksum="sha256:0000bad"))
        assert not r["ok"]
        assert "mismatch" in r["error"].lower()

    def test_complete_without_checksum_still_succeeds(self):
        data = b"data"
        self.dl.handle_start(self._start_msg(size=len(data), chunks=1))
        self.dl.handle_chunk(self._chunk_msg(data=data))
        r = self.dl.handle_complete(self._complete_msg(checksum=""))
        assert r["ok"]

    def test_multi_chunk_download(self):
        chunk1 = b"A" * 100
        chunk2 = b"B" * 50
        combined = chunk1 + chunk2
        sha = hashlib.sha256(combined).hexdigest()

        self.dl.handle_start(self._start_msg(size=len(combined), chunks=2))
        self.dl.handle_chunk(self._chunk_msg(index=0, data=chunk1))
        self.dl.handle_chunk(self._chunk_msg(index=1, data=chunk2))
        r = self.dl.handle_complete(self._complete_msg(checksum=f"sha256:{sha}"))
        assert r["ok"]
        assert r["size_bytes"] == 150

    def test_active_count(self):
        assert self.dl.active_count == 0
        self.dl.handle_start(self._start_msg(file_id="a"))
        assert self.dl.active_count == 1
        self.dl.handle_start(self._start_msg(file_id="b", target="out/b.json"))
        assert self.dl.active_count == 2

    def test_cleanup_all(self):
        self.dl.handle_start(self._start_msg(file_id="a"))
        self.dl.handle_start(self._start_msg(file_id="b", target="out/b.json"))
        self.dl.cleanup_all()
        assert self.dl.active_count == 0

    def test_file_too_large_start(self):
        r = self.dl.handle_start(self._start_msg(size=MAX_FILE_SIZE_BYTES + 1))
        assert not r["ok"]
        assert "too large" in r["error"].lower()


# ──────────────────────────────────────────── Uploader Tests ──


class TestFileUploader:
    def setup_method(self):
        self.sent_messages: list[dict] = []

        async def mock_send(msg):
            self.sent_messages.append(msg)

        self.uploader = FileUploader(send_fn=mock_send, executor_id="exec-test-1")
        self.tmpdir = tempfile.mkdtemp()

    def teardown_method(self):
        import shutil
        shutil.rmtree(self.tmpdir, ignore_errors=True)

    def _create_file(self, name: str, content: bytes) -> str:
        path = os.path.join(self.tmpdir, name)
        with open(path, "wb") as f:
            f.write(content)
        return path

    @pytest.mark.asyncio
    async def test_upload_small_file(self):
        data = b"hello upload test"
        path = self._create_file("small.txt", data)

        result = await self.uploader.upload(path, "results/small.txt", "req-1")

        assert result["ok"]
        assert result["size_bytes"] == len(data)
        assert result["checksum"] == f"sha256:{hashlib.sha256(data).hexdigest()}"

        # Should have sent: start, 1 chunk, complete
        types = [m["type"] for m in self.sent_messages]
        assert types == ["upload_file_start", "upload_file_chunk", "upload_file_complete"]

    @pytest.mark.asyncio
    async def test_upload_multi_chunk(self):
        # Create a file larger than one chunk
        data = b"X" * (CHUNK_SIZE_RAW + 100)
        path = self._create_file("big.bin", data)

        result = await self.uploader.upload(path, "results/big.bin", "req-2")

        assert result["ok"]
        assert result["size_bytes"] == len(data)

        types = [m["type"] for m in self.sent_messages]
        assert types[0] == "upload_file_start"
        assert types[-1] == "upload_file_complete"
        chunk_msgs = [m for m in self.sent_messages if m["type"] == "upload_file_chunk"]
        assert len(chunk_msgs) == 2  # ceil(CHUNK_SIZE_RAW + 100 / CHUNK_SIZE_RAW) = 2

    @pytest.mark.asyncio
    async def test_upload_file_not_found(self):
        result = await self.uploader.upload("/nonexistent/file.txt", "results/x.txt", "req-3")
        assert not result["ok"]
        assert "not found" in result["error"].lower()

    @pytest.mark.asyncio
    async def test_upload_empty_file(self):
        path = self._create_file("empty.txt", b"")
        result = await self.uploader.upload(path, "results/empty.txt", "req-4")
        assert result["ok"]
        assert result["size_bytes"] == 0

    @pytest.mark.asyncio
    async def test_upload_preserves_request_id(self):
        data = b"test"
        path = self._create_file("t.txt", data)
        await self.uploader.upload(path, "results/t.txt", "req-42")

        for msg in self.sent_messages:
            assert msg["request_id"] == "req-42"
            assert msg["executor_id"] == "exec-test-1"

    @pytest.mark.asyncio
    async def test_upload_content_type_detection(self):
        path = self._create_file("data.json", b'{"a":1}')
        await self.uploader.upload(path, "results/data.json", "req-5")
        start_msg = self.sent_messages[0]
        assert start_msg["content_type"] == "application/json"

    @pytest.mark.asyncio
    async def test_upload_chunk_data_integrity(self):
        data = b"verify chunk data"
        path = self._create_file("check.txt", data)
        await self.uploader.upload(path, "results/check.txt", "req-6")

        chunk_msg = [m for m in self.sent_messages if m["type"] == "upload_file_chunk"][0]
        decoded = base64.b64decode(chunk_msg["data_b64"])
        assert decoded == data
