from __future__ import annotations

import io
import sys
import time
from datetime import datetime
from typing import Any, Dict, Optional, Tuple

import msgspec
import requests
import socketio
import torch
from tqdm import tqdm

from ... import __IPYTHON__, CONFIG, remote_logger
from ...schema.request import RequestModel, StreamValueModel
from ...schema.response import ResponseModel
from ...schema.result import RESULT, ResultModel
from ...tracing.backends import Backend
from ...tracing.graph import Graph
from ...util import NNsightError
from ..contexts.local import LocalContext, RemoteContext


class RemoteBackend(Backend):
    """Backend to execute a context object via a remote service.

    Context object must inherit from RemoteMixin and implement its methods.

    Attributes:

        url (str): Remote host url. Defaults to that set in CONFIG.API.HOST.
    """

    def __init__(
        self,
        model_key: str,
        host: str = None,
        blocking: bool = True,
        job_id: str = None,
        ssl: bool = None,
        api_key: str = "",
    ) -> None:

        self.model_key = model_key

        self.job_id = job_id or CONFIG.API.JOB_ID
        self.ssl = CONFIG.API.SSL if ssl is None else ssl
        self.zlib = CONFIG.API.ZLIB
        self.format = CONFIG.API.FORMAT
        self.api_key = api_key or CONFIG.API.APIKEY
        self.blocking = blocking

        self.host = host or CONFIG.API.HOST
        self.address = f"http{'s' if self.ssl else ''}://{self.host}"
        self.ws_address = f"ws{'s' if CONFIG.API.SSL else ''}://{self.host}"

    def request(self, graph: Graph) -> Tuple[bytes, Dict[str, str]]:

        data = RequestModel.serialize(graph, self.format, self.zlib)

        headers = {
            "model_key": self.model_key,
            "format": self.format,
            "zlib": str(self.zlib),
            "ndif-api-key": self.api_key,
            "sent-timestamp": str(time.time()),
        }

        return data, headers

    def __call__(self, graph: Graph):

        if self.blocking:

            # Do blocking request.
            result = self.blocking_request(graph)

        else:

            # Otherwise we are getting the status / result of the existing job.
            result = self.non_blocking_request(graph)

        if result is not None:
            ResultModel.inject(graph, result)

    def handle_response(
        self, response: ResponseModel, graph: Optional[Graph] = None
    ) -> Optional[RESULT]:
        """Handles incoming response data.

        Logs the response object.
        If the job is completed, retrieve and stream the result from the remote endpoint.
        Use torch.load to decode and load the `ResultModel` into memory.
        Use the backend object's .handle_result method to handle the decoded result.

        Args:
            response (Any): Json data to concert to `ResponseModel`

        Raises:
            Exception: If the job's status is `ResponseModel.JobStatus.ERROR`

        Returns:
            ResponseModel: ResponseModel.
        """

        # Log response for user
        response.log(remote_logger)

        # If job is completed:
        if response.status == ResponseModel.JobStatus.COMPLETED:

            # If the response has no result data, it was too big and we need to stream it from the server.
            if response.data is None:

                result = self.get_result(response.id)
            else:

                result = response.data

            return result

        # If were receiving a streamed value:
        elif response.status == ResponseModel.JobStatus.STREAM:

            # Second item is index of LocalContext node.
            # First item is the streamed value from the remote service.

            index, dependencies = response.data

            ResultModel.inject(graph, dependencies)

            node = graph.nodes[index]

            node.execute()

        elif response.status == ResponseModel.JobStatus.NNSIGHT_ERROR:
            if graph.debug:
                error_node = graph.nodes[response.data["node_id"]]
                try:
                    raise NNsightError(
                        response.data["err_message"],
                        error_node.index,
                        response.data["traceback"],
                    )
                except NNsightError as nns_err:
                    if (
                        __IPYTHON__
                    ):  # in IPython the traceback content is rendered by the Error itself
                        # add the error node traceback to the the error's traceback
                        nns_err.traceback_content += "\nDuring handling of the above exception, another exception occurred:\n\n"
                        nns_err.traceback_content += error_node.meta_data["traceback"]
                    else:  # else we print the traceback manually
                        print(f"\n{response.data['traceback']}")
                        print(
                            "During handling of the above exception, another exception occurred:\n"
                        )
                        print(f"{error_node.meta_data['traceback']}")

                    sys.tracebacklimit = 0
                    raise nns_err from None
                finally:
                    if __IPYTHON__:
                        sys.tracebacklimit = None
            else:
                print(f"\n{response.data['traceback']}")
                raise SystemExit("Remote exception.")

    def submit_request(
        self, data: bytes, headers: Dict[str, Any]
    ) -> Optional[ResponseModel]:
        """Sends request to the remote endpoint and handles the response object.

        Raises:
            Exception: If there was a status code other than 200 for the response.

        Returns:
            (ResponseModel): Response.
        """

        from ...schema.response import ResponseModel

        headers["Content-Type"] = "application/octet-stream"

        response = requests.post(
            f"{self.address}/request",
            data=data,
            headers=headers,
        )

        if response.status_code == 200:

            response = ResponseModel(**response.json())

            self.handle_response(response)

            return response

        else:
            msg = response.reason
            raise ConnectionError(msg)

    def get_response(self) -> Optional[RESULT]:
        """Retrieves and handles the response object from the remote endpoint.

        Raises:
            Exception: If there was a status code other than 200 for the response.

        Returns:
            (ResponseModel): Response.
        """

        from ...schema.response import ResponseModel

        response = requests.get(
            f"{self.address}/response/{self.job_id}",
            headers={"ndif-api-key": self.api_key},
        )

        if response.status_code == 200:

            response = ResponseModel(**response.json())

            return self.handle_response(response)

        else:

            raise Exception(response.reason)

    def get_result(self, id: str) -> RESULT:

        result_bytes = io.BytesIO()
        result_bytes.seek(0)

        # Get result from result url using job id.
        with requests.get(
            url=f"{self.address}/result/{id}",
            stream=True,
        ) as stream:
            # Total size of incoming data.
            total_size = float(stream.headers["Content-length"])

            with tqdm(
                total=total_size,
                unit="B",
                unit_scale=True,
                desc="Downloading result",
            ) as progress_bar:
                # chunk_size=None so server determines chunk size.
                for data in stream.iter_content(chunk_size=None):
                    progress_bar.update(len(data))
                    result_bytes.write(data)

        # Move cursor to beginning of bytes.
        result_bytes.seek(0)

        # Decode bytes with pickle and then into pydantic object.
        result = torch.load(result_bytes, map_location="cpu", weights_only=False)

        result = ResultModel(**result).result

        # Close bytes
        result_bytes.close()

        return result

    def blocking_request(self, graph: Graph) -> Optional[RESULT]:
        """Send intervention request to the remote service while waiting for updates via websocket.

        Args:
            request (RequestModel):Request.
        """

        # We need to do some processing / optimizations on both the graph were sending remotely
        # and our local intervention graph. In order handle the more complex Protocols for streaming.

        # Create a socketio connection to the server.
        with socketio.SimpleClient(reconnection_attempts=10) as sio:
            # Connect
            sio.connect(
                self.ws_address,
                socketio_path="/ws/socket.io",
                transports=["websocket"],
                wait_timeout=10,
            )

            remote_graph = preprocess(graph)

            data, headers = self.request(remote_graph)

            headers["session_id"] = sio.sid

            # Submit request via
            response = self.submit_request(data, headers)

            LocalContext.set(
                lambda *args: self.stream_send(*args, job_id=response.id, sio=sio)
            )

            try:
                # Loop until
                while True:

                    # Get pickled bytes value from the websocket.
                    response = sio.receive()[1]
                    # Convert to pydantic object.
                    response = ResponseModel.unpickle(response)
                    # Handle the response.
                    result = self.handle_response(response, graph=graph)
                    # Break when completed.
                    if result is not None:
                        return result

            except Exception as e:

                raise e

            finally:
                LocalContext.set(None)

    def stream_send(
        self, values: Dict[int, Any], job_id: str, sio: socketio.SimpleClient
    ):
        """Upload some value to the remote service for some job id.

        Args:
            value (Any): Value to upload
            job_id (str): Job id.
            sio (socketio.SimpleClient): Connected websocket client.
        """

        sio.emit(
            "stream_upload",
            data=(StreamValueModel.serialize(values, self.format, self.zlib), job_id),
        )

    def non_blocking_request(self, graph: Graph):
        """Send intervention request to the remote service if request provided. Otherwise get job status.

        Sets CONFIG.API.JOB_ID on initial request as to later get the status of said job.

        When job is completed, clear CONFIG.API.JOB_ID to request a new job.

        Args:
            request (RequestModel): Request if submitting a new request. Defaults to None
        """

        if self.job_id is None:

            data, headers = self.request(graph)

            # Submit request via
            response = self.submit_request(data, headers)

            CONFIG.API.JOB_ID = response.id

            CONFIG.save()

        else:

            try:

                result = self.get_response()

                if result is not None:

                    CONFIG.API.JOB_ID = None

                    CONFIG.save()

                    return result

            except Exception as e:

                CONFIG.API.JOB_ID = None

                CONFIG.save()

                raise e


def preprocess(graph: Graph):

    new_graph = graph.copy()

    for node in new_graph.nodes:

        if node.target is LocalContext:

            graph.nodes[node.index].kwargs["uploads"] = RemoteContext.from_local(node)

    return new_graph
