import base64
import hashlib
import io
import json
import mimetypes
import os
import re
import shutil
import tempfile
import typing as t
from pathlib import Path

from binaryornot.check import is_binary
from fastapi.encoders import jsonable_encoder
from furl import furl
from PIL import Image as PILImage
from pydantic import BaseModel
from pydantic import Field as PydanticField
from pydantic.fields import ModelField, Undefined
from typing_extensions import Literal
from w3lib.url import parse_data_uri

from tungstenkit._internal import contexts
from tungstenkit._internal.utils.requests import download_file
from tungstenkit._internal.utils.uri import get_path_from_file_uri

if t.TYPE_CHECKING:
    from _typeshed import StrPath

RE_BASE64 = "^([A-Za-z0-9+/]{4})*([A-Za-z0-9+/]{3}=|[A-Za-z0-9+/]{2}==)?$"
SUPPORTED_URL_SCHEMES_FOR_FILES = ["http", "https", "data", "file"]
IMAGE_MODES_IN_PILLOW = ["RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "1", "L", "P", "I", "F"]
BUFFER_SIZE = 4 * 1024 * 1024


class URIForFile(str):
    validate_always = True

    def get_scheme(self):
        return _get_uri_scheme(self)

    def to_file_uri(self) -> "URIForFile":
        scheme = self.get_scheme()

        if scheme == "data":
            try:
                parsed_data_uri = parse_data_uri(self)
            except Exception:
                err_msg = f"Invalid data uri: '{self[:100]}'"
                if len(self) > 100:
                    err_msg += "..."
                raise ValueError(err_msg)
            mime_type = parsed_data_uri.media_type
            ext = mimetypes.guess_extension(type=mime_type)
            fd, path = tempfile.mkstemp(suffix=ext, dir=".")
            with os.fdopen(fd, mode="wb") as f:
                f.write(parsed_data_uri.data)
            return URIForFile(Path(path).resolve().as_uri())

        if scheme == "http" or scheme == "https":
            with tempfile.TemporaryDirectory() as tmp_dir:
                downloaded_path = download_file(url=self, out_path=tmp_dir)
                fd, path = tempfile.mkstemp(suffix=downloaded_path.parts[-1], dir=".")
                os.close(fd)
                shutil.move(downloaded_path, path)
            return URIForFile(Path(path).resolve().as_uri())

        return self

    def to_data_uri(self) -> "URIForFile":
        scheme = self.get_scheme()

        if scheme == "data":
            return self

        file_uri = self.to_file_uri()
        path = get_path_from_file_uri(file_uri)
        mimetype = mimetypes.guess_type(url=file_uri, strict=False)[0]
        if mimetype is None:
            mimetype = "application/octet-stream" if is_binary(str(path)) else "text/plain"
        return URIForFile(
            f"data:{mimetype};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}"
        )

    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v: t.Any):
        if not isinstance(v, str):
            raise ValueError(f"Type of uri should be 'str', not '{type(v)}'.")

        if isinstance(v, URIForFile):
            return cls(v)

        scheme = _get_uri_scheme(v)
        if not scheme:
            if (
                contexts.APP == contexts.Application.MODEL_SERVER
                or contexts.APP == contexts.Application.TASK_SERVER  # noqa: W503
            ):
                is_base64 = _check_base64(v)
                if not is_base64:
                    raise ValueError(
                        "Invalid file input. File input should be either an url or "
                        "a base64-encoded string"
                    )
                return cls._from_base64(v)
            else:
                return Path(v).resolve().as_uri()

        if scheme not in SUPPORTED_URL_SCHEMES_FOR_FILES:
            raise ValueError(
                f"Unsupported URI scheme: '{scheme}'. "
                f"Supported schemes: {', '.join(SUPPORTED_URL_SCHEMES_FOR_FILES)}"
            )
        if scheme == "file":
            parsed = furl(v)
            if parsed.netloc:
                raise ValueError("Remote files are not supported")

        return cls(v)

    @classmethod
    def _from_base64(cls, base64_str: str):
        try:
            data_uri_header = "data:application/octet-stream;base64"
            return cls(data_uri_header + "," + base64_str)
        except Exception:
            raise ValueError(f"Not a valid base64-encoded string: '{base64_str[:100]}...'")

    @classmethod
    def __modify_schema__(cls, field_schema: t.Dict[str, t.Any]) -> None:
        field_schema["format"] = "uri"


