import os
import queue
import threading
from concurrent.futures import ThreadPoolExecutor, Future
from dataclasses import dataclass, field
from typing import Iterator, Iterable, Optional, Union, TYPE_CHECKING

import httpx

if TYPE_CHECKING:
    from lqs.common.facade import CoreFacade

from .utils import get_relative_object_path, decompress_chunk_bytes
from lqs.transcode import Transcode
from lqs.interface.core.models import Ingestion, Record, Topic


@dataclass
class _RecordSetResult:
    """Carries the results of processing one record set, tagged with a sequence number."""

    sequence: int
    results: list[tuple[Record, Union[bytes, dict]]] = field(default_factory=list)


@dataclass
class _ErrorSentinel:
    """Wraps an exception for propagation through the result queue."""

    exception: BaseException


_DONE = object()


class RecordFetcher:
    """
    High-performance replacement for RecordUtils.iter_record_data.

    Uses a queue-based producer-consumer pipeline with httpx connection pooling,
    proper thread synchronization, bounded backpressure, and ordered yielding.
    """

    def __init__(
        self,
        app: "CoreFacade",
        max_workers: int = 2,
        max_pending_results: int = 64,
    ):
        self.app = app
        self.max_workers = max_workers
        self.max_pending_results = max_pending_results

        # Shared caches protected by _cache_lock
        self._ingestions: dict[str, Ingestion] = {}
        self._topics: dict[str, Topic] = {}
        self._presigned_urls: dict[str, str] = {}
        self._cache_lock = threading.Lock()

    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_000_000,
        max_contiguous_records: int = 1000,
        max_workers: Optional[int] = None,
        order_by_timestamp: bool = True,
        stop_event: Optional[threading.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.

        Drop-in replacement for RecordUtils.iter_record_data with:
        - Thread-safe queue-based pipeline (no shared mutable dicts)
        - httpx connection pooling (reuses TCP+TLS connections)
        - Bounded backpressure (prevents unbounded memory growth)
        - Proper error propagation (errors surface immediately)
        - No polling loops (all blocking via queue operations)
        """
        if stop_event is None:
            stop_event = threading.Event()
        if transcoder is None:
            transcoder = Transcode()

        effective_workers = max_workers if max_workers is not None else self.max_workers

        # Normalize input to an iterator
        if isinstance(records, {}.values().__class__):
            records = iter(records)
        if isinstance(records, list):
            records = iter(records)

        shutdown_event = threading.Event()
        result_queue: queue.Queue = queue.Queue(maxsize=self.max_pending_results)

        client = httpx.Client(
            timeout=httpx.Timeout(60.0, connect=10.0),
            limits=httpx.Limits(
                max_connections=effective_workers + 2,
                max_keepalive_connections=effective_workers + 2,
            ),
        )

        generator_thread = threading.Thread(
            target=self._generate_and_dispatch,
            args=(
                records,
                result_queue,
                shutdown_event,
                stop_event,
                density_threshold,
                max_contiguous_size,
                max_contiguous_records,
                deserialize_results,
                transcoder,
                client,
                use_cache,
                cache_dir,
                read_from_cache,
                write_to_cache,
                effective_workers,
            ),
            daemon=True,
        )

        try:
            generator_thread.start()
            yield from self._consume_results(
                result_queue, shutdown_event, stop_event, order_by_timestamp
            )
        except GeneratorExit:
            shutdown_event.set()
        finally:
            shutdown_event.set()
            generator_thread.join(timeout=10.0)
            client.close()

    # -------------------------------------------------------------------------
    # Thread 1: Record set generation + dispatch to worker pool
    # -------------------------------------------------------------------------

    def _generate_and_dispatch(
        self,
        records: Iterable[Record],
        result_queue: queue.Queue,
        shutdown_event: threading.Event,
        stop_event: threading.Event,
        density_threshold: float,
        max_contiguous_size: int,
        max_contiguous_records: int,
        deserialize_results: bool,
        transcoder: Transcode,
        client: httpx.Client,
        use_cache: bool,
        cache_dir: str,
        read_from_cache: bool,
        write_to_cache: bool,
        max_workers: int,
    ):
        try:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                sequence = 0
                futures: list[Future] = []
                leftover_record = None

                while not shutdown_event.is_set() and not stop_event.is_set():
                    try:
                        record_set, leftover_record = self.app.utils.get_record_set(
                            records=records,
                            carryover_record=leftover_record,
                            density_threshold=density_threshold,
                            max_contiguous_size=max_contiguous_size,
                            max_contiguous_records=max_contiguous_records,
                        )
                    except Exception:
                        if leftover_record is None:
                            # Input iterator exhausted with no leftover
                            break
                        raise

                    if len(record_set) == 0:
                        break

                    future = executor.submit(
                        self._fetch_record_set,
                        record_set=record_set,
                        sequence=sequence,
                        result_queue=result_queue,
                        shutdown_event=shutdown_event,
                        stop_event=stop_event,
                        deserialize_results=deserialize_results,
                        transcoder=transcoder,
                        client=client,
                        use_cache=use_cache,
                        cache_dir=cache_dir,
                        read_from_cache=read_from_cache,
                        write_to_cache=write_to_cache,
                    )
                    futures.append(future)
                    sequence += 1

                    if leftover_record is None:
                        break

                # Wait for all futures and check for errors
                for future in futures:
                    if shutdown_event.is_set():
                        break
                    try:
                        future.result()
                    except Exception as e:
                        self._put_error(result_queue, e)
                        shutdown_event.set()
                        return

            # Signal completion
            self._put_safe(result_queue, _DONE, shutdown_event)

        except Exception as e:
            self._put_error(result_queue, e)
            shutdown_event.set()

    # -------------------------------------------------------------------------
    # Worker threads: Fetch a single record set with retries
    # -------------------------------------------------------------------------

    def _fetch_record_set(
        self,
        record_set: list[Record],
        sequence: int,
        result_queue: queue.Queue,
        shutdown_event: threading.Event,
        stop_event: threading.Event,
        deserialize_results: bool,
        transcoder: Transcode,
        client: httpx.Client,
        use_cache: bool,
        cache_dir: str,
        read_from_cache: bool,
        write_to_cache: bool,
    ):
        max_attempts = 3
        for attempt in range(1, max_attempts + 1):
            if shutdown_event.is_set() or stop_event.is_set():
                return
            try:
                results = self._fetch_record_set_inner(
                    record_set=record_set,
                    deserialize_results=deserialize_results,
                    transcoder=transcoder,
                    client=client,
                    use_cache=use_cache,
                    cache_dir=cache_dir,
                    read_from_cache=read_from_cache,
                    write_to_cache=write_to_cache,
                    shutdown_event=shutdown_event,
                    stop_event=stop_event,
                )
                self._put_safe(
                    result_queue,
                    _RecordSetResult(sequence=sequence, results=results),
                    shutdown_event,
                )
                return
            except Exception as e:
                if attempt == max_attempts:
                    raise
                self.app.logger.debug(
                    f"Error on attempt {attempt} fetching record set: {e}"
                )

    # -------------------------------------------------------------------------
    # Core fetch logic: HTTP GET, decompress, deserialize
    # -------------------------------------------------------------------------

    def _fetch_record_set_inner(
        self,
        record_set: list[Record],
        deserialize_results: bool,
        transcoder: Transcode,
        client: httpx.Client,
        use_cache: bool,
        cache_dir: str,
        read_from_cache: bool,
        write_to_cache: bool,
        shutdown_event: threading.Event,
        stop_event: threading.Event,
    ) -> list[tuple[Record, Union[bytes, dict]]]:
        results: list[tuple[Record, Union[bytes, dict]]] = []
        datastore_id = self.app.get_datastore_id() if use_cache else None

        # All records in a set share the same ingestion and source
        first_record = record_set[0]
        ingestion = self._get_ingestion(str(first_record.ingestion_id))
        object_key = str(ingestion.object_key)
        if first_record.source is not None:
            object_key = get_relative_object_path(
                object_key=object_key, source=first_record.source
            )

        # Separate cached vs needs-fetch records (single pass)
        cached_results: dict[int, tuple[Record, Union[bytes, dict]]] = {}
        records_needing_fetch: list[tuple[int, Record]] = []

        for idx, record in enumerate(record_set):
            if use_cache and read_from_cache:
                cache_path = self._cache_path(cache_dir, datastore_id, record)
                if os.path.exists(cache_path):
                    with open(cache_path, "rb") as f:
                        record_bytes = f.read()
                    if deserialize_results:
                        topic = self._get_topic(str(record.topic_id))
                        record_data = transcoder.deserialize(
                            type_encoding=topic.type_encoding,
                            type_name=topic.type_name,
                            type_data=topic.type_data,
                            message_bytes=record_bytes,
                        )
                        cached_results[idx] = (record, record_data)
                    else:
                        cached_results[idx] = (record, record_bytes)
                    continue
            records_needing_fetch.append((idx, record))

        # If all cached, return immediately
        if len(records_needing_fetch) == 0:
            return [cached_results[i] for i in sorted(cached_results)]

        # Calculate byte range from records that need fetching
        first_fetch = records_needing_fetch[0][1]
        last_fetch = records_needing_fetch[-1][1]
        start_offset = first_fetch.data_offset
        end_offset = last_fetch.data_offset + last_fetch.data_length

        # Get presigned URL (lock-protected cache)
        presigned_url = self._get_presigned_url(
            object_key=object_key, ingestion=ingestion
        )

        # HTTP GET with Range header
        data_bytes = self._http_range_get(
            client=client,
            presigned_url=presigned_url,
            start_offset=start_offset,
            end_offset=end_offset,
            object_key=object_key,
            ingestion=ingestion,
        )

        # Parse records from the fetched byte range
        data_view = memoryview(data_bytes)
        current_offset = start_offset
        decompressed_bytes: Optional[bytes] = None
        compressed_chunk_offset: Optional[int] = None
        stream_pos = 0  # position within data_view

        for idx, record in enumerate(record_set):
            if shutdown_event.is_set() or stop_event.is_set():
                break

            if idx in cached_results:
                results.append(cached_results[idx])
                continue

            data_offset = record.data_offset
            data_length = record.data_length

            if (
                compressed_chunk_offset is not None
                and record.chunk_compression is not None
                and data_offset == compressed_chunk_offset
            ):
                # Reuse previously decompressed chunk
                message_bytes = decompressed_bytes[
                    record.chunk_offset : record.chunk_offset + record.chunk_length
                ]
            else:
                # Skip gap bytes and extract record data
                skip = data_offset - current_offset
                stream_pos += skip
                message_bytes = bytes(data_view[stream_pos : stream_pos + data_length])
                stream_pos += data_length
                current_offset = data_offset + data_length

                # Decompress if needed
                if record.chunk_compression not in (None, "", "none"):
                    decompressed_bytes = decompress_chunk_bytes(
                        chunk_bytes=message_bytes,
                        chunk_compression=record.chunk_compression,
                        chunk_length=record.chunk_length,
                    )
                    message_bytes = decompressed_bytes[
                        record.chunk_offset : record.chunk_offset + record.chunk_length
                    ]
                    compressed_chunk_offset = data_offset

            # Write to cache if requested
            if use_cache and write_to_cache:
                cache_path = self._cache_path(cache_dir, datastore_id, record)
                os.makedirs(os.path.dirname(cache_path), exist_ok=True)
                with open(cache_path, "wb") as f:
                    f.write(message_bytes)

            # Deserialize if requested
            if deserialize_results:
                topic = self._get_topic(str(record.topic_id))
                record_data = transcoder.deserialize(
                    type_encoding=topic.type_encoding,
                    type_name=topic.type_name,
                    type_data=topic.type_data,
                    message_bytes=message_bytes,
                )
                results.append((record, record_data))
            else:
                results.append((record, message_bytes))

        return results

    # -------------------------------------------------------------------------
    # HTTP fetch with Range header and presigned URL retry
    # -------------------------------------------------------------------------

    def _http_range_get(
        self,
        client: httpx.Client,
        presigned_url: str,
        start_offset: int,
        end_offset: int,
        object_key: str,
        ingestion: Ingestion,
    ) -> bytes:
        headers = {"Range": f"bytes={start_offset}-{end_offset - 1}"}
        try:
            response = client.get(presigned_url, headers=headers)
            response.raise_for_status()
            return response.content
        except (httpx.HTTPStatusError, httpx.TransportError):
            # Presigned URL may have expired; evict, regenerate, retry
            self.app.logger.debug(
                "Presigned URL request failed, regenerating and retrying."
            )
            with self._cache_lock:
                self._presigned_urls.pop(object_key, None)
            new_url = self._get_presigned_url(
                object_key=object_key, ingestion=ingestion
            )
            response = client.get(new_url, headers=headers)
            response.raise_for_status()
            return response.content

    # -------------------------------------------------------------------------
    # Main thread: consume results in order
    # -------------------------------------------------------------------------

    def _consume_results(
        self,
        result_queue: queue.Queue,
        shutdown_event: threading.Event,
        stop_event: threading.Event,
        order_by_timestamp: bool,
    ) -> Iterator[tuple[Record, Union[bytes, dict]]]:
        next_sequence = 0
        buffer: dict[int, _RecordSetResult] = {}

        while not shutdown_event.is_set() and not stop_event.is_set():
            # Drain buffered results that are ready (in sequence order)
            while next_sequence in buffer:
                result = buffer.pop(next_sequence)
                if order_by_timestamp:
                    result.results.sort(
                        key=lambda x: (x[0].timestamp, str(x[0].topic_id))
                    )
                yield from result.results
                next_sequence += 1

            # Wait for next item from queue
            try:
                item = result_queue.get(timeout=1.0)
            except queue.Empty:
                continue

            if item is _DONE:
                # Drain any remaining buffered results
                while next_sequence in buffer:
                    result = buffer.pop(next_sequence)
                    if order_by_timestamp:
                        result.results.sort(
                            key=lambda x: (x[0].timestamp, str(x[0].topic_id))
                        )
                    yield from result.results
                    next_sequence += 1
                return

            if isinstance(item, _ErrorSentinel):
                raise item.exception

            # Buffer or yield immediately based on sequence
            if item.sequence == next_sequence:
                if order_by_timestamp:
                    item.results.sort(
                        key=lambda x: (x[0].timestamp, str(x[0].topic_id))
                    )
                yield from item.results
                next_sequence += 1
            else:
                buffer[item.sequence] = item

    # -------------------------------------------------------------------------
    # Lock-protected cache helpers
    # -------------------------------------------------------------------------

    def _get_ingestion(self, ingestion_id: str) -> Ingestion:
        with self._cache_lock:
            cached = self._ingestions.get(ingestion_id)
            if cached is not None:
                return cached
        ingestion = self.app.fetch.ingestion(ingestion_id).data
        with self._cache_lock:
            self._ingestions[ingestion_id] = ingestion
        return ingestion

    def _get_topic(self, topic_id: str) -> Topic:
        with self._cache_lock:
            cached = self._topics.get(topic_id)
            if cached is not None:
                return cached
        topic = self.app.fetch.topic(topic_id).data
        with self._cache_lock:
            self._topics[topic_id] = topic
        return topic

    def _get_presigned_url(self, object_key: str, ingestion: Ingestion) -> str:
        with self._cache_lock:
            cached = self._presigned_urls.get(object_key)
            if cached is not None:
                return cached
        presigned_url = self.app.utils.get_presigned_url(
            object_key=object_key,
            object_store_id=ingestion.object_store_id,
            log_id=ingestion.log_id,
        )
        with self._cache_lock:
            self._presigned_urls[object_key] = presigned_url
        return presigned_url

    # -------------------------------------------------------------------------
    # Utility helpers
    # -------------------------------------------------------------------------

    @staticmethod
    def _cache_path(cache_dir: str, datastore_id, record: Record) -> str:
        return (
            f"{cache_dir}/datastores/{datastore_id}/logs/{record.log_id}"
            f"/topics/{record.topic_id}/records/{record.timestamp:09}.bin"
        )

    @staticmethod
    def _put_safe(q: queue.Queue, item, shutdown_event: threading.Event):
        """Put an item on a queue, respecting shutdown."""
        while not shutdown_event.is_set():
            try:
                q.put(item, timeout=1.0)
                return
            except queue.Full:
                continue

    @staticmethod
    def _put_error(q: queue.Queue, exception: BaseException):
        """Best-effort put of an error sentinel on the queue."""
        try:
            q.put(_ErrorSentinel(exception), timeout=2.0)
        except queue.Full:
            pass
