from __future__ import annotations

import io
import logging
import os
import typing as t
from functools import lru_cache

from multipart.multipart import parse_options_header
from starlette.datastructures import UploadFile
from starlette.requests import Request
from starlette.responses import Response

from ...exceptions import BadInput
from ...exceptions import BentoMLException
from ...exceptions import InvalidArgument
from ...grpc.utils import import_generated_stubs
from ..service.openapi import SUCCESS_DESCRIPTION
from ..service.openapi.specification import MediaType
from ..service.openapi.specification import Schema
from ..types import FileLike
from ..utils.filesystem import resolve_user_filepath
from ..utils.http import set_cookies
from .base import IODescriptor

logger = logging.getLogger(__name__)

if t.TYPE_CHECKING:
    from bentoml.grpc.v1 import service_pb2 as pb
    from bentoml.grpc.v1alpha1 import service_pb2 as pb_v1alpha1

    from ..context import ServiceContext as Context
    from .base import OpenAPIResponse

    FileKind: t.TypeAlias = t.Literal["binaryio", "textio"]
else:
    pb, _ = import_generated_stubs("v1")
    pb_v1alpha1, _ = import_generated_stubs("v1alpha1")

FileType = t.Union[io.IOBase, t.IO[bytes], FileLike[bytes]]


