from __future__ import annotations

import re

import pytest
from typing_extensions import Literal


def get_url(read_mode: Literal["r", "rb"]) -> str:
    if read_mode == "rb":
        return "https://picsum.photos/200"
    if read_mode == "r":
        return "https://www.lipsum.com/feed/html/"
    raise NotImplementedError


def test_import_version():
    from qqabc import __version__

    assert re.match(r"^v?\d+\.\d+\.\d+([.-]\w+)?$", __version__)


def test_usage():
    """測試add_wait基本功能

    測試使用add_wait方法來加入一個URL下載任務
    完成後需要成功取得下載的資料。
    """
    from qqabc.rurl import resolve

    url = get_url("rb")
    with resolve() as resolver:
        od = resolver.add_wait(url)
        assert od
        b = od.data
    assert b.seek(0, 2) > 1024


def test_usage2():
    """測試add + completed的使用情境

    測試加入多個任務，然後使用completed方法來取得所有完成的任務結果。
    """
    from qqabc.rurl import resolve

    url = get_url("rb")
    tasks = set()
    with resolve() as resolver:
        for _ in range(2):
            tasks.add(resolver.add(url))
        for task in resolver.iter_and_close():
            b = task.data
            tasks.remove(task.task_id)
            assert b.seek(0, 2) > 1024
        assert len(tasks) == 0


def test_usage21():
    """測試add + completed的使用情境

    測試加入多個任務，然後使用completed方法來取得所有完成的任務結果。
    這次加上邊跑邊加任務的情境。
    """
    from qqabc.rurl import resolve

    url = get_url("rb")
    tasks = set()
    with resolve() as resolver:
        todos = set(range(4))
        for _ in range(2):
            todos.pop()
            tasks.add(resolver.add(url))
        for task in resolver.completed(timeout=2):
            if todos:
                todos.pop()
                tasks.add(resolver.add(url))
            b = task.data
            tasks.remove(task.task_id)
            assert b.seek(0, 2) > 1024
        assert len(tasks) == 0


def test_usage3():
    """測試add + wait的使用情境

    測試加入多個任務，然後使用wait方法來等待每個任務完成並取得結果。
    """
    from qqabc.rurl import resolve

    url = get_url("rb")
    tasks = set()
    with resolve() as resolver:
        for _ in range(2):
            tasks.add(resolver.add(url))
        for task_id in tasks:
            od = resolver.wait(task_id)
            assert od
            b = od.data
            assert b.seek(0, 2) > 1024


@pytest.mark.parametrize("read_mode", ["rb", "r"])
def test_usage4(tmpdir, read_mode: str):
    """測試open方法的使用情境

    測試使用open方法來打開一個包含URL的檔案, 會回傳IO物件來讀取下載的資料。
    預設cache size足夠大到裝的下測試的檔案, 此時不會馬上把下載的資料存到硬碟。
    resolver關閉後, 資料會被存回硬碟, 此時打開會讀到下載過後的檔案而非URL。
    """
    url = get_url(read_mode)
    with open(tmpdir / "urls.txt", "w") as f:
        f.write(url)
    from qqabc.rurl import resolve

    with resolve() as resolver:
        with resolver.open(tmpdir / "urls.txt", read_mode) as fp:
            assert isinstance(fp.read(1), bytes if read_mode == "rb" else str)
            assert fp.seek(0, 2) > 1024
        with open(tmpdir / "urls.txt") as fp:
            assert fp.read() == url
    with open(tmpdir / "urls.txt", read_mode) as fp:
        assert fp.seek(0, 2) > 1024


@pytest.mark.parametrize("read_mode", ["rb", "r"])
def test_usage5(tmpdir, read_mode: str):
    """測試open方法的使用情境, cache size設為0

    測試使用open方法來打開一個包含URL的檔案, 會回傳IO物件來讀取下載的資料。
    這次cache size設為0, 代表不會把下載的資料存在記憶體中, 會直接存到硬碟。
    resolver關閉後, 資料已經存在硬碟中, 此時
    打開會讀到下載過後的檔案而非URL。
    """
    url = get_url(read_mode)
    with open(tmpdir / "urls.txt", "w") as f:
        f.write(url + "\n")
    from qqabc.rurl import resolve

    with resolve(cache_size=0) as resolver:
        with resolver.open(tmpdir / "urls.txt", read_mode) as fp:
            assert fp.seek(0, 2) > 1024

    with open(tmpdir / "urls.txt", read_mode) as fp:
        assert fp.seek(0, 2) > 1024


