import json
from datetime import datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Generator, Optional
from urllib.parse import urlparse

from highlighter.agent.capabilities._buffered_video_reader import BufferedVideoCapture
from highlighter.core.logging import get_default_logger

try:
    import cv2
except ModuleNotFoundError as _:
    cv2 = None

import numpy as np
from PIL import Image

from highlighter.client import download_bytes
from highlighter.client.base_models import DataFile
from highlighter.client.io import (
    _pil_open_image_bytes,
    _pil_open_image_path,
    _pil_open_image_url,
)
from highlighter.core.exceptions import require_package

from .base_capability import DataSourceCapability, DataSourceType, StreamEvent

__all__ = [
    "ImageDataSource",
    "TextDataSource",
    "JsonArrayDataSource",
    "VideoDataSource",
]


def _iter_buffer(
    buffer,
    sep: bytes,
) -> Generator[bytes, None, None]:
    """
    Streams data from stdin until the Nth occurrence of a separator byte sequence.

    sep: The separator byte sequence (e.g., b'\n' or b'\x00\x01')
    """
    if sep is None or len(sep) == 0:
        raise ValueError("Separator must be a non-empty bytes object")

    byte_result = bytearray()
    buffer_remainder = bytearray()

    while True:
        chunk = buffer.read(4096)  # Adjusted for testing, increase for performance
        if not chunk:
            if byte_result:
                byte_result = bytes(byte_result)
                yield byte_result
            break

        buffer_remainder.extend(chunk)

        while True:
            sep_index = buffer_remainder.find(sep)
            if sep_index == -1:
                break

            byte_result.extend(buffer_remainder[:sep_index])
            yield bytes(byte_result)

            byte_result.clear()
            buffer_remainder = buffer_remainder[sep_index + len(sep) :]

        byte_result.extend(buffer_remainder)
        buffer_remainder.clear()


class TextFrameIterator:

    def __init__(
        self,
        source_buffers=None,
        source_urls=None,
        logger=None,
    ):

        self.logger = logger if logger is not None else get_default_logger("TextFrameIterator")
        self.source_buffers = source_buffers
        self.source_urls = source_urls

        if self.source_urls is not None:
            self.logger.info("TextFrameIterator using source_url")
            self._source = iter([str(u) for u, _ in self.source_urls])
        elif self.source_buffers is not None:

            def iter_buffers(bufs):
                for b in bufs:
                    for item in _iter_buffer(b, "===END===".encode("utf-8")):
                        yield item.decode()

            self._source = iter_buffers(self.source_buffers)
            self.logger.info(f"TextFrameIterator using source_buffer")
        else:
            raise ValueError("Must provide source_buffer &| source_url")

        self.frame_index = 0

    def _is_url(self, p):
        return all([urlparse(p), urlparse(p).netloc])

    def _is_local_path(self, p):
        return Path(p).exists()

    def _read_text(self, text_src):
        if self._is_local_path(text_src):
            with open(text_src, "r") as f:
                text = f.read()
            original_source_url = text_src
        elif self._is_url(text_src):
            text = download_bytes(text_src).decode("utf-8")
            original_source_url = text_src
        else:
            text = text_src
            original_source_url = None
        return text, original_source_url

    def __iter__(self):
        return self

    def __next__(self):
        text_src = next(self._source)
        text_content, original_source_url = self._read_text(text_src)

        data_file = DataFile(
            # file_id=self.ds.id,
            content=text_content,
            content_type="text",
            media_frame_index=0,
            original_source_url=original_source_url,
        )
        return data_file


