# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import asyncio
import os
import operator
from concurrent.futures import Future, ThreadPoolExecutor
from typing import AsyncGenerator, Callable, Iterator, List, Optional

from pypaimon.common.options.core_options import ChangelogProducer
from pypaimon.common.predicate import Predicate
from pypaimon.consumer.consumer import Consumer
from pypaimon.manifest.manifest_file_manager import ManifestFileManager
from pypaimon.manifest.manifest_list_manager import ManifestListManager
from pypaimon.read.plan import Plan
from pypaimon.read.query_auth_split import wrap_plan_with_auth
from pypaimon.execution.python.read.scanner.append_table_split_generator import \
    AppendTableSplitGenerator
from pypaimon.execution.python.read.scanner.changelog_follow_up_scanner import \
    ChangelogFollowUpScanner
from pypaimon.execution.python.read.scanner.delta_follow_up_scanner import DeltaFollowUpScanner
from pypaimon.execution.python.read.scanner.file_scanner import FileScanner
from pypaimon.execution.python.read.scanner.follow_up_scanner import FollowUpScanner
from pypaimon.execution.python.read.scanner.primary_key_table_split_generator import \
    PrimaryKeyTableSplitGenerator
from pypaimon.snapshot.snapshot import Snapshot
from pypaimon.table.bucket_mode import BucketMode


