import time
from collections.abc import Callable
from dataclasses import dataclass, field
from graphlib import CycleError, TopologicalSorter
from typing import Any

from dbslice.adapters.base import DatabaseAdapter
from dbslice.config import (
    DatabaseType,
    ExtractConfig,
    OutputFormat,
    SeedSpec,
    TraversalDirection,
)
from dbslice.constants import DEFAULT_ANONYMIZATION_SEED, DEFAULT_TRAVERSAL_DEPTH
from dbslice.core.graph import GraphTraverser, TraversalConfig, TraversalResult
from dbslice.exceptions import (
    ExtractionError,
    NoRowsFoundError,
    TableNotFoundError,
)
from dbslice.logging import get_logger
from dbslice.models import ForeignKey, SchemaGraph
from dbslice.output.sql import SQLGenerator
from dbslice.utils.anonymizer import DeterministicAnonymizer
from dbslice.utils.connection import get_adapter_for_url, parse_database_url
from dbslice.validation import ExtractionValidator, ValidationResult

logger = get_logger(__name__)

# Type alias for progress callback functions.
#
# Progress callbacks are called during extraction to report status updates.
# Signature: (stage: str, message: str, current: int, total: int) -> None
#
# Args:
#     stage: Current extraction stage (e.g., "schema", "fetch", "sort")
#     message: Human-readable status message
#     current: Current item number (0 if not applicable)
#     total: Total items (0 if not applicable)
ProgressCallback = Callable[[str, str, int, int], None]


@dataclass
class ExtractedRow:
    """A single extracted record."""

    table: str
    pk_values: tuple[Any, ...]
    data: dict[str, Any]


@dataclass
class ExtractionResult:
    """
    Complete extraction result with data, metadata, and cycle information.

    Attributes:
        tables: Mapping of table names to lists of row data dictionaries
        insert_order: Topologically sorted table names for safe INSERT order
        stats: Row count statistics per table
        traversal_path: List of traversal steps for debugging (verbose mode)
        has_cycles: Whether circular FK dependencies were detected
        broken_fks: List of ForeignKey objects that were broken to resolve cycles
        deferred_updates: List of DeferredUpdate objects to restore broken FK values
        cycle_infos: List of CycleInfo objects describing detected cycles
        validation_result: Result of extraction validation (None if validation skipped)
    """

    tables: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
    insert_order: list[str] = field(default_factory=list)
    stats: dict[str, int] = field(default_factory=dict)
    traversal_path: list[str] = field(default_factory=list)
    has_cycles: bool = False
    broken_fks: list[Any] = field(default_factory=list)  # list[ForeignKey]
    deferred_updates: list[Any] = field(default_factory=list)  # list[DeferredUpdate]
    cycle_infos: list[Any] = field(default_factory=list)  # list[CycleInfo]
    validation_result: ValidationResult | None = None
    profiler: Any = None  # Optional QueryProfiler
    used_deferred_cycle_strategy: bool = False

    def total_rows(self) -> int:
        """Get total number of extracted rows."""
        if self.tables:
            return sum(len(rows) for rows in self.tables.values())
        # Streaming/dry-run mode: tables is empty, use stats
        return sum(self.stats.values())

    def table_count(self) -> int:
        """Get number of tables with data."""
        if self.tables:
            return len(self.tables)
        # Streaming/dry-run mode: tables is empty, use stats
        return len(self.stats)