class TextDataSource(DataSourceCapability):
    """

    TODO: Check/update this

    Example:
        # process a single string
        hl agent run --data-source TextDataSource PIPELINE.json "tell me a joke."

        # process many text files
        ToDo

        # Read from stdin
        cat file | hl agent run --data-source TextDataSource -sp read_stdin=true PIPELINE.json
    """

    stream_media_type = "text"

    class DefaultStreamParameters(DataSourceCapability.DefaultStreamParameters):
        byte_encoding: Optional[str] = "utf-8"

    @property
    def byte_encoding(self) -> str:
        value, _ = self._get_parameter("byte_encoding")
        return value

    def get_text_frame_generator(self, stream):

        if "source_paths_generator" in stream.variables:
            for frame_data in TextFrameIterator(
                source_buffers=None,
                source_urls=stream.variables["source_paths_generator"],
                logger=self.logger,
            ):
                yield frame_data
        elif "source_buffers" in stream.variables:
            for frame_data in TextFrameIterator(
                source_buffers=stream.variables["source_buffers"],
                source_urls=None,
                logger=self.logger,
            ):
                yield frame_data
        elif "records" in stream.variables:
            raise NotImplementedError("#### ToDo records ####")
        else:
            raise NotImplementedError("#### ToDo ####")

    def get_next_frame_data(self, stream):

        frame_generator = stream.variables.get("text_frame_generator", None)
        if frame_generator is None:
            frame_generator = self.get_text_frame_generator(stream)
            stream.variables["text_frame_generator"] = frame_generator

        data_file = next(frame_generator)
        entities = {}  # ToDo, how/where to load existing entities from
        return data_file, entities


class JsonArrayFrameIterator:

    def __init__(self, source_buffers=None, source_urls=None, key="", logger=None):

        self.logger = logger if logger is not None else get_default_logger("JsonArrayFrameIterator")
        self.source_urls = source_urls
        self.source_buffers = source_buffers
        self.key = key

        if self.source_urls is not None:
            self.logger.info("Using source_urls")
            self._source = self._iter_source_urls(self.source_urls)

        elif self.source_buffers is not None:
            self.logger.info("Using source_buffers")
            self._source = self._iter_source_buffers(self.source_buffers)
        else:
            raise ValueError("Must provide source_buffer &| source_url")

    def _is_url(self, p):
        return all([urlparse(p), urlparse(p).netloc])

    def _is_local_path(self, p):
        return Path(p).exists()

    def _iter_source_urls(self, source_urls):
        frame_counter = 0
        for url in source_urls:
            url = str(url)
            if self._is_url(url):
                arr = json.loads(download_bytes(url).decode("utf-8"))
            elif self._is_local_path(url):
                with open(url, "r") as f:
                    arr = json.load(f)
            else:
                raise ValueError(f"Unable to load {url} as json")

            for k in self.key.split("."):
                arr = arr[k]

            for item in arr:
                yield item, url, frame_counter
                frame_counter += 1

    def _iter_source_buffers(self, source_buffers):
        frame_counter = 0
        for buf in source_buffers:
            for item in _iter_buffer(buf, "===END===".encode("utf-8")):
                arr = json.loads(item)

                for k in self.key.split("."):
                    arr = arr[k]

                for item in arr:
                    yield item, None, frame_counter
                    frame_counter += 1

    def __iter__(self):
        return self

    def __next__(self):

        content, original_source_url, frame_index = next(self._source)
        data_file = DataFile(
            # file_id=self.ds.id,
            content=content,
            content_type="text",
            media_frame_index=frame_index,
            original_source_url=original_source_url,
        )
        return data_file


class JsonArrayDataSource(DataSourceCapability):

    stream_media_type = "text"

    class DefaultStreamParameters(DataSourceCapability.DefaultStreamParameters):
        key: str = ""

    @property
    def key(self) -> str:
        value, _ = self._get_parameter("key")
        return value

    def get_json_array_frame_generator(self, stream):

        if "source_paths_generator" in stream.variables:
            source_urls = [path for path, _ in stream.variables["source_paths_generator"]]
            for frame_data in JsonArrayFrameIterator(
                source_buffers=None,
                source_urls=source_urls,
                key=self.key,
                logger=self.logger,
            ):
                yield frame_data
        elif "source_buffers" in stream.variables:
            for frame_data in JsonArrayFrameIterator(
                source_buffers=stream.variables["source_buffers"],
                source_urls=None,
                key=self.key,
                logger=self.logger,
            ):
                yield frame_data
        elif "records" in stream.variables:
            raise NotImplementedError("#### ToDo records ####")
        else:
            raise NotImplementedError("#### ToDo ####")

    def get_next_frame_data(self, stream):

        json_array_frame_generator = stream.variables.get("json_array_frame_generator", None)
        if json_array_frame_generator is None:
            json_array_frame_generator = self.get_json_array_frame_generator(stream)
            stream.variables["json_array_frame_generator"] = json_array_frame_generator

        data_file = next(json_array_frame_generator)
        entities = {}  # ToDo, how/where to load existing entities from
        return data_file, entities