def test_usage6():
    """測試Worker啟動失敗的情境

    測試使用一個會在啟動時拋出例外的Worker類別,
    來確保Resolver能正確處理Worker啟動失敗的情況,會退出process並拋出WorkersDiedOutError例外。
    """
    from qqabc.rurl import InData, IWorker, OutData, WorkersDiedOutError, resolve

    class BadWorker(IWorker):
        def start(self, worker_id: int):
            raise RuntimeError("Failed to start worker")

        def resolve(self, indata: InData) -> OutData:
            pass

    url = get_url("rb")
    with resolve(worker=BadWorker) as resolver, pytest.raises(WorkersDiedOutError):
        resolver.add_wait(url)


def test_usage61():
    """測試Worker啟動失敗的情境

    測試使用一個會在啟動時拋出例外的Worker類別,
    來確保Resolver能正確處理Worker啟動失敗的情況。
    這裡只要還有一個Worker啟動成功, 就能正常下載資料。
    """
    from qqabc.rurl import DefaultWorker, resolve

    class Worker(DefaultWorker):
        def start(self, worker_id: int):
            if worker_id == 3:
                import time

                time.sleep(1)
            return super().start(worker_id)

        def resolve(self, indata):
            if self.worker_id != 3:
                raise RuntimeError("Failed to download")
            return super().resolve(indata)

        @property
        def input_timeout(self) -> float | None:
            return 4.0

    url = get_url("rb")
    with resolve(num_workers=4, worker=Worker) as resolver:
        resolver.add_wait(url)


def test_usage7(tmpdir):
    """測試open方法的使用情境

    測試使用open方法來打開多個包含URL的檔案, 會回傳IO物件來讀取下載的資料。
    使用cache size能裝一個但裝不下兩個的情境。
    在第二個檔案打開時, 會把第一個檔案的下載資料存回硬碟。
    resolver關閉後, 資料會被存回硬碟, 此時打開會讀到下載過後的檔案而非URL。

    這次使用自訂的Worker來限制下載內容的大小, 以確保測試的穩定性。
    下載的內容會是4500 bytes, 因此cache size設為5000 bytes能裝下一個但裝不下兩個。
    """
    from io import BytesIO

    from qqabc.rurl import DefaultWorker, InData, OutData, resolve

    class Worker(DefaultWorker):
        def resolve(self, indata: InData) -> OutData:
            resp = self.client.get(indata.url)
            resp.raise_for_status()
            content = resp.content[:4500]
            assert len(content) == 4500
            b = BytesIO(content)
            return OutData(task_id=indata.task_id, data=b)

    url = get_url("rb")
    with open(tmpdir / "urls1.txt", "w") as f:
        f.write(url)
    with open(tmpdir / "urls2.txt", "w") as f:
        f.write(url)

    with resolve(cache_size=5000, worker=Worker) as resolver:
        with resolver.open(tmpdir / "urls1.txt", "rb") as fp:
            assert isinstance(fp.read(1), bytes)
            assert fp.seek(0, 2) > 1024
        with open(tmpdir / "urls1.txt") as fp:
            assert fp.read() == url
        with resolver.open(tmpdir / "urls2.txt", "rb") as fp:
            assert isinstance(fp.read(1), bytes)
            assert fp.seek(0, 2) > 1024
        with open(tmpdir / "urls1.txt", "rb") as fp:
            assert fp.seek(0, 2) > 1024
        with open(tmpdir / "urls2.txt") as fp:
            assert fp.read() == url
    with open(tmpdir / "urls2.txt", "rb") as fp:
        assert fp.seek(0, 2) > 1024
