import io
import logging

import grpc
from asgiref.sync import sync_to_async
from celery.result import AsyncResult
from django_socio_grpc import mixins
from django_socio_grpc.decorators import grpc_action

logger = logging.getLogger(__name__)

from django_socio_grpc.grpc_actions.actions import GRPCActionMixin
from django_socio_grpc.grpc_actions.placeholders import (
    FnPlaceholder,
    LookupField,
    SelfSerializer,
    StrTemplatePlaceholder,
)
from django_socio_grpc.grpc_actions.utils import get_serializer_base_name
from django_socio_grpc.utils.constants import DEFAULT_LIST_FIELD_NAME, PARTIAL_UPDATE_FIELD_NAME, REQUEST_SUFFIX


from ..tasks import extract_metadata_from_parquet_file


class AsyncUploadModelMixin(GRPCActionMixin):
    @grpc_action(
        request=[
            {"name": "filename", "type": "string"},
            {"name": "update_item_id", "type": "string"},
            {"name": "data", "type": "bytes"},
        ],
        request_name="FileUploadChunk",
        response=[{"name": "success", "type": "bool"}],
        response_name="FileUploadInfo",
        request_stream=True,
    )
    async def UploadFile(self, request, context):
        # logger.info(f"File upload started:")

        filename = ""
        try:
            with io.BytesIO() as f:
                async for file_chunk in request:
                    filename = file_chunk.filename
                    # logger.info("Reading file chunk:", len(file_chunk.data))
                    f.write(file_chunk.data)

                queryset = self.get_queryset()

                item = await queryset.filter(pk=file_chunk.update_item_id).aget()
                # logger.info("--- data found:", data.name_display)
                res = await sync_to_async(item.file.save)(filename, f, save=True)

                f.seek(0)
                file_content = f.getvalue()
                if filename.endswith(".parquet"):
                    task = extract_metadata_from_parquet_file.delay(parquet_file_str=file_content)
                    #logger.info(f"+++ mixin - metadata task: {metadata}")
                    # task_id = task.id
                    # item.metadata = AsyncResult(task_id)  #metadata.collect()
                    # await item.asave()

            logger.info(f"File upload successful: {len(file_content)} bytes")
            return self.pb2.FileUploadInfo(success=True)

        except Exception as e:
            logger.exception(f"Document ({filename}) upload has failed {e}")
            return self.pb2.FileUploadInfo(success=False)

    # image upload
    @grpc_action(
        request=[
            {"name": "filename", "type": "string"},
            {"name": "update_item_id", "type": "string"},
            {"name": "data", "type": "bytes"},
        ],
        request_name="ImageUploadChunk",
        response=[{"name": "success", "type": "bool"}],
        response_name="ImageUploadInfo",
        request_stream=True,
    )
    async def UploadImage(self, request, context):
        logger.info("File upload started:")

        try:
            with io.BytesIO() as f:
                async for file_chunk in request:
                    filename = file_chunk.filename
                    #logger.info("Reading file chunk:", len(file_chunk.data))
                    f.write(file_chunk.data)

                queryset = self.get_queryset()

                item = await queryset.filter(pk=file_chunk.update_item_id).aget()
                #logger.info("--- upload item found:", file_chunk.update_item_id, item.name_display)
                res = await sync_to_async(item.image.save)(filename, f, save=True)

                f.seek(0)
                file_content = f.getvalue()

            logger.info(f"Image File upload successful: {len(file_content)} bytes")
            return self.pb2.ImageUploadInfo(success=True)

        except Exception as e:
            logger.exception(f"Image File ({filename}) upload has failed {e}")
            return self.pb2.ImageUploadInfo(success=True)


