from __future__ import annotations

import io
import shutil
import tempfile
import traceback
from abc import ABC, abstractmethod
from contextlib import AbstractContextManager, contextmanager
from dataclasses import dataclass
from io import BytesIO
from queue import Empty
from typing import TYPE_CHECKING

from typing_extensions import Self

import qqabc.qq

if TYPE_CHECKING:
    from collections.abc import Callable


@dataclass
class InData:
    task_id: int
    url: str
    fpath: str | None = None


@dataclass
class OutData:
    task_id: int
    data: BytesIO


class DataDeletedError(KeyError):
    def __init__(self, task_id: int):
        super().__init__(f"Output data for task_id {task_id} has been deleted.")


class WorkersDiedOutError(RuntimeError):
    def __init__(self):
        super().__init__("All workers have stopped unexpectedly.")


class QQBugError(RuntimeError):
    def __init__(self, msg: str):
        super().__init__(msg)


class InvalidTaskError(ValueError):
    def __init__(self, task_id: int):
        super().__init__(f"Invalid task_id: {task_id}")


class IStorage(ABC):
    @abstractmethod
    def register(self, indata: InData):
        pass

    @abstractmethod
    def save(self, task_id: int, outdata: OutData):
        pass

    @abstractmethod
    def load(self, task_id: int) -> OutData:
        """Load the output data associated with the given task ID.

        Raises ValueError if the data has been deleted.
        """

    @abstractmethod
    def delete(self, task_id: int) -> None:
        pass

    @abstractmethod
    def delete_all(self) -> None:
        pass

    @abstractmethod
    def has(self, task_id: int) -> bool:
        pass


class Storage(IStorage):
    def __init__(self, cached_size: int):
        self.cached_size = cached_size
        self.indata_storage: dict[int, InData] = {}
        self.outdata_storage: dict[int, OutData] = {}
        self.size = 0
        self.saved: set[int] = set()

    def register(self, indata: InData):
        self.indata_storage[indata.task_id] = indata

    def save(self, task_id: int, outdata: OutData):
        if task_id in self.saved:
            raise ValueError(
                f"Output data for task_id {task_id} has already been saved."
            )
        this_size = outdata.data.getbuffer().nbytes
        while self.size + this_size > self.cached_size and self.outdata_storage:
            oldest_task_id = min(self.outdata_storage)
            self.delete(oldest_task_id)
        self.saved.add(task_id)
        if self.size + this_size > self.cached_size:
            indata = self.indata_storage.pop(task_id)
            self._save_to_disk(indata, outdata)
        else:
            self.size += this_size
            self.outdata_storage[task_id] = outdata

    def load(self, task_id: int) -> OutData:
        if task_id not in self.outdata_storage and task_id in self.saved:
            raise DataDeletedError(task_id)
        return self.outdata_storage[task_id]

    def _save_to_disk(self, indata: InData, outdata: OutData):
        if indata.fpath is not None:
            with tempfile.NamedTemporaryFile(delete=False) as tmpf:
                tmpf.write(outdata.data.getbuffer())
            shutil.move(tmpf.name, indata.fpath)
            self.size -= outdata.data.getbuffer().nbytes

    def delete(self, task_id: int) -> None:
        outdata = self.outdata_storage.pop(task_id)
        indata = self.indata_storage.pop(task_id)
        self._save_to_disk(indata, outdata)

    def delete_all(self) -> None:
        for task_id in list(self.outdata_storage.keys()):
            self.delete(task_id)

    def has(self, task_id: int) -> bool:
        return task_id in self.saved


class IWorker(ABC):
    @abstractmethod
    def start(self, worker_id: int) -> AbstractContextManager[IWorker]:
        pass

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

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


class DefaultWorker(IWorker):
    @contextmanager
    def start(self, worker_id: int):
        self.worker_id = worker_id
        import httpx  # noqa: PLC0415

        with httpx.Client(follow_redirects=True) as client:
            self.client = client
            yield self

    def resolve(self, indata: InData) -> OutData:
        resp = self.client.get(indata.url)
        resp.raise_for_status()
        b = BytesIO(resp.content)
        return OutData(task_id=indata.task_id, data=b)