class AsyncStreamingTableScan:
    """
    Async streaming table scan for continuous reads from Paimon tables.

    This class provides an async iterator that continuously polls for new
    snapshots and yields Plans containing splits for new data.

    Usage:
        scan = AsyncStreamingTableScan(table)

        async for plan in scan.stream():
            for split in plan.splits():
                # Process the data
                pass

    For synchronous usage:
        for plan in scan.stream_sync():
            process(plan)
    """

    def __init__(
        self,
        table,
        predicate: Optional[Predicate] = None,
        poll_interval_ms: int = 1000,
        follow_up_scanner: Optional[FollowUpScanner] = None,
        bucket_filter: Optional[Callable[[int], bool]] = None,
        prefetch_enabled: bool = True,
        diff_threshold: int = 10,
        consumer_id: Optional[str] = None
    ):
        """Initialize the streaming table scan."""
        self.table = table
        self.predicate = predicate
        self.poll_interval = poll_interval_ms / 1000.0

        # Bucket filter for parallel consumption
        self._bucket_filter = bucket_filter

        # Retain the constructor argument without replacing physical stream
        # events with a net snapshot diff when a consumer falls behind.
        self._diff_threshold = diff_threshold

        # Prefetching configuration
        self._prefetch_enabled = prefetch_enabled
        self._prefetch_future: Optional[Future] = None
        self._prefetch_snapshot_id: Optional[int] = None
        self._lookahead_skips = 0
        self._prefetch_executor = ThreadPoolExecutor(max_workers=1) if prefetch_enabled else None
        self._lookahead_size = 10  # How many snapshots to look ahead

        # Initialize managers
        self._snapshot_manager = table.snapshot_manager()
        self._manifest_list_manager = ManifestListManager(table)
        self._manifest_file_manager = ManifestFileManager(table)

        # Consumer management for persisting streaming progress
        self._consumer_id = consumer_id
        self._read_type = None
        self._query_auth_fn = self.table.catalog_environment.table_query_auth(
            self.table.options, self.table.identifier)
        self._consumer_manager = (
            table.consumer_manager()
            if consumer_id else None
        )

        # Scanner for determining which snapshots to read
        # Auto-select based on changelog-producer if not explicitly provided
        self.follow_up_scanner = follow_up_scanner or self._create_follow_up_scanner()

        # State tracking
        self._next_snapshot_id: Optional[int] = None
        self._pending_consumer_snapshot: Optional[int] = None
        self._pending_consumer_owner = None
        self._consumer_restored = False
        self._watermark: Optional[int] = None
        self._closed = False
        self._generation = 0
        self._missing_snapshot_polls = 0

    @property
    def next_snapshot_id(self):
        return self.checkpoint()

    @next_snapshot_id.setter
    def next_snapshot_id(self, next_snapshot_id):
        self.restore(next_snapshot_id)

    def checkpoint(self):
        """Return the next snapshot to consume, without acknowledging it."""
        return self._next_snapshot_id

    def watermark(self):
        """Return the watermark of the last successfully planned frame."""
        return self._watermark

    def restore(self, next_snapshot_id=None):
        """Restore a checkpoint, or reset to the initial/consumer position."""
        self._require_open()
        next_snapshot_id = self._snapshot_id(next_snapshot_id)
        if next_snapshot_id is not None and next_snapshot_id < 1:
            raise ValueError("next snapshot id must be positive")
        self._generation += 1
        self._discard_prefetch()
        self._next_snapshot_id = next_snapshot_id
        self._consumer_restored = next_snapshot_id is not None
        self._pending_consumer_snapshot = None
        self._pending_consumer_owner = None
        self._lookahead_skips = 0
        self._missing_snapshot_polls = 0

    def notify_checkpoint_complete(self, next_snapshot):
        """Persist progress only after the caller has processed its plan."""
        self._require_open()
        next_snapshot = self._snapshot_id(next_snapshot)
        if self._consumer_manager and self._consumer_id and next_snapshot is not None:
            if next_snapshot < 1:
                raise ValueError("consumer next snapshot must be positive")
            self._consumer_manager.reset_consumer(
                self._consumer_id, Consumer(next_snapshot=next_snapshot))

    def plan(self):
        """Plan one available frame; return None instead of waiting for data.

        Initial state may be empty. Empty follow-up frames and irrelevant
        commits advance the checkpoint without emitting a frame, as in Java
        and Rust. A planning failure leaves the selected snapshot retryable.
        """
        self._require_open()
        if self._next_snapshot_id is None and not self._consumer_restored:
            if self._consumer_manager:
                consumer = self._consumer_manager.consumer(self._consumer_id)
                if consumer is not None:
                    self.restore(consumer.next_snapshot)
            self._consumer_restored = True
        if self._next_snapshot_id is None:
            generation = self._generation
            latest = self._snapshot_manager.get_latest_snapshot()
            if latest is None:
                self.__auth_query()
                return None
            plan = self._create_initial_plan(latest)
            self._check_generation(generation)
            self._next_snapshot_id = latest.id + 1
            self._watermark = latest.watermark
            return plan

        while True:
            generation = self._generation
            if self._prefetch_future is not None:
                future = self._prefetch_future
                self._prefetch_future = None
                self._prefetch_snapshot_id = None
                result = future.result()
            else:
                result = self._fetch_plan_with_lookahead(self._next_snapshot_id)
            self._check_generation(generation)
            plan, next_id, skipped, watermark = result
            self._lookahead_skips += skipped
            self._next_snapshot_id = next_id
            if plan is not None:
                self._missing_snapshot_polls = 0
                self._watermark = watermark
                if plan.splits():
                    if self._prefetch_enabled:
                        self._start_prefetch(next_id)
                    return plan
            elif skipped == 0:
                self.__auth_query()
                self._check_missing_snapshot()
                return None
            else:
                self._missing_snapshot_polls = 0

    async def stream(self) -> AsyncGenerator[Plan, None]:
        """Poll plan(), acknowledging each yielded frame on iterator resume."""
        owner = object()
        try:
            while True:
                self._flush_pending_consumer()
                checkpoint = self.checkpoint()
                plan = self.plan()
                if plan is None:
                    next_snapshot = self.checkpoint()
                    if next_snapshot is not None and next_snapshot != checkpoint:
                        self._stage_consumer(owner)
                        self._flush_pending_consumer()
                    await asyncio.sleep(self.poll_interval)
                    continue
                self._stage_consumer(owner)
                yield plan
        except (GeneratorExit, asyncio.CancelledError):
            if self._pending_consumer_owner is owner:
                self._pending_consumer_snapshot = None
                self._pending_consumer_owner = None
                self._discard_prefetch()
            raise

    def stream_sync(self) -> Iterator[Plan]:
        """Synchronously poll the same scan, closing its async iterator."""
        loop = asyncio.new_event_loop()
        iterator = self.stream()
        try:
            while True:
                yield loop.run_until_complete(iterator.__anext__())
        finally:
            try:
                loop.run_until_complete(iterator.aclose())
            finally:
                loop.close()

    def _require_open(self):
        if self._closed:
            raise RuntimeError("stream scan is closed")

    @staticmethod
    def _snapshot_id(value):
        if value is None:
            return None
        value = operator.index(value)
        if not -(1 << 63) <= value < (1 << 63):
            raise OverflowError("snapshot id is outside the signed 64-bit range")
        return value

    def _check_generation(self, generation):
        self._require_open()
        if generation != self._generation:
            raise RuntimeError("stream scan was restored during planning")

    def _check_missing_snapshot(self):
        # Java NextSnapshotFetcher and Rust periodically distinguish delayed
        # publication from an expired checkpoint or a recreated table.
        self._missing_snapshot_polls += 1
        if self._missing_snapshot_polls % 16:
            return
        latest = self._snapshot_manager.get_latest_snapshot()
        earliest = (self._snapshot_manager.try_get_earliest_snapshot(latest.id)
                    if latest is not None else None)
        checkpoint = self._next_snapshot_id
        if ((earliest is not None and earliest.id > checkpoint)
                or (latest is not None and checkpoint > latest.id + 1)
                or (latest is None and checkpoint > 1)):
            raise ValueError("stream checkpoint %s is outside the available snapshots %s..%s" % (
                checkpoint, earliest.id if earliest is not None else None,
                latest.id if latest is not None else None))

    def _discard_prefetch(self):
        future = self._prefetch_future
        self._prefetch_future = None
        self._prefetch_snapshot_id = None
        if future is not None:
            future.cancel()

    def cancel_prefetch(self):
        """Discard speculative work when a public polling iterator stops."""
        self._discard_prefetch()

    def close(self):
        """Discard speculative work without acknowledging pending rows.

        Running Python I/O cannot be interrupted. Its detached result cannot
        advance the checkpoint or be reused after close/restore.
        """
        if self._closed:
            return
        self._closed = True
        self._generation += 1
        self._discard_prefetch()
        self._pending_consumer_snapshot = None
        self._pending_consumer_owner = None
        if self._prefetch_executor is not None:
            self._prefetch_executor.shutdown(wait=False)

    def _stage_consumer(self, owner) -> None:
        """Stage next_snapshot_id to be written to disk on the next generator resume."""
        self._pending_consumer_owner = owner
        if self._consumer_manager and self._consumer_id and self.next_snapshot_id is not None:
            self._pending_consumer_snapshot = self.next_snapshot_id

    def _flush_pending_consumer(self) -> None:
        """Flush the staged consumer position to disk.

        Called at the resume point after each yield — i.e. when the caller calls
        __anext__() to request the next plan, which happens after the caller's loop
        body (to_arrow + sink write) has completed. This gives at-least-once semantics:
        the consumer file is only advanced after the caller has processed the prior plan.
        """
        if self._consumer_manager and self._consumer_id and self._pending_consumer_snapshot is not None:
            self.notify_checkpoint_complete(self._pending_consumer_snapshot)
            self._pending_consumer_snapshot = None
            self._pending_consumer_owner = None

    def __apply_auth(self, plan) -> Plan:
        return wrap_plan_with_auth(self.__auth_query(), plan)

    def __auth_query(self):
        from pypaimon.read.auth_support import authorize
        return authorize(self.table, self._query_auth_fn, self._read_type, self.predicate)

    def _start_prefetch(self, snapshot_id: int) -> None:
        """Start prefetching the next scannable snapshot in a background thread."""
        if self._prefetch_future is not None or self._prefetch_executor is None:
            return  # Already prefetching or executor not available

        self._prefetch_snapshot_id = snapshot_id
        # Submit to thread pool - this starts immediately, not when event loop runs
        self._prefetch_future = self._prefetch_executor.submit(
            self._fetch_plan_with_lookahead,
            snapshot_id
        )

    def _fetch_plan_with_lookahead(self, start_id: int) -> tuple:
        """Plan without mutating progress; the caller owns checkpoint updates."""
        snapshot, next_id, skipped = self._snapshot_manager.find_next_scannable(
            start_id, self.follow_up_scanner.should_scan,
            lookahead_size=self._lookahead_size)
        if snapshot is None:
            return None, next_id, skipped, None
        plan = self._create_follow_up_plan(snapshot)
        return plan, next_id, skipped, snapshot.watermark

    def _create_follow_up_plan(self, snapshot: Snapshot) -> Plan:
        """Route to changelog or delta plan based on scanner type."""
        if isinstance(self.follow_up_scanner, ChangelogFollowUpScanner):
            plan = self._create_changelog_plan(snapshot)
        else:
            plan = self._create_delta_plan(snapshot)
        return self.__apply_auth(plan)

    def _create_follow_up_scanner(self) -> FollowUpScanner:
        """Create the appropriate follow-up scanner based on changelog-producer option."""
        changelog_producer = self.table.options.changelog_producer()
        if changelog_producer == ChangelogProducer.NONE:
            return DeltaFollowUpScanner()
        else:
            # INPUT, FULL_COMPACTION, LOOKUP all use changelog scanner
            return ChangelogFollowUpScanner()

    def _filter_entries_for_shard(self, entries: List) -> List:
        """Filter manifest entries by bucket filter, if set."""
        if self._bucket_filter is not None:
            return [e for e in entries if self._bucket_filter(e.bucket)]
        return entries

    def _only_read_real_buckets(self) -> bool:
        # Java DataTableStreamScan leaves postpone pending data visible only
        # when there is no changelog producer to emit it after assignment.
        return (self.table.options.bucket() == BucketMode.POSTPONE_BUCKET.value
                and self.table.options.changelog_producer() != ChangelogProducer.NONE)

    def __create_initial_plan_raw(self, snapshot, auth_result=None):
        def all_manifests():
            return self._manifest_list_manager.read_all(snapshot), snapshot

        starting_scanner = FileScanner(
            self.table,
            all_manifests,
            predicate=self.predicate,
            limit=None
        )
        starting_scanner.only_read_real_buckets = self._only_read_real_buckets()
        producer = self.table.options.changelog_producer()
        if producer == ChangelogProducer.LOOKUP:
            # L0 will produce its changelog after lookup compaction, so the
            # initial state must contain only the already compacted levels.
            starting_scanner.level_filter = lambda level: level > 0
        elif producer == ChangelogProducer.FULL_COMPACTION:
            options = self.table.options.options.to_map()
            levels = options.get('num-levels')
            if levels is None:
                # Java CoreOptions.numLevels uses incrementSafely(trigger).
                trigger = int(options.get('num-sorted-run.compaction-trigger', 5))
                levels = min(trigger + 1, (1 << 31) - 1)
            last_level = int(levels) - 1
            starting_scanner.level_filter = lambda level: level == last_level
        if auth_result is not None:
            from pypaimon.execution.python.read.auth_support import prune_scanner_by_auth
            prune_scanner_by_auth(self.table, starting_scanner, auth_result)
        plan = starting_scanner.scan()
        if self._bucket_filter is not None:
            plan = Plan([split for split in plan.splits() if self._bucket_filter(split.bucket)],
                        snapshot_id=plan.snapshot_id)
        return plan

    def _create_initial_plan(self, snapshot: Snapshot) -> Plan:
        """Create a Plan for the initial full scan of the latest snapshot."""
        auth_result = self.__auth_query()
        plan = self.__create_initial_plan_raw(snapshot, auth_result)
        return wrap_plan_with_auth(auth_result, plan)

    def _create_delta_plan(self, snapshot: Snapshot) -> Plan:
        """Read new files from delta_manifest_list (changelog-producer=none)."""
        manifest_files = self._manifest_list_manager.read_delta(snapshot)
        return self._create_plan_from_manifests(manifest_files, snapshot.id)

    def _create_changelog_plan(self, snapshot: Snapshot) -> Plan:
        """Java ChangelogFollowUpScanner consumes any snapshot with a changelog."""
        manifest_files = self._manifest_list_manager.read_changelog(snapshot)
        return self._create_plan_from_manifests(manifest_files, snapshot.id)

    def _create_plan_from_manifests(self, manifest_files: List, snapshot_id=None) -> Plan:
        """Create splits from manifest files, applying shard filtering."""
        if not manifest_files:
            return Plan([], snapshot_id=snapshot_id)

        # Use configurable parallelism from table options
        max_workers = max(8, self.table.options.scan_manifest_parallelism(os.cpu_count() or 8))

        def require_add(entry):
            if entry.kind != 0:
                raise ValueError("Incremental manifests must contain only ADD entries")
            return True

        # Validate before the manifest reader reconciles ADD/DELETE entries.
        entries = self._manifest_file_manager.read_entries_parallel(
            manifest_files,
            manifest_entry_filter=require_add,
            max_workers=max_workers
        )

        # Apply shard/bucket filtering for parallel consumption
        if self._only_read_real_buckets():
            entries = [entry for entry in entries if entry.bucket >= 0]
        entries = self._filter_entries_for_shard(entries) if entries else []
        if not entries:
            return Plan([], snapshot_id=snapshot_id)

        # Get split options from table
        options = self.table.options
        target_split_size = options.source_split_target_size()
        open_file_cost = options.source_split_open_file_cost()

        # Create appropriate split generator based on table type
        if self.table.is_primary_key_table:
            split_generator = PrimaryKeyTableSplitGenerator(
                self.table,
                target_split_size,
                open_file_cost,
                deletion_files_map={},
                snapshot_id=snapshot_id,
            )
        else:
            split_generator = AppendTableSplitGenerator(
                self.table,
                target_split_size,
                open_file_cost,
                deletion_files_map={},
                snapshot_id=snapshot_id,
            )

        splits = split_generator.create_splits(entries)
        for split in splits:
            split.is_streaming = True
        return Plan(splits, snapshot_id=snapshot_id)