class AsyncRetrieveByNameModelMixin(GRPCActionMixin):
    @grpc_action(
        #request=LookupField,
        request_name=StrTemplatePlaceholder(
            f"{{}}RetrieveByNameNS{REQUEST_SUFFIX}", get_serializer_base_name
        ),
        request=[
            {"name": "name", "type": "string"},
            {"name": "namespace_uri", "type": "string"},
        ],
        response=SelfSerializer,
    )
    async def RetrieveByNameNamespace(self, request, context):
        """
        Retrieve a model instance.

        The request have to include a field corresponding to
        ``lookup_request_field``.  If an object can be retrieved this returns
        a proto message of ``serializer.Meta.proto_class``.
        """
        queryset = self.get_queryset()
        queryset = await self.afilter_queryset(queryset)

        if request.namespace_uri is None or request.namespace_uri == "":
            queryset = await sync_to_async(queryset.filter)(name=request.name)
        else:
            queryset = await sync_to_async(queryset.filter)(namespace__uri=request.namespace_uri, name=request.name)

        instance = await queryset.afirst()

        serializer = await self.aget_serializer(instance)
        return await serializer.amessage

    @staticmethod
    def get_default_method(model_name):
        return {
            "Retrieve": {
                "request": {"is_stream": False, "message": f"{model_name}RetrieveRequest"},
                "response": {"is_stream": False, "message": model_name},
            },
        }

    @staticmethod
    def get_default_message(model_name, fields="__pk__"):
        return {
            f"{model_name}RetrieveRequest": fields,
        }

class AsyncDownloadModelMixin(mixins.GRPCActionMixin):
    @grpc_action(
        request=[{"name": "item_id", "type": "string"}, {"name": "chunk_size", "type": "uint64"}],
        request_name="FileDownloadInfo",
        response=[
            {"name": "filename", "type": "string"},
            {"name": "data", "type": "bytes"},
            {"name": "success", "type": "bool"},
        ],
        response_name="FileDownloadChunk",
        response_stream=True,
    )
    async def DownloadFile(self, request, context):
        CHUNK_SIZE = 1024 * 1024  # 1 MB

        logger.info("File download started:")

        filename = ""

        queryset = self.get_queryset()
        instance = await queryset.filter(pk=request.item_id).aget()
        file_field = instance.file
        filename = file_field.name

        try:
            with file_field.open("rb") as f:
                while True:
                    chunk = f.read(CHUNK_SIZE)
                    if not chunk:
                        break
                    yield self.pb2.FileDownloadChunk(filename=filename, data=chunk, success=True)

            logger.info(f"File download successful: {filename}")
            #return self.pb2.FileDownloadInfo(success=True)
        except self.model.DoesNotExist:
            context.abort(grpc.StatusCode.NOT_FOUND, "File not found.")
        except Exception as e:
            logger.exception(f"Document ({filename}) upload has failed {e}")
            context.abort(grpc.StatusCode.INTERNAL, str(e))
            yield self.pb2.FileDownloadInfo(success=False)


    @grpc_action(
        request=[{"name": "item_id", "type": "string"},
                 {"name": "chunk_size", "type": "uint64"}],
        request_name="ImageDownloadInfo",
        response=[
            {"name": "filename_image", "type": "string"},
            {"name": "data", "type": "bytes"},
            {"name": "success", "type": "bool"},
        ],
        response_name="ImageDownloadChunk",
        response_stream=True,
    )
    async def DownloadImage(self, request, context):
        CHUNK_SIZE = 1024 * 1024  # 1 MB

        logger.info("File download started:")

        filename = ""

        queryset = self.get_queryset()

        instance = await queryset.filter(pk=request.item_id).aget()
        file_field = instance.image
        filename = file_field.name

        try:
            with file_field.open("rb") as f:
                while True:
                    chunk = f.read(CHUNK_SIZE)
                    if not chunk:
                        break
                    yield self.pb2.ImageDownloadChunk(filename_image=filename, data=chunk, success=True)
            
            logger.info(f"Image download successful: {filename}")
            # return self.pb2.ImageDownloadInfo(success=True)
        except self.model.DoesNotExist:
            context.abort(grpc.StatusCode.NOT_FOUND, "Image not found.")
        except Exception as e:
            logger.exception(f"Image ({filename}) download has failed {e}")
            context.abort(grpc.StatusCode.INTERNAL, str(e))
            # return self.pb2.ImageDownloadInfo(success=False)
            