# TODO remove inheritance from BaseModel and __root__
class File(BaseModel):
    _schema_prefix: t.ClassVar[str] = "#/tungsten/"

    __root__: URIForFile

    @classmethod
    def from_path(cls, path: "StrPath"):
        return cls(__root__=URIForFile(Path(path).resolve().as_uri()))

    @property
    def path(self) -> str:
        self.__root__ = self.__root__.to_file_uri()
        return str(get_path_from_file_uri(self.__root__))

    @classmethod
    def __modify_schema__(
        cls, field_schema: t.Dict[str, t.Any], field: t.Optional[ModelField]
    ) -> None:
        if (
            contexts.APP != contexts.Application.MODEL_SERVER
            and contexts.APP != contexts.Application.TASK_SERVER
        ):
            ref = cls._schema_prefix + cls.__name__
            if field and field.field_info.description:
                field_schema["allOf"] = [{"$ref": ref}]
                field_schema["description"] = field.field_info.description
            else:
                field_schema["$ref"] = cls._schema_prefix + cls.__name__


class Image(File):
    @staticmethod
    def from_pil_image(pil_image: PILImage.Image) -> "Image":
        # TODO log warning when this is called while building
        if not isinstance(pil_image, PILImage.Image):
            raise TypeError(f"Invalid type for 'pil_image': {type(pil_image)}")
        try:
            mime_type: str = pil_image.get_format_mimetype()
            fmt = pil_image.format
        except AttributeError:
            mime_type = "image/png"
            fmt = "PNG"

        data_uri_header = f"data:{mime_type};base64"
        bytes_io = io.BytesIO()
        pil_image.save(bytes_io, format=fmt)
        image_bytes = bytes_io.getvalue()
        image_b64_str = base64.b64encode(image_bytes).decode()
        bytes_io.close()
        return Image(__root__=URIForFile(data_uri_header + "," + image_b64_str))

    def to_pil_image(
        self,
        mode: Literal[
            "RGB", "RGBA", "CMYK", "YCbCr", "LAB", "HSV", "1", "L", "P", "I", "F"
        ] = "RGB",
    ) -> PILImage.Image:
        """
        Convert to PIL Image.
        """
        if mode not in IMAGE_MODES_IN_PILLOW:
            raise ValueError("Unsupported image mode: '{}'")

        scheme = _get_uri_scheme(self.__root__)
        if scheme == "data":
            try:
                data = parse_data_uri(self.__root__).data
            except Exception:
                err_msg = f"Invalid data uri: '{self.__root__[:100]}'"
                if len(self.__root__) > 100:
                    err_msg += "..."
                raise ValueError(err_msg)
            pil_image = PILImage.open(io.BytesIO(data))
        else:
            pil_image = PILImage.open(self.path)

        return pil_image.convert(mode)

    class Config:
        schema_extra = {"example": "https://picsum.photos/200.jpg"}


class Binary(File):
    class Config:
        schema_extra = {"example": "data:text/plain;base64,aGVsbG8gd29ybGQ="}


class Video(File):
    class Config:
        schema_extra = {
            "example": "http://commondatastorage.googleapis.com/gtv-videos-bucket/sample/ForBiggerBlazes.mp4"  # noqa
        }


class Audio(File):
    class Config:
        schema_extra = {"example": "https://download.samplelib.com/mp3/sample-3s.mp3"}


class BaseIO(BaseModel):
    def _hash_for_batching(self):
        m = hashlib.sha256()
        to_be_hashsed = ""
        for field_name, field in self.__fields__.items():
            if not field.required:
                serialized_field_value = json.dumps(
                    jsonable_encoder(getattr(self, field_name))
                ).replace("\n", "")
                to_be_hashsed += field_name + ":" + serialized_field_value + "\n"

        m.update(to_be_hashsed.encode("utf-8"))
        return "sha256:" + m.hexdigest()

    class Config:
        validate_assignment = True
        validate_all = True
        arbitrary_types_allowed = True


def Field(
    *,
    description: str = None,
    choices: t.Sequence[t.Union[str, int]] = None,
    ge: float = None,
    le: float = None,
    min_length: int = None,
    max_length: int = None,
):
    extras: t.Dict[str, t.Any] = dict()
    if choices is not None:
        extras.update(choices=choices)
    return PydanticField(
        default=Undefined,
        description=description,
        ge=ge,
        le=le,
        min_length=min_length,
        max_length=max_length,
        **extras,
    )


def Option(
    default: t.Any,
    *,
    description: str = None,
    choices: t.Sequence[t.Union[str, int]] = None,
    ge: float = None,
    le: float = None,
    min_length: int = None,
    max_length: int = None,
    default_factory: t.Callable[[], t.Any] = None,
):
    extras: t.Dict[str, t.Any] = dict()
    if choices is not None:
        extras.update(choices=choices)
    return PydanticField(
        default=Undefined if default is Ellipsis else default,
        description=description,
        ge=ge,
        le=le,
        min_length=min_length,
        max_length=max_length,
        default_factory=default_factory,
        **extras,
    )


def _get_uri_scheme(uri_str: str) -> str:
    return furl(uri_str[:50]).scheme


def _check_base64(s: str) -> bool:
    return False if s is None or not re.search(RE_BASE64, s) else True