def _worker_download(
    input_q: qqabc.qq.Q[InData],
    output_q: qqabc.qq.Q[int],
    worker: IWorker,
    worker_id: int,
    storage: IStorage,
):
    try:
        with worker.start(worker_id):
            for msg in input_q.iter(worker.input_timeout):
                ind = msg.data
                try:
                    outd = worker.resolve(ind)
                    storage.save(ind.task_id, outd)
                    output_q.put(ind.task_id)
                except Exception:
                    input_q.put(msg)
                    raise
    except Empty:
        pass
    except Exception:
        traceback.print_exc()


class Resolver:
    def __init__(
        self,
        num_workers: int,
        *,
        storage: IStorage,
        worker_factory: Callable[[], IWorker],
    ):
        input_q = qqabc.qq.Q[InData]("thread")
        output_q = qqabc.qq.Q[int]("thread")
        self.input_q = input_q
        self.output_q = output_q
        self.storage = storage
        self.workers = [
            qqabc.qq.run_thread(
                _worker_download, input_q, output_q, worker_factory(), w, self.storage
            )
            for w in range(num_workers)
        ]
        self.task_cnt = 0
        self.url_max = 512

    def __enter__(self) -> Self:
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()
        self.storage.delete_all()

    def _get_task_id(self):
        self.task_cnt += 1
        return self.task_cnt

    @contextmanager
    def open(self, filepath: str, mode: str = "r"):
        outd = None
        with open(filepath, "rb") as f:
            if f.seek(0, 2) <= self.url_max:
                try:
                    f.seek(0)
                    url = f.read(self.url_max).decode("utf-8")
                    outd = self.add_wait(url.strip(), fname=filepath)
                except UnicodeDecodeError:
                    pass
                except DataDeletedError:
                    pass
        if outd is None:
            with open(filepath, mode) as f:
                yield f
        else:
            outd.data.seek(0)
            if "b" in mode:
                yield outd.data
            else:
                yield io.StringIO(outd.data.read().decode("utf-8"))

    def add_wait(self, url, fname: str | None = None):
        task_id = self.add(url, fname=fname)
        return self.wait(task_id)

    def add(self, url: str, fname: str | None = None) -> int:
        task_id = self._get_task_id()
        indata = InData(task_id=task_id, url=url, fpath=fname)
        self.storage.register(indata)
        self.input_q.put(indata)
        return task_id

    def _get_result(self, task_id: int):
        return self.storage.load(task_id)

    def iter_and_close(self):
        self.close()
        for msg in self.output_q:
            task_id = msg.data
            yield self._get_result(task_id)

    def close(self):
        self.input_q.stop(self.workers)
        self.output_q.end()

    def completed(self, timeout: float = 0):
        """Generator that yields completed tasks as they finish.

        If timeout is set to a float value, the generator will yield
        completed tasks until no more tasks are available within the timeout period.
        """
        for msg in self._iter(timeout=timeout, empty_ok=True):
            task_id = msg.data
            yield self._get_result(task_id)

    def _iter(self, timeout: float = 0.05, *, empty_ok: bool):
        """Generator that yields completed tasks as they finish.

        If timeout is set to a float value, the generator will yield
        completed tasks until no more tasks are available within the timeout period.
        """
        while True:
            try:
                yield from self.output_q.iter(timeout=timeout)
            except Empty:  # noqa: PERF203
                if all(not worker.is_alive() for worker in self.workers):
                    raise WorkersDiedOutError from None
                if empty_ok:
                    break

    def wait(self, task_id: int):
        if not (0 < task_id <= self.task_cnt):
            raise InvalidTaskError(task_id)
        if self.storage.has(task_id):
            return self._get_result(task_id)
        for msg in self._iter(timeout=0.05, empty_ok=False):
            completed_task_id = msg.data
            if completed_task_id == task_id:
                return self._get_result(task_id)
        raise QQBugError("Unreachable code reached.")


def resolve(
    num_workers: int = 4,
    *,
    cache_size: int = 1 << 20,
    worker: type[IWorker] | Callable[[], IWorker] | None = None,
) -> Resolver:
    storage: IStorage = Storage(cached_size=cache_size)

    return Resolver(
        num_workers,
        storage=storage,
        worker_factory=worker if worker is not None else DefaultWorker,
    )