class File(
    IODescriptor[FileType], descriptor_id="bentoml.io.File", proto_fields=("file",)
):
    """
    :obj:`File` defines API specification for the inputs/outputs of a Service, where either
    inputs will be converted to or outputs will be converted from file-like objects as
    specified in your API function signature.

    A sample ViT service:

    .. code-block:: python
       :caption: `service.py`

       from __future__ import annotations

       import io
       from typing import TYPE_CHECKING
       from typing import Any

       import bentoml
       from bentoml.io import File

       if TYPE_CHECKING:
           from numpy.typing import NDArray

       runner = bentoml.tensorflow.get('image-classification:latest').to_runner()

       svc = bentoml.Service("vit-pdf-classifier", runners=[runner])

       @svc.api(input=File(), output=NumpyNdarray(dtype="float32"))
       async def predict(input_pdf: io.BytesIO[Any]) -> NDArray[Any]:
           return await runner.async_run(input_pdf)

    Users then can then serve this service with :code:`bentoml serve`:

    .. code-block:: bash

        % bentoml serve ./service.py:svc --reload

    Users can then send requests to the newly started services with any client:

    .. tab-set::

        .. tab-item:: Bash

            .. code-block:: bash

               % curl -H "Content-Type: multipart/form-data" \\
                      -F 'fileobj=@test.pdf;type=application/pdf' \\
                      http://0.0.0.0:3000/predict

        .. tab-item:: Python

           .. code-block:: python
              :caption: `request.py`

              import requests

              requests.post(
                  "http://0.0.0.0:3000/predict",
                  files = {"upload_file": open('test.pdf', 'rb')},
                  headers = {"content-type": "multipart/form-data"}
              ).text

    Args:
        kind: The kind of file-like object to be used. Currently, the only accepted value is ``binaryio``.
        mime_type: Return MIME type of the :code:`starlette.response.Response`, only available when used as output descriptor

    Returns:
        :obj:`File`: IO Descriptor that represents file-like objects.

    """

    def __new__(
        cls, kind: FileKind = "binaryio", mime_type: str | None = None, **kwargs: t.Any
    ) -> File:
        if kind == "binaryio":
            res = super().__new__(BytesIOFile, **kwargs)
        else:
            raise ValueError(f"invalid File kind '{kind}'")
        res._mime_type = mime_type
        return res

    def _from_sample(self, sample: FileType | str) -> FileType:
        """
        Create a :class:`~bentoml._internal.io_descriptors.file.File` IO Descriptor from given inputs.

        Args:
            sample: Given File-like object, or a path to a file.
            kind: The kind of file-like object to be used. Currently, the only accepted value is ``binaryio``.
            mime_type: Optional MIME type for the descriptor. If not provided, ``from_sample``
                       will try to infer the MIME type from the file extension.

        Returns:
            :class:`~bentoml._internal.io_descriptors.file.File`: IODescriptor from given users inputs.

        Example:

        .. code-block:: python
           :caption: `service.py`

           from __future__ import annotations

           import bentoml
           from typing import Any
           from bentoml.io import File

           input_spec = File.from_sample("/path/to/file.pdf")
           @svc.api(input=input_spec, output=File())
           async def predict(input: t.IO[t.Any]) -> t.IO[t.Any]:
               return await runner.async_run(input)

        Raises:
            :class:`MissingDependencyException`: If ``filetype`` is not installed.
            ValueError: If the MIME type cannot be inferred automatically.
            :class:`BadInput`: Any other errors that may occur during the process of figure out the MIME type.
        """
        import mimetypes

        if isinstance(sample, t.IO):
            sample = FileLike[bytes](sample, "<sample>")
        elif isinstance(sample, (str, os.PathLike)):
            p = resolve_user_filepath(sample, ctx=None)
            mime = mimetypes.guess_type(p)[0]
            self._mime_type = mime
            with open(p, "rb") as f:
                sample = FileLike[bytes](f, "<sample>")
        return sample

    @classmethod
    def from_spec(cls, spec: dict[str, t.Any]) -> t.Self:
        if "args" not in spec:
            raise InvalidArgument(f"Missing args key in File spec: {spec}")
        return cls(**spec["args"])

    def input_type(self) -> t.Type[t.Any]:
        return FileLike[bytes]

    def openapi_schema(self) -> Schema:
        return Schema(type="string", format="binary")

    def openapi_components(self) -> dict[str, t.Any] | None:
        pass

    def openapi_example(self):
        pass

    def openapi_request_body(self) -> dict[str, t.Any]:
        return {
            "content": {
                "*/*" if self._mime_type is None else self._mime_type: MediaType(
                    schema=self.openapi_schema()
                )
            },
            "required": True,
            "x-bentoml-io-descriptor": self.to_spec(),
        }

    def openapi_responses(self) -> OpenAPIResponse:
        return {
            "description": SUCCESS_DESCRIPTION,
            "content": {
                "*/*" if self._mime_type is None else self._mime_type: MediaType(
                    schema=self.openapi_schema()
                )
            },
            "x-bentoml-io-descriptor": self.to_spec(),
        }

    async def to_http_response(self, obj: FileType, ctx: Context | None = None):
        if isinstance(obj, bytes):
            body = obj
        else:
            body = obj.read()

        if ctx is not None:
            res = Response(
                body,
                headers=ctx.response.metadata,  # type: ignore (bad starlette types)
                status_code=ctx.response.status_code,
            )
            set_cookies(res, ctx.response.cookies)
        else:
            res = Response(
                body,
                headers={
                    "content-type": (
                        self._mime_type
                        if self._mime_type
                        else "application/octet-stream"
                    )
                },
            )
        return res

    async def to_proto(self, obj: FileType) -> pb.File:
        if isinstance(obj, bytes):
            body = obj
        else:
            body = obj.read()

        return pb.File(
            kind=(
                "appliction/octet-stream"
                if self._mime_type is None
                else self._mime_type
            ),
            content=body,
        )

    async def to_proto_v1alpha1(self, obj: FileType) -> pb_v1alpha1.File:
        if isinstance(obj, bytes):
            body = obj
        else:
            body = obj.read()

        try:
            kind = mimetype_to_filetype_pb_map()[
                (
                    "application/octet-stream"
                    if self._mime_type is None
                    else self._mime_type
                )
            ]
        except KeyError:
            raise BadInput(
                f"{self._mime_type} doesn't have a corresponding File 'kind'"
            ) from None

        return pb_v1alpha1.File(kind=kind, content=body)

    async def from_proto(self, field: pb.File | bytes) -> FileLike[bytes]:
        raise NotImplementedError

    async def from_http_request(self, request: Request) -> FileLike[bytes]:
        raise NotImplementedError


