from __future__ import annotations

import io
import logging
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, Union

import torch
from pydantic import BaseModel, ConfigDict

from .result import RESULT


class ResponseModel(BaseModel):

    model_config = ConfigDict(arbitrary_types_allowed=True, protected_namespaces=())

    class JobStatus(Enum):
        RECEIVED = "RECEIVED"
        APPROVED = "APPROVED"
        RUNNING = "RUNNING"
        COMPLETED = "COMPLETED"
        LOG = "LOG"
        STREAM = "STREAM"
        ERROR = "ERROR"
        NNSIGHT_ERROR = "NNSIGHT_ERROR"

    id: str
    status: JobStatus

    description: Optional[str] = ""
    data: Optional[Union[RESULT, Any]] = None
    received: Optional[datetime] = None
    session_id: Optional[str] = None

    def __str__(self) -> str:
        return f"{self.id} - {self.status.name}: {self.description}"

    def log(self, logger: logging.Logger) -> ResponseModel:
        if self.status == ResponseModel.JobStatus.ERROR:
            logger.error(str(self))
            raise SystemExit("Remote exception.")
        elif self.status == ResponseModel.JobStatus.STREAM:
            pass
        else:
            logger.info(str(self))

        return self

    def pickle(self) -> bytes:
        """Pickles self and returns bytes.

        Returns:
            bytes: Pickled ResponseModel
        """

        with io.BytesIO() as file:

            torch.save(self.model_dump(exclude_unset=True), file)

            file.seek(0)

            return file.read()

    @classmethod
    def unpickle(cls, data: bytes) -> ResponseModel:
        """Loads a ResponseModel from pickled bytes.

        Args:
            data (bytes): Pickled ResoonseModel.

        Returns:
            ResponseModel: Response.
        """

        with io.BytesIO(data) as file:
            return ResponseModel(
                **torch.load(file, map_location="cpu", weights_only=False)
            )
