import io
import base64
from uuid import UUID
from typing import Iterator, Iterable, List, Optional, Union, TYPE_CHECKING
from concurrent.futures import ThreadPoolExecutor
from threading import Event

from PIL import Image as ImagePIL

if TYPE_CHECKING:
    from lqs.common.facade import CoreFacade
from .utils import (
    get_relative_object_path,
    decompress_chunk_bytes,
    get_record_image,
)
from lqs.transcode import Transcode
from lqs.interface.core.models import (
    Ingestion,
    Object,
    Record,
    Topic,
)
from .record_fetcher import RecordFetcher


class RecordUtils:
    def __init__(self, app: "CoreFacade"):
        self.app = app

    def get_record_image(
        self,
        record_data: dict | bytes,
        max_size: Optional[int] = None,
        format: str = "WEBP",
        format_params: dict = {},
        renormalize: bool = False,
        reset_position: bool = True,
        return_bytes: bool = False,
        **kwargs,
    ) -> Union[ImagePIL.Image, io.BytesIO, None]:
        """
        A convenience method which takes deserialized record data from a standard image topic and returns the image as a PIL Image or BytesIO object.

        :param record_data: The record data.
        :type record_data: dict | bytes
        :param max_size: The maximum width or height to downscale to. Defaults to None, which means no downscaling.
        :type max_size: int, optional
        :param format: The output format to use. Defaults to "WEBP".
        :type format: str, optional
        :param format_params: The format parameters to use. Defaults to {}.
        :type format_params: dict, optional
        :param renormalize: Whether to renormalize the image, which is necessary for visualization in some cases. Defaults to False.
        :type renormalize: bool, optional
        :param reset_position: Whether to reset the position offset position of the BytesIO object. Defaults to True.
        :type reset_position: bool, optional
        :param return_bytes: Whether to return the image as a BytesIO object. Defaults to False.
        :type return_bytes: bool, optional
        :returns: The image, either as a PIL Image or BytesIO object, or None if the record data does not contain an image.
        :rtype: Union[ImagePIL.Image, io.BytesIO, None]
        """
        return get_record_image(
            record_data=record_data,
            max_size=max_size,
            format=format,
            format_params=format_params,
            renormalize=renormalize,
            reset_position=reset_position,
            return_bytes=return_bytes,
            **kwargs,
        )

    def iter_topic_records(
        self,
        topic: Union[Topic, UUID, str],
        start_time: Optional[int] = None,
        end_time: Optional[int] = None,
        per_request_limit: int = 1000,
        frequency: Optional[float] = None,
        include_auxiliary_data: bool = False,
    ):
        """
        Iterate over records for a topic.

        :param topic: The topic to use.
        :type topic: Topic | UUID | str
        :param start_time: The start time to use. Defaults to None.
        :type start_time: int, optional
        :param end_time: The end time to use. Defaults to None.
        :type end_time: int, optional
        :param per_request_limit: The limit to use for each request. Defaults to 1000.
        :type per_request_limit: int, optional
        :param frequency: The frequency to use for each request. Defaults to None.
        :type frequency: float, optional
        :param include_auxiliary_data: Whether to include auxiliary data. Defaults to False.
        :type include_auxiliary_data: bool, optional
        :yields: The record.
        :rtype: Record
        """
        # TODO: there's a bug when using frequency with pagination where records can be skipped
        if isinstance(topic, Topic):
            topic_id = topic.id
        else:
            topic_id = topic
        with ThreadPoolExecutor() as executor:
            records = self.app.list.record(
                topic_id=topic_id,
                timestamp_gte=start_time,
                timestamp_lte=end_time,
                limit=per_request_limit,
                frequency=frequency,
                include_auxiliary_data=include_auxiliary_data,
            ).data
            if len(records) == 0:
                return
            last_record = records[-1]
            last_record_timestamp = last_record.timestamp
            records_future = executor.submit(
                self.app.list.record,
                topic_id=topic_id,
                timestamp_gt=last_record_timestamp,
                timestamp_lte=end_time,
                limit=per_request_limit,
                frequency=frequency,
                include_auxiliary_data=include_auxiliary_data,
            )
            while len(records) > 0:
                yield from records
                records_res = records_future.result()
                records = records_res.data
                if len(records) == 0:
                    break
                last_record = records[-1]
                last_record_timestamp = last_record.timestamp
                records_future = executor.submit(
                    self.app.list.record,
                    topic_id=topic_id,
                    timestamp_gt=last_record_timestamp,
                    timestamp_lte=end_time,
                    limit=per_request_limit,
                    frequency=frequency,
                    include_auxiliary_data=include_auxiliary_data,
                )

    def iter_topics_records(
        self,
        topics: List[Union[Topic, UUID, str]],
        start_time: Optional[int] = None,
        end_time: Optional[int] = None,
        per_request_limit: int = 1000,
        frequency: Optional[float] = None,
        include_auxiliary_data: bool = False,
    ):
        """
        Iterate over records for multiple topics.

        :param topics: The topics to use.
        :type topics: List[Topic | UUID | str]
        :param start_time: The start time to use. Defaults to None.
        :type start_time: int, optional
        :param end_time: The end time to use. Defaults to None.
        :type end_time: int, optional
        :param per_request_limit: The limit to use for each request. Defaults to 1000.
        :type per_request_limit: int, optional
        :param frequency: The frequency to use for each request. Defaults to None.
        :type frequency: float, optional
        :param include_auxiliary_data: Whether to include auxiliary data. Defaults to False.
        :type include_auxiliary_data: bool, optional
        :yields: The record.
        :rtype: Record
        """
        topic_ids = []
        for topic in topics:
            if isinstance(topic, Topic):
                topic_ids.append(topic.id)
            else:
                topic_ids.append(topic)
        record_iters = {
            topic_id: self.iter_topic_records(
                topic=topic_id,
                start_time=start_time,
                end_time=end_time,
                per_request_limit=per_request_limit,
                frequency=frequency,
                include_auxiliary_data=include_auxiliary_data,
            )
            for topic_id in topic_ids
        }
        next_records = {topic_id: None for topic_id in topic_ids}
        while True:
            for topic_id, record_iter in record_iters.items():
                if next_records[topic_id] is None:
                    next_records[topic_id] = next(record_iter, None)
            next_topic_id = None
            next_record = None
            for topic_id, record in next_records.items():
                if record is not None:
                    if next_record is None or record.timestamp < next_record.timestamp:
                        next_topic_id = topic_id
                        next_record = record
            if next_record is None:
                break
            yield next_record
            next_records[next_topic_id] = next(record_iters[next_topic_id], None)

    def load_auxiliary_data_image(self, source: Union[Record, dict]):
        if isinstance(source, Record):
            auxiliary_data = source.get_auxiliary_data()
        else:
            auxiliary_data = source

        if auxiliary_data is None:
            return None
        if "image" not in auxiliary_data:
            return None
        encoded_webp_data = auxiliary_data["image"]
        decoded_webp_data = base64.b64decode(encoded_webp_data)
        image = ImagePIL.open(io.BytesIO(decoded_webp_data))
        return image

    def get_deserialized_record_data(
        self,
        record: Record,
        topic: Optional[Topic] = None,
        ingestion: Optional[Ingestion] = None,
        transcoder: Optional[Transcode] = None,
    ) -> dict:
        if transcoder is None:
            transcoder = Transcode()

        if topic is None:
            topic = self.app.fetch.topic(record.topic_id).data

        message_bytes = self.fetch_record_bytes(record=record, ingestion=ingestion)

        return transcoder.deserialize(
            type_encoding=topic.type_encoding,
            type_name=topic.type_name,
            type_data=topic.type_data,
            message_bytes=message_bytes,
        )

    def fetch_record_bytes(
        self,
        record: Record,
        ingestion: Optional[Ingestion] = None,
        decompress_chunk: bool = True,
        return_full_chunk: bool = False,
    ) -> bytes:

        if ingestion is None:
            ingestion = self.app.fetch.ingestion(record.ingestion_id).data

        object_store_id = (
            str(ingestion.object_store_id)
            if ingestion.object_store_id is not None
            else None
        )
        object_key = str(ingestion.object_key)

        if record.source is not None:
            # if the record has a source, we need to get the relative path from the object_key
            object_key = get_relative_object_path(
                object_key=object_key, source=record.source
            )

        if object_store_id is None:
            # the data is coming from a log object
            message_bytes: bytes = self.app.fetch.log_object(
                object_key=object_key,
                log_id=record.log_id,
                redirect=True,
                offset=record.data_offset,
                length=record.data_length,
            )
        else:
            # the data is coming from an object store
            message_bytes: bytes = self.app.fetch.object(
                object_key=object_key,
                object_store_id=object_store_id,
                redirect=True,
                offset=record.data_offset,
                length=record.data_length,
            )

        if record.chunk_compression is not None and record.chunk_compression not in [
            "",
            "none",
        ]:
            if decompress_chunk:
                # if the record is compressed, we need to decompress it
                message_bytes = decompress_chunk_bytes(
                    chunk_bytes=message_bytes,
                    chunk_compression=record.chunk_compression,
                    chunk_length=record.chunk_length,
                )
                if not return_full_chunk:
                    # we only return the relevant part of the chunk
                    message_bytes = message_bytes[
                        record.chunk_offset : record.chunk_offset + record.chunk_length
                    ]
            else:
                if not return_full_chunk:
                    raise Exception(
                        "Cannot return partial chunk without decompressing it."
                    )

        return message_bytes

    def get_record_set(
        self,
        records: Iterable[Record],
        carryover_record: Optional[Record] = None,
        density_threshold: float = 0.9,
        max_contiguous_size: int = 100 * 1000 * 1000,  # 100 MB
        max_contiguous_records: int = 1000,
    ) -> tuple[list[list[Record]], Optional[Record]]:
        record_set: list[Record] = []
        relevant_length = 0
        full_length = 0
        start_offset = None
        last_ingestion_id = None
        last_source = None
        last_offset = None

        if carryover_record is not None:
            record_set.append(carryover_record)
            start_offset = carryover_record.data_offset
            last_ingestion_id = carryover_record.ingestion_id
            last_source = carryover_record.source
            last_offset = carryover_record.data_offset + carryover_record.data_length
            carryover_record = None
        leftover_record: Optional[Record] = None

        for record in records:
            if start_offset is None:
                start_offset = record.data_offset
            if last_ingestion_id is None:
                last_ingestion_id = record.ingestion_id
            if last_source is None:
                last_source = record.source
            if last_offset is None:
                last_offset = record.data_offset + record.data_length

            if (record.data_offset + record.data_length) < last_offset:
                # ensure records are ordered properly by offset
                leftover_record = record
                break

            relevant_length += record.data_length
            full_length = record.data_offset + record.data_length - start_offset
            if (
                relevant_length / full_length > density_threshold
                and last_ingestion_id == record.ingestion_id
                and last_source == record.source
                and len(record_set) < max_contiguous_records
                and full_length < max_contiguous_size
                and record.data_offset + record.data_length >= last_offset
            ):
                record_set.append(record)
                last_offset = record.data_offset + record.data_length
            else:
                if len(record_set) == 0:
                    raise Exception("Record set cannot be empty.")
                leftover_record = record
                break
        return record_set, leftover_record

    def get_presigned_url(
        self,
        object_key: str,
        object_store_id: Union[UUID, str, None] = None,
        log_id: Union[UUID, str, None] = None,
        start_offset: Optional[int] = None,
        end_offset: Optional[int] = None,
    ) -> str:
        params = dict(object_key=object_key, redirect=False)
        if object_store_id is None:
            # the data is coming from a log object
            params["log_id"] = log_id
            if start_offset is not None:
                params["offset"] = start_offset
                if end_offset is not None:
                    params["length"] = end_offset - start_offset
            object_meta: Object = self.app.fetch.log_object(**params).data
        else:
            # the data is coming from an object store
            params["object_store_id"] = object_store_id
            if start_offset is not None:
                params["offset"] = start_offset
                if end_offset is not None:
                    params["length"] = end_offset - start_offset
            object_meta: Object = self.app.fetch.object(**params).data
        presigned_url = object_meta.presigned_url
        return presigned_url

    def iter_record_data(
        self,
        records: Iterable[Record],
        deserialize_results: bool = False,
        transcoder: Optional[Transcode] = None,
        density_threshold: float = 0.9,
        max_contiguous_size: int = 100 * 1000 * 1000,
        max_contiguous_records: int = 1000,
        max_workers: Optional[int] = 2,
        order_by_timestamp: bool = True,
        stop_event: Optional[Event] = None,
        use_cache: bool = False,
        cache_dir: str = "/tmp/lqs",
        read_from_cache: bool = True,
        write_to_cache: bool = True,
    ) -> Iterator[tuple[Record, Union[bytes, dict]]]:
        """
        Given a set of records, yield the record and its data.

        :param records: The records to use.
        :type records: Iterable[Record]
        :param deserialize_results: Whether to deserialize the results. Defaults to False.
        :type deserialize_results: bool, optional
        :param transcoder: The transcoder to use. Defaults to None.
        :type transcoder: Transcode, optional
        :param density_threshold: The density threshold to use. Defaults to 0.9.
        :type density_threshold: float, optional
        :param max_contiguous_size: The maximum contiguous size to use. Defaults to 100 * 1000 * 1000.
        :type max_contiguous_size: int, optional
        :param max_contiguous_records: The maximum contiguous records to use. Defaults to 1000.
        :type max_contiguous_records: int, optional
        :param max_workers: The maximum number of workers to use. Defaults to 2.
        :type max_workers: int | None, optional
        :param order_by_timestamp: Whether to order the records by timestamp. Defaults to True.
        :type order_by_timestamp: bool, optional
        :param stop_event: An event to signal stopping the iteration. Defaults to None.
        :type stop_event: Event, optional
        :yields: The record and the record data.
        :rtype: tuple[Record, dict | bytes]
        :param use_cache: Whether to use caching for record data. Defaults to False.
        :type use_cache: bool, optional
        :param cache_dir: The directory to use for caching. Defaults to "/tmp/lqs".
        :type cache_dir: str, optional
        :param read_from_cache: Whether to read from cache if available. Defaults to True.
        :type read_from_cache: bool, optional
        :param write_to_cache: Whether to write to cache after fetching. Defaults to True.
        :type write_to_cache: bool, optional
        """
        fetcher = RecordFetcher(
            app=self.app,
            max_workers=max_workers or 2,
        )
        yield from fetcher.iter_record_data(
            records=records,
            deserialize_results=deserialize_results,
            transcoder=transcoder,
            density_threshold=density_threshold,
            max_contiguous_size=max_contiguous_size,
            max_contiguous_records=max_contiguous_records,
            order_by_timestamp=order_by_timestamp,
            stop_event=stop_event,
            use_cache=use_cache,
            cache_dir=cache_dir,
            read_from_cache=read_from_cache,
            write_to_cache=write_to_cache,
        )