class BytesIOFile(File, descriptor_id=None):
    def to_spec(self) -> dict[str, t.Any]:
        return {
            "id": super().descriptor_id,
            "args": {
                "kind": "binaryio",
                "mime_type": self._mime_type,
            },
        }

    async def from_http_request(self, request: Request) -> FileLike[bytes]:
        content_type, _ = parse_options_header(request.headers["content-type"])
        if content_type.decode("utf-8") == "multipart/form-data":
            form = await request.form()
            found_mimes: t.List[str] = []
            val: t.Union[str, UploadFile]
            for val in form.values():  # type: ignore
                if isinstance(val, UploadFile):
                    found_mimes.append(val.content_type)  # type: ignore (bad starlette types)
                    if self._mime_type is None or val.content_type == self._mime_type:  # type: ignore (bad starlette types)
                        res = FileLike[bytes](val.file, val.filename)  # type: ignore (bad starlette types)
                        break
            else:
                if len(found_mimes) == 0:
                    raise BentoMLException("no File found in multipart form")
                else:
                    raise BentoMLException(
                        f"The File IO descriptor requires input to be of Content-Type '{self._mime_type}', got files with content types {', '.join(found_mimes)}"
                    )
            return res  # type: ignore
        if self.mime_type is None or content_type.decode("utf-8") == self._mime_type:
            body = await request.body()
            return FileLike[bytes](io.BytesIO(body), "<request body>")
        raise BentoMLException(
            f"File should have Content-Type '{self._mime_type}' or 'multipart/form-data', got {content_type} instead"
        )

    async def from_proto(
        self, field: pb.File | pb_v1alpha1.File | bytes
    ) -> FileLike[bytes]:
        # check if the request message has the correct field
        if isinstance(field, bytes):
            content = field
        elif isinstance(field, pb_v1alpha1.File):
            mapping = filetype_pb_to_mimetype_map()
            if field.kind:
                try:
                    mime_type = mapping[field.kind]
                    if self._mime_type is not None and mime_type != self._mime_type:
                        raise BadInput(
                            f"Inferred mime_type from 'kind' is '{mime_type}', while '{self!r}' is expecting '{self._mime_type}'",
                        )
                except KeyError:
                    raise BadInput(
                        f"{field.kind} is not a valid File kind. Accepted file kind: {[names for names, _ in pb_v1alpha1.File.FileType.items()]}",
                    ) from None
            content = field.content
            if not content:
                raise BadInput("Content is empty!") from None
            return FileLike[bytes](io.BytesIO(content), "<content>")
        else:
            assert isinstance(field, pb.File)
            if (
                field.kind is not None
                and self._mime_type is not None
                and field.kind != self._mime_type
            ):
                raise BadInput(
                    f"MIME type from 'kind' is '{field.kind}', while '{self!r}' is expecting '{self._mime_type}'",
                )
            content = field.content
            if not content:
                raise BadInput("Content is empty!") from None

        return FileLike[bytes](io.BytesIO(content), "<content>")


# v1alpha1 backward compatibility
@lru_cache(maxsize=1)
def filetype_pb_to_mimetype_map() -> dict[pb_v1alpha1.File.FileType.ValueType, str]:
    return {
        pb_v1alpha1.File.FILE_TYPE_CSV: "text/csv",
        pb_v1alpha1.File.FILE_TYPE_PLAINTEXT: "text/plain",
        pb_v1alpha1.File.FILE_TYPE_JSON: "application/json",
        pb_v1alpha1.File.FILE_TYPE_BYTES: "application/octet-stream",
        pb_v1alpha1.File.FILE_TYPE_PDF: "application/pdf",
        pb_v1alpha1.File.FILE_TYPE_PNG: "image/png",
        pb_v1alpha1.File.FILE_TYPE_JPEG: "image/jpeg",
        pb_v1alpha1.File.FILE_TYPE_GIF: "image/gif",
        pb_v1alpha1.File.FILE_TYPE_TIFF: "image/tiff",
        pb_v1alpha1.File.FILE_TYPE_BMP: "image/bmp",
        pb_v1alpha1.File.FILE_TYPE_WEBP: "image/webp",
        pb_v1alpha1.File.FILE_TYPE_SVG: "image/svg+xml",
    }


@lru_cache(maxsize=1)
def mimetype_to_filetype_pb_map() -> dict[str, pb_v1alpha1.File.FileType.ValueType]:
    return {v: k for k, v in filetype_pb_to_mimetype_map().items()}