class ExtractionEngine:
    """
    Main extraction engine that orchestrates the extraction process.

    Flow:
    1. Connect to database
    2. Introspect schema
    3. Parse seeds and fetch seed rows
    4. Traverse FK graph to find related records
    5. Fetch all identified rows
    6. Topologically sort tables
    7. Return result for output generation
    """

    def __init__(
        self,
        config: ExtractConfig,
        progress_callback: ProgressCallback | None = None,
    ):
        self.config = config
        self.adapter: DatabaseAdapter | None = None
        self.schema: SchemaGraph | None = None
        self.progress_callback = progress_callback

        # Initialize anonymizer if needed (schema will be set after introspection)
        self.anonymizer: DeterministicAnonymizer | None = None
        if config.anonymize or config.redact_fields:
            self.anonymizer = DeterministicAnonymizer(
                seed=config.anonymization_seed or DEFAULT_ANONYMIZATION_SEED
            )
            self.anonymizer.configure(
                config.redact_fields,
                field_providers=config.anonymization_field_providers,
                patterns=config.anonymization_patterns,
                security_null_fields=config.security_null_fields,
            )

    def _log(self, stage: str, message: str, current: int = 0, total: int = 0) -> None:
        """Send progress update to callback if configured."""
        if self.progress_callback:
            self.progress_callback(stage, message, current, total)

    def extract(self) -> tuple[ExtractionResult, SchemaGraph]:
        """
        Perform the extraction.

        Returns:
            Tuple of (ExtractionResult, SchemaGraph) - schema is returned
            to avoid needing to reconnect for SQL generation.
        """
        start_time = time.time()

        db_config = parse_database_url(self.config.database_url)
        logger.info(
            "Starting extraction",
            database=db_config.database,
            db_type=db_config.db_type.value,
            seed_count=len(self.config.seeds),
            depth=self.config.depth,
            direction=self.config.direction.value,
        )

        profiler = None
        if self.config.profile:
            from dbslice.utils.profiling import QueryProfiler

            profiler = QueryProfiler()
            logger.info("Query profiling enabled")

        # Get appropriate adapter with profiler if available
        from dbslice.adapters.postgresql import PostgreSQLAdapter

        if db_config.db_type.value == "postgresql":
            self.adapter = PostgreSQLAdapter(
                batch_size=self.config.db_batch_size,
                profiler=profiler,
                schema=self.config.schema,
                allow_unsafe_where=self.config.allow_unsafe_where,
            )
        else:
            self.adapter = get_adapter_for_url(self.config.database_url)

        try:
            with logger.timed_operation("database_connection", database=db_config.database):
                self.adapter.connect(self.config.database_url)
            self._log("connect", "Connected successfully")

            # Use snapshot transaction for consistent reads
            with self.adapter.snapshot_transaction():
                result = self._do_extract(db_config.db_type)

                # Log extraction summary
                elapsed_ms = int((time.time() - start_time) * 1000)
                logger.info(
                    "Extraction completed successfully",
                    total_rows=result.total_rows(),
                    table_count=result.table_count(),
                    duration_ms=elapsed_ms,
                    has_cycles=result.has_cycles,
                )

                if profiler:
                    result.profiler = profiler

                # Return schema along with result to avoid reconnection
                assert self.schema is not None
                return result, self.schema
        except Exception as e:
            elapsed_ms = int((time.time() - start_time) * 1000)
            logger.error(
                "Extraction failed",
                error=str(e),
                duration_ms=elapsed_ms,
                exc_info=True,
            )
            raise
        finally:
            if self.adapter:
                self.adapter.close()
                logger.debug("Database connection closed")

    def _do_extract(self, db_type: DatabaseType) -> ExtractionResult:
        """Perform extraction within a snapshot transaction."""
        assert self.adapter is not None

        self._log("schema", "Introspecting database schema...")
        with logger.timed_operation("schema_introspection"):
            self.schema = self.adapter.get_schema()

        for vfk in self.config.virtual_foreign_keys:
            self.schema.add_virtual_fk(vfk)

        logger.info(
            "Schema introspection complete",
            table_count=len(self.schema.tables),
            fk_count=len(self.schema.edges),
            virtual_fk_count=len(self.schema.virtual_edges),
        )
        self._log(
            "schema",
            f"Found {len(self.schema.tables)} tables, {len(self.schema.edges)} foreign keys"
            + (
                f", {len(self.schema.virtual_edges)} virtual foreign keys"
                if self.schema.virtual_edges
                else ""
            ),
        )

        if self.anonymizer:
            self.anonymizer.schema = self.schema
            logger.debug("Anonymizer configured with schema")

        all_records: dict[str, set[tuple[Any, ...]]] = {}
        all_paths: list[str] = []

        total_seeds = len(self.config.seeds)
        for i, seed in enumerate(self.config.seeds):
            seed_desc = (
                f"{seed.table}.{seed.column}={seed.value}"
                if seed.column
                else f"{seed.table}:{seed.where_clause}"
            )
            self._log("seed", f"Processing seed: {seed_desc}", i + 1, total_seeds)
            logger.debug("Processing seed", seed=seed_desc, index=i + 1, total=total_seeds)

            traversal_result = self._process_seed(seed)
            all_paths.extend(traversal_result.traversal_path)

            for table, pks in traversal_result.records.items():
                if table not in all_records:
                    all_records[table] = set()
                all_records[table].update(pks)

            logger.info(
                "Seed traversal complete",
                seed=seed_desc,
                records_found=traversal_result.total_records(),
                tables_affected=traversal_result.table_count(),
            )
            self._log(
                "seed",
                f"Found {traversal_result.total_records()} records across {traversal_result.table_count()} tables",
                i + 1,
                total_seeds,
            )

        self._log("sort", "Sorting tables by dependencies...")
        with logger.timed_operation("topological_sort", table_count=len(all_records)):
            insert_order, broken_fks, cycle_infos, used_deferred_cycle_strategy = (
                self._topological_sort(set(all_records.keys()), db_type)
            )

        if broken_fks:
            logger.warning(
                "Circular dependencies detected",
                broken_fk_count=len(broken_fks),
                cycle_count=len(cycle_infos),
            )

        # Dry-run mode: skip data fetch and return stats only
        if self.config.dry_run:
            dry_run_stats = {table: len(pks) for table, pks in all_records.items()}
            total_rows_estimate = sum(dry_run_stats.values())

            logger.info(
                "Dry-run complete",
                total_rows=total_rows_estimate,
                table_count=len(all_records),
            )
            self._log(
                "dry_run",
                f"Dry-run summary: {total_rows_estimate} rows across {len(all_records)} tables would be extracted",
            )

            for table in insert_order:
                if table in dry_run_stats:
                    self._log("dry_run", f"  {table}: {dry_run_stats[table]} rows")

            if all_paths:
                self._log("dry_run", "Traversal path:")
                for path_entry in all_paths:
                    self._log("dry_run", f"  {path_entry}")

            if cycle_infos:
                self._log("dry_run", f"Detected {len(cycle_infos)} circular dependency cycle(s)")

            return ExtractionResult(
                tables={},
                insert_order=insert_order,
                stats=dry_run_stats,
                traversal_path=all_paths,
                has_cycles=bool(cycle_infos),
                broken_fks=broken_fks,
                deferred_updates=[],
                cycle_infos=cycle_infos,
                used_deferred_cycle_strategy=used_deferred_cycle_strategy,
            )

        total_rows_estimate = sum(len(pks) for pks in all_records.values())
        use_streaming = self._should_use_streaming(total_rows_estimate)

        if use_streaming:
            if not self.config.output_file:
                raise ValueError("Streaming mode requires --out-file to be specified")

            logger.info(
                "Using streaming mode",
                estimated_rows=total_rows_estimate,
                threshold=self.config.streaming_threshold,
                output_file=self.config.output_file,
            )
            self._log(
                "stream",
                f"Using streaming mode for {total_rows_estimate} rows (threshold: {self.config.streaming_threshold})",
            )

            return self._do_streaming_extract(
                db_type,
                all_records,
                insert_order,
                broken_fks,
                cycle_infos,
                all_paths,
                used_deferred_cycle_strategy,
            )

        logger.info(
            "Using in-memory mode",
            estimated_rows=total_rows_estimate,
            threshold=self.config.streaming_threshold,
        )

        # Fetch all row data
        self._log("fetch", f"Fetching data from {len(all_records)} tables...")
        logger.info("Starting data fetch phase", table_count=len(all_records))

        tables_data: dict[str, list[dict[str, Any]]] = {}
        stats: dict[str, int] = {}

        total_tables = len(all_records)
        for i, (table, pk_values) in enumerate(all_records.items()):
            if not pk_values:
                continue

            table_info = self.schema.get_table(table)
            if not table_info:
                continue

            self._log("fetch", f"Fetching {len(pk_values)} rows from {table}", i + 1, total_tables)

            pk_columns = table_info.primary_key
            with logger.timed_operation(
                "fetch_table_data",
                table=table,
                row_count=len(pk_values),
            ):
                rows = list(self.adapter.fetch_by_pk(table, pk_columns, pk_values))

            # Anonymize if enabled
            if self.anonymizer:
                with logger.timed_operation("anonymize_table_data", table=table):
                    rows = self._anonymize_table_data(table, rows)
                logger.debug("Table data anonymized", table=table, row_count=len(rows))

            tables_data[table] = rows
            stats[table] = len(rows)
            logger.debug("Table data fetched", table=table, row_count=len(rows))

        if self._has_row_limits():
            self._log("limits", "Applying deterministic row limits with integrity closure...")
            with logger.timed_operation("apply_row_limits"):
                tables_data = self._apply_row_limits(tables_data)
            stats = {table: len(rows) for table, rows in tables_data.items()}
            logger.info(
                "Row limits applied",
                global_limit=self.config.row_limit_global,
                per_table_limits=len(self.config.row_limit_per_table),
                total_rows=sum(stats.values()),
            )
            self._log("limits", "Row limits applied")

        deferred_updates = []
        if broken_fks:
            from dbslice.core.cycles import build_deferred_updates

            self._log("cycles", f"Breaking {len(broken_fks)} circular reference(s)...")
            with logger.timed_operation("build_deferred_updates"):
                deferred_updates = build_deferred_updates(broken_fks, tables_data, self.schema)

            logger.info(
                "Circular references resolved",
                broken_fks=len(broken_fks),
                deferred_updates=len(deferred_updates),
            )
            self._log("cycles", f"Generated {len(deferred_updates)} deferred UPDATE(s)")

        total_rows = sum(len(rows) for rows in tables_data.values())
        self._log(
            "complete", f"Extraction complete: {total_rows} rows from {len(tables_data)} tables"
        )

        validation_result = None
        if self.config.validate:
            self._log("validate", "Validating extraction for referential integrity...")
            validator = ExtractionValidator(self.schema)

            with logger.timed_operation("validate_extraction"):
                validation_result = validator.validate(tables_data, broken_fks)

            if validation_result.is_valid:
                logger.info(
                    "Validation passed",
                    records_checked=validation_result.total_records_checked,
                    fk_checks=validation_result.total_fk_checks,
                )
                self._log("validate", "Validation passed: all FK references are intact")
            else:
                logger.warning(
                    "Validation failed",
                    orphaned_records=len(validation_result.orphaned_records),
                    records_checked=validation_result.total_records_checked,
                )
                self._log(
                    "validate",
                    f"Validation failed: {len(validation_result.orphaned_records)} orphaned record(s) found",
                )

                if self.config.fail_on_validation_error:
                    from dbslice.exceptions import ExtractionError

                    error_msg = (
                        f"Extraction validation failed: {len(validation_result.orphaned_records)} "
                        f"orphaned record(s) detected. These records have foreign key references to "
                        f"parent records that are not included in the extraction.\n\n"
                        f"{validation_result.format_report()}"
                    )
                    raise ExtractionError(error_msg)

        return ExtractionResult(
            tables=tables_data,
            insert_order=insert_order,
            stats=stats,
            traversal_path=all_paths,
            has_cycles=bool(cycle_infos),
            broken_fks=broken_fks,
            deferred_updates=deferred_updates,
            cycle_infos=cycle_infos,
            validation_result=validation_result,
            used_deferred_cycle_strategy=used_deferred_cycle_strategy,
        )

    def _anonymize_table_data(self, table: str, rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
        """
        Anonymize sensitive fields in extracted rows.

        This is called after data fetching but before SQL generation,
        ensuring anonymization happens once and consistently.

        Args:
            table: Table name
            rows: List of row dictionaries

        Returns:
            List of anonymized row dictionaries
        """
        if not self.anonymizer:
            return rows

        return [self.anonymizer.anonymize_row(table, row) for row in rows]

    def _has_row_limits(self) -> bool:
        """Check whether any row-limit configuration is active."""
        return self.config.row_limit_global is not None or bool(self.config.row_limit_per_table)

    def _row_limit_for_table(self, table: str) -> int | None:
        """Resolve effective row limit for a table."""
        if table in self.config.row_limit_per_table:
            return self.config.row_limit_per_table[table]
        return self.config.row_limit_global

    def _row_sort_key(self, table: str, row: dict[str, Any]) -> tuple[Any, ...]:
        """Build a deterministic sort key for row selection."""
        assert self.schema is not None
        table_info = self.schema.get_table(table)
        if table_info and table_info.primary_key:
            return tuple((row.get(col) is None, repr(row.get(col))) for col in table_info.primary_key)
        # Fallback for tables without primary key
        return tuple((k, repr(v)) for k, v in sorted(row.items()))

    def _row_identity(
        self, table: str, row: dict[str, Any], index: int
    ) -> tuple[Any, ...]:
        """Build a stable identity key for selected/candidate rows."""
        assert self.schema is not None
        table_info = self.schema.get_table(table)
        if table_info and table_info.primary_key:
            return ("pk",) + tuple(row.get(col) for col in table_info.primary_key)
        return ("idx", index)

    def _apply_row_limits(
        self, tables_data: dict[str, list[dict[str, Any]]]
    ) -> dict[str, list[dict[str, Any]]]:
        """
        Apply deterministic row limits and expand parent closure to preserve FK integrity.

        Limits are treated as soft caps: required parent rows are added even if
        that exceeds the configured cap.
        """
        assert self.schema is not None
        if not tables_data or not self._has_row_limits():
            return tables_data

        candidate_rows_by_id: dict[str, dict[tuple[Any, ...], dict[str, Any]]] = {}
        selected_ids: dict[str, dict[tuple[Any, ...], None]] = {}
        sorted_rows: dict[str, list[dict[str, Any]]] = {}
        effective_limits: dict[str, int] = {}

        for table, rows in tables_data.items():
            ordered_rows = sorted(rows, key=lambda row: self._row_sort_key(table, row))
            sorted_rows[table] = ordered_rows

            table_candidates: dict[tuple[Any, ...], dict[str, Any]] = {}
            table_selected: dict[tuple[Any, ...], None] = {}
            limit = self._row_limit_for_table(table)
            if limit is not None:
                effective_limits[table] = limit

            for idx, row in enumerate(ordered_rows):
                row_id = self._row_identity(table, row, idx)
                table_candidates[row_id] = row

            selected_count = len(ordered_rows) if limit is None else min(limit, len(ordered_rows))
            for idx in range(selected_count):
                row_id = self._row_identity(table, ordered_rows[idx], idx)
                table_selected[row_id] = None

            candidate_rows_by_id[table] = table_candidates
            selected_ids[table] = table_selected

        lookup_cache: dict[tuple[str, tuple[str, ...]], dict[tuple[Any, ...], list[tuple[Any, ...]]]] = {}

        def _lookup_for(table: str, columns: tuple[str, ...]) -> dict[tuple[Any, ...], list[tuple[Any, ...]]]:
            cache_key = (table, columns)
            if cache_key in lookup_cache:
                return lookup_cache[cache_key]

            lookup: dict[tuple[Any, ...], list[tuple[Any, ...]]] = {}
            for row_id, row in candidate_rows_by_id.get(table, {}).items():
                key = tuple(row.get(col) for col in columns)
                if None in key:
                    continue
                lookup.setdefault(key, []).append(row_id)
            lookup_cache[cache_key] = lookup
            return lookup

        changed = True
        while changed:
            changed = False
            for source_table in sorted(selected_ids.keys()):
                if source_table not in self.schema.tables:
                    continue
                source_selected_ids = list(selected_ids[source_table].keys())
                if not source_selected_ids:
                    continue

                parent_edges = sorted(
                    self.schema.get_parents(source_table),
                    key=lambda item: (item[0], item[1].name),
                )
                for parent_table, fk in parent_edges:
                    if parent_table not in candidate_rows_by_id:
                        continue
                    assert isinstance(fk, ForeignKey)
                    parent_lookup = _lookup_for(parent_table, fk.target_columns)
                    parent_selected = selected_ids.setdefault(parent_table, {})

                    for source_id in source_selected_ids:
                        source_row = candidate_rows_by_id[source_table][source_id]
                        source_key = tuple(source_row.get(col) for col in fk.source_columns)
                        if None in source_key:
                            continue
                        parent_row_ids = parent_lookup.get(source_key, [])
                        for parent_row_id in parent_row_ids:
                            if parent_row_id not in parent_selected:
                                parent_selected[parent_row_id] = None
                                changed = True

        limited_tables: dict[str, list[dict[str, Any]]] = {}
        for table in tables_data.keys():
            table_selected = selected_ids.get(table, {})
            limited_rows = [candidate_rows_by_id[table][row_id] for row_id in table_selected.keys()]
            limited_tables[table] = limited_rows

            limit = effective_limits.get(table)
            if limit is not None and len(limited_rows) > limit:
                logger.info(
                    "Row limit soft cap exceeded to preserve FK integrity",
                    table=table,
                    configured_limit=limit,
                    final_rows=len(limited_rows),
                )

        return limited_tables

    def _process_seed(self, seed: SeedSpec) -> TraversalResult:
        """Process a single seed and return traversal result."""
        assert self.schema is not None
        assert self.adapter is not None

        if not self.schema.has_table(seed.table):
            raise TableNotFoundError(seed.table, self.schema.get_table_names())

        table_info = self.schema.get_table(seed.table)
        assert table_info is not None  # Guaranteed by has_table check above
        pk_columns = table_info.primary_key
        if not pk_columns:
            raise ExtractionError(
                "Seed table has no primary key; choose a seed table with a primary key",
                table=seed.table,
            )

        where_clause, params = seed.to_where_clause(
            allow_unsafe_subqueries=self.config.allow_unsafe_where
        )
        seed_rows = list(self.adapter.fetch_rows(seed.table, where_clause, params))

        if not seed_rows:
            seed_str = (
                f"{seed.table}.{seed.column}={seed.value}"
                if seed.column
                else f"{seed.table}:{seed.where_clause}"
            )
            raise NoRowsFoundError(seed_str, seed.table)

        seed_pks: set[tuple[Any, ...]] = set()
        for row in seed_rows:
            pk_tuple = tuple(row[col] for col in pk_columns)
            seed_pks.add(pk_tuple)

        traversal_config = TraversalConfig(
            max_depth=self.config.depth,
            direction=self.config.direction,
            exclude_tables=self.config.exclude_tables,
            passthrough_tables=self.config.passthrough_tables,
            table_depth_overrides=self.config.table_depth_overrides,
            table_direction_overrides=self.config.table_direction_overrides,
        )

        traverser = GraphTraverser(self.schema, self.adapter)
        return traverser.traverse(seed.table, seed_pks, traversal_config)

    def _should_use_streaming(self, total_rows: int) -> bool:
        """
        Determine whether to use streaming mode based on configuration and row count.

        Args:
            total_rows: Estimated total number of rows to extract

        Returns:
            True if streaming mode should be used, False otherwise
        """
        # Force streaming if explicitly enabled
        if self.config.stream:
            if self._has_row_limits():
                logger.warning(
                    "Streaming disabled because row limits are configured and require in-memory integrity closure"
                )
                return False
            logger.debug("Streaming enabled via --stream flag")
            return True

        # Auto-enable streaming if above threshold and output file is specified
        if total_rows >= self.config.streaming_threshold and self.config.output_file:
            if self._has_row_limits():
                logger.warning(
                    "Streaming disabled because row limits are configured and require in-memory integrity closure"
                )
                return False
            logger.debug(
                "Auto-enabling streaming mode",
                total_rows=total_rows,
                threshold=self.config.streaming_threshold,
            )
            return True

        return False

    def _do_streaming_extract(
        self,
        db_type: DatabaseType,
        all_records: dict[str, set[tuple[Any, ...]]],
        insert_order: list[str],
        broken_fks: list[Any],
        cycle_infos: list[Any],
        all_paths: list[str],
        used_deferred_cycle_strategy: bool,
    ) -> ExtractionResult:
        """
        Perform streaming extraction to file.

        This method writes extraction results directly to a file in chunks,
        avoiding loading the entire dataset into memory.
        Perform streaming extraction that writes directly to output file.

        This method is called when the dataset is large enough to warrant
        streaming mode (avoiding OOM errors).

        Args:
            db_type: Database type
            all_records: Map of table -> set of PK tuples to extract
            insert_order: Topologically sorted table names
            broken_fks: Foreign keys broken to handle cycles
            cycle_infos: Information about detected cycles
            all_paths: Traversal paths (for verbose mode)

        Returns:
            ExtractionResult with statistics (but empty tables dict)
        """
        assert self.schema is not None
        assert self.adapter is not None

        from dbslice.core.streaming import StreamingExtractionEngine

        deferred_updates = []
        if broken_fks:
            from dbslice.core.cycles import build_deferred_updates

            self._log("cycles", f"Breaking {len(broken_fks)} circular reference(s)...")

            # For streaming mode, we need to build deferred updates without having
            # all data in memory. We fetch the necessary data on-demand.
            with logger.timed_operation("build_deferred_updates_streaming"):
                temp_data = {}
                for fk in broken_fks:
                    table = fk.source_table
                    if table not in temp_data and table in all_records:
                        pk_values = all_records[table]
                        table_info = self.schema.get_table(table)
                        if table_info:
                            pk_columns = table_info.primary_key
                            rows = list(self.adapter.fetch_by_pk(table, pk_columns, pk_values))
                            temp_data[table] = rows

                deferred_updates = build_deferred_updates(broken_fks, temp_data, self.schema)

            logger.info(
                "Circular references resolved",
                broken_fks=len(broken_fks),
                deferred_updates=len(deferred_updates),
            )
            self._log("cycles", f"Generated {len(deferred_updates)} deferred UPDATE(s)")

        streaming_engine = StreamingExtractionEngine(
            config=self.config,
            adapter=self.adapter,
            schema=self.schema,
            records=all_records,
            insert_order=insert_order,
            broken_fks=broken_fks,
            deferred_updates=deferred_updates,
            db_type=db_type,
            progress_callback=self.progress_callback,
            chunk_size=self.config.streaming_chunk_size,
        )

        assert self.config.output_file is not None
        result = streaming_engine.stream_to_file(self.config.output_file)

        result.traversal_path = all_paths
        result.cycle_infos = cycle_infos
        result.has_cycles = bool(cycle_infos)
        result.used_deferred_cycle_strategy = used_deferred_cycle_strategy

        return result

    def _topological_sort(
        self, tables: set[str], db_type: DatabaseType
    ) -> tuple[list[str], list[Any], list[Any], bool]:
        """
        Topologically sort tables based on FK dependencies with cycle handling.

        Tables with no dependencies come first (parents before children).
        If cycles are detected, breaks them at nullable foreign keys.

        Returns:
            Tuple of (insert_order, broken_fks, cycle_infos, used_deferred_cycle_strategy)
        """
        assert self.schema is not None

        from dbslice.core.cycles import (
            CycleInfo,
            break_cycles_at_nullable_fks,
            find_cycles_dfs,
            identify_cycle_fks,
        )

        dependencies: dict[str, set[str]] = {t: set() for t in tables}

        for fk in self.schema.edges:
            if fk.source_table in tables and fk.target_table in tables:
                dependencies[fk.source_table].add(fk.target_table)

        ts = TopologicalSorter(dependencies)

        try:
            insert_order = list(ts.static_order())
            return insert_order, [], [], False
        except CycleError:
            try:
                fks_to_break, cycle_infos = break_cycles_at_nullable_fks(
                    self.schema, tables, dependencies
                )
            except ValueError as e:
                cycles = find_cycles_dfs(dependencies)
                cycle_infos = [
                    CycleInfo(
                        tables=cycle,
                        fks_in_cycle=identify_cycle_fks(self.schema, cycle),
                    )
                    for cycle in cycles
                ]
                non_deferrable: list[str] = []

                # For JSON/CSV output, deterministic ordering is enough.
                if self.config.output_format != OutputFormat.SQL:
                    logger.warning(
                        "Cycle detected with no nullable FK; using deterministic output order for non-SQL format",
                        cycle_count=len(cycle_infos),
                        format=self.config.output_format.value,
                    )
                    return sorted(tables), [], cycle_infos, False

                # SQL output can fallback only when explicitly requested and feasible.
                if (
                    self.config.disable_fk_checks
                    and db_type == DatabaseType.POSTGRESQL
                    and cycle_infos
                ):
                    for cycle_info in cycle_infos:
                        for fk in cycle_info.fks_in_cycle:
                            if not fk.is_deferrable:
                                non_deferrable.append(
                                    f"{fk.source_table}.{', '.join(fk.source_columns)} -> "
                                    f"{fk.target_table}.{', '.join(fk.target_columns)}"
                                )
                    if not non_deferrable:
                        logger.warning(
                            "Cycle detected with no nullable FK; using deferred-constraint strategy",
                            cycle_count=len(cycle_infos),
                        )
                        return sorted(tables), [], cycle_infos, True

                from dbslice.exceptions import CircularReferenceError

                if self.config.disable_fk_checks and db_type == DatabaseType.POSTGRESQL:
                    if non_deferrable:
                        non_deferrable_msg = "\n".join(f"  - {item}" for item in non_deferrable[:10])
                        if len(non_deferrable) > 10:
                            non_deferrable_msg += f"\n  ... and {len(non_deferrable) - 10} more"
                    elif cycle_infos:
                        non_deferrable_msg = (
                            "  (cycle details available, but deferrability could not be determined)"
                        )
                    else:
                        non_deferrable_msg = "  (no cycle details available)"
                    raise CircularReferenceError(
                        "Circular dependency requires deferrable constraints for SQL fallback.\n\n"
                        "Non-deferrable FKs in detected cycles:\n"
                        f"{non_deferrable_msg}\n\n"
                        "Make at least one FK nullable (preferred), or change cycle FKs to DEFERRABLE."
                    )

                raise CircularReferenceError(str(e))

            modified_deps = {t: set(deps) for t, deps in dependencies.items()}
            for fk in fks_to_break:
                if fk.source_table in modified_deps:
                    modified_deps[fk.source_table].discard(fk.target_table)

            ts = TopologicalSorter(modified_deps)
            insert_order = list(ts.static_order())

            return insert_order, fks_to_break, cycle_infos, False


def extract_subset(
    database_url: str,
    seed: str,
    depth: int = DEFAULT_TRAVERSAL_DEPTH,
    direction: str = "up",
) -> str:
    """
    Convenience function to extract a database subset as SQL.

    Args:
        database_url: Database connection URL
        seed: Seed specification (e.g., "orders.id=12345")
        depth: Maximum FK traversal depth
        direction: Traversal direction ("up", "down", "both")

    Returns:
        SQL string with INSERT statements
    """
    seed_spec = SeedSpec.parse(seed)
    direction_enum = TraversalDirection(direction)
    config = ExtractConfig(
        database_url=database_url,
        seeds=[seed_spec],
        depth=depth,
        direction=direction_enum,
    )

    engine = ExtractionEngine(config)
    result, schema = engine.extract()

    db_config = parse_database_url(database_url)
    generator = SQLGenerator(db_type=db_config.db_type)
    return generator.generate(
        result.tables,
        result.insert_order,
        schema.tables,
    )
