import dataclasses
import io
import multiprocessing as _mp
import uuid
import zipfile
from concurrent.futures import Future
from multiprocessing.connection import Connection
from typing import List, Optional, Tuple

import numpy

from tiktorch import log
from tiktorch.rpc import Shutdown
from tiktorch.rpc import mp as _mp_rpc
from tiktorch.rpc.mp import MPServer
from tiktorch.server.reader import eval_model_zip

from .backend import base
from .rpc_interface import IRPCModelSession


@dataclasses.dataclass
class ModelInfo:
    # TODO: Test for model info
    name: str
    input_axes: str
    output_axes: str
    valid_shapes: List[List[Tuple[str, int]]]
    halo: List[Tuple[str, int]]
    offset: List[Tuple[str, int]]
    scale: List[Tuple[str, float]]


class ModelSessionProcess(IRPCModelSession):
    def __init__(self, model_zip: bytes, devices: List[str]) -> None:
        with zipfile.ZipFile(io.BytesIO(model_zip)) as model_file:
            self._model = eval_model_zip(model_file, devices)

        self._datasets = {}
        self._worker = base.SessionBackend(self._model)

    def forward(self, input_tensor: numpy.ndarray) -> Future:
        res = self._worker.forward(input_tensor)
        return res

    def create_dataset(self, mean, stddev):
        id_ = uuid.uuid4().hex
        self._datasets[id_] = {"mean": mean, "stddev": stddev}
        return id_

    def get_model_info(self) -> ModelInfo:
        return ModelInfo(
            self._model.name,
            self._model.input_axes,
            self._model.output_axes,
            valid_shapes=[self._model.input_shape],
            halo=self._model.halo,
            scale=self._model.scale,
            offset=self._model.offset,
        )

    def shutdown(self) -> Shutdown:
        self._worker.shutdown()
        return Shutdown()


def _run_model_session_process(
    conn: Connection, model_zip: bytes, devices: List[str], log_queue: Optional[_mp.Queue] = None
):
    try:
        # from: https://github.com/pytorch/pytorch/issues/973#issuecomment-346405667
        import resource

        rlimit = resource.getrlimit(resource.RLIMIT_NOFILE)
        resource.setrlimit(resource.RLIMIT_NOFILE, (4096, rlimit[1]))
    except ModuleNotFoundError:
        pass  # probably running on windows

    if log_queue:
        log.configure(log_queue)

    session_proc = ModelSessionProcess(model_zip, devices)
    srv = MPServer(session_proc, conn)
    srv.listen()


def start_model_session_process(
    model_zip: bytes, devices: List[str], log_queue: Optional[_mp.Queue] = None
) -> Tuple[_mp.Process, IRPCModelSession]:
    client_conn, server_conn = _mp.Pipe()
    proc = _mp.Process(
        target=_run_model_session_process,
        name="ModelSessionProcess",
        kwargs={"conn": server_conn, "devices": devices, "log_queue": log_queue, "model_zip": model_zip},
    )
    proc.start()
    return proc, _mp_rpc.create_client(IRPCModelSession, client_conn)

[((83, 10, 83, 38), 'tiktorch.rpc.mp.MPServer', 'MPServer', ({(83, 19, 83, 31): 'session_proc', (83, 33, 83, 37): 'conn'}, {}), '(session_proc, conn)', False, 'from tiktorch.rpc.mp import MPServer\n'), ((90, 31, 90, 41), 'multiprocessing.Pipe', '_mp.Pipe', ({}, {}), '()', True, 'import multiprocessing as _mp\n'), ((91, 11, 95, 5), 'multiprocessing.Process', '_mp.Process', (), '', True, 'import multiprocessing as _mp\n'), ((64, 15, 64, 25), 'tiktorch.rpc.Shutdown', 'Shutdown', ({}, {}), '()', False, 'from tiktorch.rpc import Shutdown\n'), ((74, 17, 74, 59), 'resource.getrlimit', 'resource.getrlimit', ({(74, 36, 74, 58): 'resource.RLIMIT_NOFILE'}, {}), '(resource.RLIMIT_NOFILE)', False, 'import resource\n'), ((75, 8, 75, 69), 'resource.setrlimit', 'resource.setrlimit', ({(75, 27, 75, 49): 'resource.RLIMIT_NOFILE', (75, 51, 75, 68): '(4096, rlimit[1])'}, {}), '(resource.RLIMIT_NOFILE, (4096, rlimit[1]))', False, 'import resource\n'), ((80, 8, 80, 32), 'tiktorch.log.configure', 'log.configure', ({(80, 22, 80, 31): 'log_queue'}, {}), '(log_queue)', False, 'from tiktorch import log\n'), ((97, 17, 97, 69), 'tiktorch.rpc.mp.create_client', '_mp_rpc.create_client', ({(97, 39, 97, 55): 'IRPCModelSession', (97, 57, 97, 68): 'client_conn'}, {}), '(IRPCModelSession, client_conn)', True, 'from tiktorch.rpc import mp as _mp_rpc\n'), ((37, 26, 37, 61), 'tiktorch.server.reader.eval_model_zip', 'eval_model_zip', ({(37, 41, 37, 51): 'model_file', (37, 53, 37, 60): 'devices'}, {}), '(model_file, devices)', False, 'from tiktorch.server.reader import eval_model_zip\n'), ((47, 14, 47, 26), 'uuid.uuid4', 'uuid.uuid4', ({}, {}), '()', False, 'import uuid\n'), ((36, 29, 36, 50), 'io.BytesIO', 'io.BytesIO', ({(36, 40, 36, 49): 'model_zip'}, {}), '(model_zip)', False, 'import io\n')]