class OutputType(str, Enum):
    numpy = "numpy"
    pillow = "pillow"


class ImageFrameIterator:
    def __init__(
        self,
        source_buffers=None,
        source_urls=None,
        output_type: OutputType = OutputType.numpy,
        logger=None,
    ):

        self.logger = logger if logger is not None else get_default_logger("ImageFrameIterator")
        self.source_buffers = source_buffers
        self.source_urls = source_urls
        self.output_type = output_type

        if self.source_urls is not None:
            self.logger.info("ImageFrameIterator using source_url")
            self._source = iter([str(u) for u, _ in self.source_urls])
        elif self.source_buffers is not None:

            def iter_buffers(bufs):
                for b in bufs:
                    for item in _iter_buffer(b, "===END===".encode("utf-8")):
                        yield item

            self._source = iter_buffers(self.source_buffers)
            self.logger.info(f"ImageFrameIterator using source_buffer")
        else:
            raise ValueError("Must provide source_buffer &| source_url")

        self.frame_index = 0

    def _is_url(self, p):
        return all([urlparse(p), urlparse(p).netloc])

    def _is_local_path(self, p):
        return Path(p).exists()

    def __iter__(self):
        return self

    def _read_image(self, img_src):
        if isinstance(img_src, str) and self._is_local_path(img_src):
            img = _pil_open_image_path(img_src)
            original_source_url = img_src
        elif isinstance(img_src, str) and self._is_url(img_src):
            img = _pil_open_image_url(img_src)
            original_source_url = img_src
        else:
            img = _pil_open_image_bytes(img_src)
            original_source_url = None
        return img, original_source_url

    def __next__(self):
        img_src = next(self._source)
        img, original_source_url = self._read_image(img_src)

        if self.output_type == OutputType.numpy:
            img = np.array(img, dtype=np.uint8)

        data_file = DataFile(
            # file_id=self.ds.id,
            content=img,
            content_type="image",
            media_frame_index=0,
            original_source_url=original_source_url,
        )
        return data_file


class ImageDataSource(DataSourceCapability):
    """

    Example:
        # process a single image
        hl agent run PIPELINE.json image.jpg

        # process many images
        find image/dir/ -n "*.jpg" | hl agent run PIPELINE.json
    """

    stream_media_type = "image"

    class DefaultStreamParameters(DataSourceCapability.DefaultStreamParameters):
        output_type: OutputType = OutputType.numpy

    @property
    def output_type(self) -> OutputType:
        value, _ = self._get_parameter("output_type")
        return value

    def get_image_frame_generator(self, stream):

        if "source_paths_generator" in stream.variables:
            for frame_data in ImageFrameIterator(
                source_buffers=None,
                source_urls=stream.variables["source_paths_generator"],
                output_type=self.output_type,
                logger=self.logger,
            ):
                yield frame_data
        elif "source_buffers" in stream.variables:
            for frame_data in ImageFrameIterator(
                source_buffers=stream.variables["source_buffers"],
                source_urls=None,
                output_type=self.output_type,
                logger=self.logger,
            ):
                yield frame_data
        elif "records" in stream.variables:
            raise NotImplementedError("#### ToDo records ####")
        else:
            raise NotImplementedError("#### ToDo ####")

    def get_next_frame_data(self, stream):

        frame_generator = stream.variables.get("image_frame_generator", None)
        if frame_generator is None:
            frame_generator = self.get_image_frame_generator(stream)
            stream.variables["image_frame_generator"] = frame_generator

        data_file = next(frame_generator)
        entities = {}  # ToDo, how/where to load existing entities from
        return data_file, entities


@require_package(cv2, "cv2", "opencv")
class VideoFrameIterator:

    def __init__(
        self, source_buffers=None, source_urls=None, output_type: OutputType = OutputType.numpy, logger=None
    ):

        self.logger = logger if logger is not None else get_default_logger("VideoFrameIterator")
        self.output_type = output_type
        self.source_urls = source_urls
        self.source_buffers = source_buffers
        self.start_time = datetime.now()
        self.frame_index = 0

        if self.source_urls is not None:
            self.logger.info("Using VideoCapture")
            self._source = iter([str(u) for u in self.source_urls])
            self._cap_cls = cv2.VideoCapture

        elif self.source_buffers is not None:
            self.logger.info("Using BufferedVideoCapture")
            self._source = iter(self.source_buffers)
            self._cap_cls = BufferedVideoCapture
        else:
            raise ValueError("Must provide source_buffer &| source_url")

        self._cap, self._original_source_url = self._init_video_capture()

    def _init_video_capture(self):
        src = next(self._source)
        _cap = self._cap_cls(src)
        if not _cap.isOpened():
            raise ValueError("Cannot open video file or buffer")
        if isinstance(src, str):
            _original_source_url = src
        else:
            _original_source_url = None
        return _cap, _original_source_url

    def __iter__(self):
        return self

    def __next__(self):

        ret, frame_img = self._cap.read()
        if not ret:  # No more frames to read
            self._cap.release()
            self._cap, self._original_source_url = self._init_video_capture()
            ret, frame_img = self._cap.read()
            if not ret:
                self._cap.release()
                raise StopIteration

        frame_img = cv2.cvtColor(frame_img, cv2.COLOR_BGR2RGB)

        if self.output_type == OutputType.pillow:
            frame_img = Image.fromarray(frame_img)

        timestamp = self._cap.get(cv2.CAP_PROP_POS_MSEC)
        data_file = DataFile(
            # file_id=self.ds.id,
            content=frame_img,
            content_type="image",
            recorded_at=self.start_time + timedelta(milliseconds=timestamp),
            media_frame_index=self.frame_index,
            original_source_url=self._original_source_url,
        )
        self.frame_index += 1
        return data_file

    def __del__(self):
        if self._cap.isOpened():
            self._cap.release()


class VideoDataSource(DataSourceCapability):

    stream_media_type = "video"

    class DefaultStreamParameters(DataSourceCapability.DefaultStreamParameters):
        output_type: OutputType = OutputType.numpy

    @property
    def output_type(self) -> OutputType:
        value, _ = self._get_parameter("output_type")
        return value

    def get_video_frame_generator(self, stream):

        if "source_paths_generator" in stream.variables:
            source_urls = [path for path, task_id in stream.variables["source_paths_generator"]]
            for frame_data in VideoFrameIterator(
                source_buffers=None,
                source_urls=source_urls,
                output_type=self.output_type,
                logger=self.logger,
            ):
                yield frame_data
        elif "source_buffers" in stream.variables:
            for frame_data in VideoFrameIterator(
                source_buffers=stream.variables["source_buffers"],
                source_urls=None,
                output_type=self.output_type,
                logger=self.logger,
            ):
                yield frame_data
        elif "records" in stream.variables:
            raise NotImplementedError("#### ToDo records ####")
        else:
            raise NotImplementedError("#### ToDo ####")

    def get_next_frame_data(self, stream):

        video_frame_generator = stream.variables.get("video_frame_generator", None)
        if video_frame_generator is None:
            video_frame_generator = self.get_video_frame_generator(stream)
            stream.variables["video_frame_generator"] = video_frame_generator

        data_file = next(video_frame_generator)
        entities = {}  # ToDo, how/where to load existing entities from
        return data_file, entities
