"""
Core profiling engine for Baselinr.

Orchestrates the profiling of database tables and columns,
collecting schema information and computing metrics.
"""

import logging
import time
from datetime import datetime
from typing import Any, Dict, List, Optional

from ..config.schema import BaselinrConfig, TablePattern
from ..connectors.base import BaseConnector
from ..connectors.factory import create_connector
from ..events import EventBus, ProfilingCompleted, ProfilingFailed, ProfilingStarted
from .metrics import MetricCalculator
from .query_builder import QueryBuilder

logger = logging.getLogger(__name__)


class ProfilingResult:
    """Container for profiling results."""

    def __init__(
        self, run_id: str, dataset_name: str, schema_name: Optional[str], profiled_at: datetime
    ):
        """
        Initialize profiling result container.

        Args:
            run_id: Unique identifier for this profiling run
            dataset_name: Name of the dataset/table profiled
            schema_name: Schema name (if applicable)
            profiled_at: Timestamp of profiling
        """
        self.run_id = run_id
        self.dataset_name = dataset_name
        self.schema_name = schema_name
        self.profiled_at = profiled_at
        self.columns: List[Dict[str, Any]] = []
        self.metadata: Dict[str, Any] = {}

    def add_column_metrics(self, column_name: str, column_type: str, metrics: Dict[str, Any]):
        """
        Add metrics for a column.

        Args:
            column_name: Name of the column
            column_type: Data type of the column
            metrics: Dictionary of metric_name -> metric_value
        """
        self.columns.append(
            {"column_name": column_name, "column_type": column_type, "metrics": metrics}
        )

    def to_dict(self) -> Dict[str, Any]:
        """Convert result to dictionary."""
        return {
            "run_id": self.run_id,
            "dataset_name": self.dataset_name,
            "schema_name": self.schema_name,
            "profiled_at": self.profiled_at.isoformat(),
            "columns": self.columns,
            "metadata": self.metadata,
        }


class ProfileEngine:
    """Main profiling engine for Baselinr."""

    def __init__(
        self,
        config: BaselinrConfig,
        event_bus: Optional[EventBus] = None,
        run_context: Optional[Any] = None,
    ):
        """
        Initialize profiling engine.

        Args:
            config: Baselinr configuration
            event_bus: Optional event bus for emitting profiling events
            run_context: Optional run context with logger and run_id
        """
        self.config = config
        self.connector: Optional[BaseConnector] = None
        self.metric_calculator: Optional[MetricCalculator] = None
        self.event_bus = event_bus
        self.run_context = run_context

        # Get logger from run_context or create fallback
        if run_context:
            self.logger = run_context.logger
            self.run_id = run_context.run_id
        else:
            import logging

            self.logger = logging.getLogger(__name__)
            import uuid

            self.run_id = str(uuid.uuid4())

        # Initialize worker pool ONLY if parallelism is enabled
        self.execution_config = config.execution
        self.worker_pool: Optional[Any] = None
        self.progress_callback: Optional[Any] = None

        # Only create worker pool if max_workers > 1
        if self.execution_config.max_workers > 1:
            # Determine warehouse-specific worker limit
            warehouse_limit = self.execution_config.warehouse_limits.get(
                self.config.source.type, self.execution_config.max_workers
            )

            # Special handling for SQLite (single writer)
            if self.config.source.type == "sqlite":
                warehouse_limit = 1  # SQLite doesn't support concurrent writes well
                self.logger.warning(
                    "SQLite does not support parallel writes. Using sequential execution."
                )

            if warehouse_limit > 1:
                from ..utils.worker_pool import WorkerPool

                self.worker_pool = WorkerPool(
                    max_workers=warehouse_limit,
                    queue_size=self.execution_config.queue_size,
                    warehouse_type=self.config.source.type,
                )
                self.logger.info(f"Parallel execution enabled with {warehouse_limit} workers")
        else:
            from ..utils.logging import log_event

            log_event(
                self.logger,
                "execution_mode",
                "Sequential execution (max_workers=1, default)",
                level="debug",
            )

    def profile(
        self,
        table_patterns: Optional[List[TablePattern]] = None,
        progress_callback: Optional[Any] = None,
    ) -> List[ProfilingResult]:
        """
        Profile tables with optional parallel execution.

        If max_workers=1 (default), uses sequential execution (existing behavior).
        If max_workers > 1, uses parallel execution via worker pool.

        Args:
            table_patterns: Optional list of table patterns to profile
                          (uses config if not provided)
            progress_callback: Optional callback function(current, total, table_name)
                             called when starting each table

        Returns:
            List of profiling results
        """
        patterns = table_patterns or self.config.profiling.tables

        if not patterns:
            logger.warning("No table patterns specified for profiling")
            return []

        # Create connector with retry config
        retry_config = self.config.retry
        execution_config = self.config.execution

        self.connector = create_connector(
            self.config.source, retry_config=retry_config, execution_config=execution_config
        )

        # Create query builder for partition/sampling support
        self.query_builder = QueryBuilder(database_type=self.config.source.type)

        # Create metric calculator
        self.metric_calculator = MetricCalculator(
            engine=self.connector.engine,
            max_distinct_values=self.config.profiling.max_distinct_values,
            compute_histograms=self.config.profiling.compute_histograms,
            histogram_bins=self.config.profiling.histogram_bins,
            enabled_metrics=self.config.profiling.metrics,
            query_builder=self.query_builder,
            enable_enrichment=self.config.profiling.enable_enrichment,
            enable_approx_distinct=self.config.profiling.enable_approx_distinct,
            enable_type_inference=self.config.profiling.enable_type_inference,
            type_inference_sample_size=self.config.profiling.type_inference_sample_size,
        )

        # Store progress callback for use in _profile_sequential and _profile_parallel
        self.progress_callback = progress_callback

        # Route to parallel or sequential execution
        try:
            if self.worker_pool:
                return self._profile_parallel(patterns)
            else:
                return self._profile_sequential(patterns)
        finally:
            # Cleanup
            if self.connector:
                self.connector.close()
            if self.worker_pool:
                self.worker_pool.shutdown(wait=True)

    def _profile_parallel(self, patterns: List[TablePattern]) -> List[ProfilingResult]:
        """
        Profile tables in parallel using worker pool.
        Only called when max_workers > 1.

        Args:
            patterns: List of table patterns to profile

        Returns:
            List of profiling results
        """
        from ..utils.worker_pool import profile_table_task

        if self.worker_pool is None:
            raise RuntimeError("Worker pool is not initialized")

        # Submit all tasks
        futures = []
        for pattern in patterns:
            future = self.worker_pool.submit(
                profile_table_task,
                self,  # Pass engine instance
                pattern,
                self.run_context,
                self.event_bus,
            )
            futures.append(future)

        # Wait for completion
        results = self.worker_pool.wait_for_completion(futures)

        # Filter out None results (failed tasks)
        successful = [r for r in results if r is not None]
        failed_count = len(results) - len(successful)

        if failed_count > 0:
            logger.warning(
                f"Parallel profiling completed: {len(successful)} succeeded, {failed_count} failed"
            )

        return successful

    def _profile_sequential(self, patterns: List[TablePattern]) -> List[ProfilingResult]:
        """
        Profile tables sequentially (existing implementation).
        This is the default behavior when max_workers=1.

        Args:
            patterns: List of table patterns to profile

        Returns:
            List of profiling results
        """
        results = []
        warehouse = self.config.source.type  # Get warehouse type for metrics
        total = len(patterns)

        for idx, pattern in enumerate(patterns):
            assert pattern.table is not None, "Table name must be set"
            fq_table = f"{pattern.schema_}.{pattern.table}" if pattern.schema_ else pattern.table
            start_time = time.time()

            # Call progress callback if provided
            if self.progress_callback:
                try:
                    self.progress_callback(idx + 1, total, fq_table)
                except Exception as e:
                    logger.debug(f"Progress callback error: {e}")

            try:
                from ..utils.logging import log_and_emit

                # Record metrics: profiling started
                if self.run_context and self.run_context.metrics_enabled:
                    from ..utils.metrics import record_profile_started

                    assert pattern.table is not None, "Table name must be set"
                    record_profile_started(warehouse, fq_table)

                # Log and emit profiling started
                log_and_emit(
                    self.logger,
                    self.event_bus,
                    "profiling_started",
                    f"Starting profiling for table: {fq_table}",
                    table=fq_table,
                    run_id=self.run_id,
                )

                result = self._profile_table(pattern)
                results.append(result)

                # Calculate duration
                duration = time.time() - start_time

                # Record metrics: profiling completed
                if self.run_context and self.run_context.metrics_enabled:
                    from ..utils.metrics import record_profile_completed

                    record_profile_completed(
                        warehouse,
                        fq_table,
                        duration,
                        row_count=result.metadata.get("row_count", 0),
                        column_count=len(result.columns),
                    )

                # Log and emit profiling completed
                log_and_emit(
                    self.logger,
                    self.event_bus,
                    "profiling_completed",
                    f"Profiling completed for table: {fq_table}",
                    table=fq_table,
                    run_id=self.run_id,
                    metadata={
                        "column_count": len(result.columns),
                        "row_count": result.metadata.get("row_count", 0),
                    },
                )
            except Exception as e:
                from ..utils.logging import log_and_emit

                # Calculate duration
                duration = time.time() - start_time

                # Record metrics: profiling failed
                if self.run_context and self.run_context.metrics_enabled:
                    from ..utils.metrics import record_profile_failed

                    record_profile_failed(warehouse, fq_table, duration)

                # Log and emit failure
                log_and_emit(
                    self.logger,
                    self.event_bus,
                    "profiling_error",
                    f"Failed to profile table {fq_table}: {e}",
                    level="error",
                    table=fq_table,
                    run_id=self.run_id,
                    metadata={"error": str(e), "error_type": type(e).__name__},
                )

                # Continue processing other tables instead of aborting
                logger.warning(f"Continuing with remaining tables after failure on {fq_table}")

        return results

    def _profile_table(self, pattern: TablePattern) -> ProfilingResult:
        """
        Profile a single table.

        Args:
            pattern: Table pattern configuration

        Returns:
            ProfilingResult for this table
        """
        # Ensure table name is set (should be after pattern expansion)
        assert pattern.table is not None, "Table name must be set for profiling"

        # Use run_id from context
        run_id = self.run_id
        profiled_at = datetime.utcnow()
        start_time = time.time()

        fq_table = f"{pattern.schema_}.{pattern.table}" if pattern.schema_ else pattern.table
        from ..utils.logging import log_event

        log_event(
            self.logger,
            "table_profiling_started",
            f"Profiling table: {fq_table}",
            table=fq_table,
            metadata={"pattern": pattern.table},
        )

        # Emit profiling started event
        if self.event_bus:
            self.event_bus.emit(
                ProfilingStarted(
                    event_type="ProfilingStarted",
                    timestamp=profiled_at,
                    table=pattern.table,
                    run_id=run_id,
                    metadata={},
                )
            )

        try:
            # Get table metadata
            if self.connector is None:
                raise RuntimeError("Connector is not initialized")
            table = self.connector.get_table(pattern.table, schema=pattern.schema_)

            # Create result container
            result = ProfilingResult(
                run_id=run_id,
                dataset_name=pattern.table,
                schema_name=pattern.schema_,
                profiled_at=profiled_at,
            )

            # Infer partition key if metadata_fallback is enabled
            partition_config = pattern.partition
            if partition_config and partition_config.metadata_fallback and not partition_config.key:
                inferred_key = self.query_builder.infer_partition_key(table)
                if inferred_key:
                    partition_config.key = inferred_key
                    logger.info(f"Using inferred partition key: {inferred_key}")

            # Add table metadata
            current_row_count = self._get_row_count(table, partition_config, pattern.sampling)
            result.metadata["row_count"] = current_row_count
            result.metadata["column_count"] = len(table.columns)
            result.metadata["partition_config"] = (
                partition_config.model_dump() if partition_config else None
            )
            result.metadata["sampling_config"] = (
                pattern.sampling.model_dump() if pattern.sampling else None
            )

            # Profile each column
            for column in table.columns:
                logger.debug(f"Profiling column: {column.name}")

                try:
                    if self.metric_calculator is None:
                        raise RuntimeError("Metric calculator is not initialized")
                    metrics = self.metric_calculator.calculate_all_metrics(
                        table=table,
                        column_name=column.name,
                        partition_config=partition_config,
                        sampling_config=pattern.sampling,
                    )

                    result.add_column_metrics(
                        column_name=column.name, column_type=str(column.type), metrics=metrics
                    )
                except Exception as e:
                    logger.error(f"Failed to profile column {column.name}: {e}")
                    # Add error marker
                    result.add_column_metrics(
                        column_name=column.name,
                        column_type=str(column.type),
                        metrics={"error": str(e)},
                    )

            # Store schema snapshot for enrichment metrics (calculated during storage write)
            if self.config.profiling.enable_enrichment:
                current_columns = {col["column_name"]: col["column_type"] for col in result.columns}
                result.metadata["column_schema"] = current_columns

            # Schema change detection happens in storage writer
            # where we have access to storage engine

            # Calculate duration
            duration = time.time() - start_time

            # Emit profiling completed event
            if self.event_bus:
                self.event_bus.emit(
                    ProfilingCompleted(
                        event_type="ProfilingCompleted",
                        timestamp=datetime.utcnow(),
                        table=pattern.table,
                        run_id=run_id,
                        row_count=result.metadata.get("row_count", 0),
                        column_count=result.metadata.get("column_count", 0),
                        duration_seconds=duration,
                        metadata={},
                    )
                )

            logger.info(
                f"Successfully profiled {pattern.table} with {len(result.columns)} "
                f"columns in {duration:.2f}s"
            )
            return result

        except Exception as e:
            # Emit profiling failed event
            if self.event_bus:
                self.event_bus.emit(
                    ProfilingFailed(
                        event_type="ProfilingFailed",
                        timestamp=datetime.utcnow(),
                        table=pattern.table,
                        run_id=run_id,
                        error=str(e),
                        metadata={},
                    )
                )
            raise

    def _get_row_count(self, table, partition_config=None, sampling_config=None) -> int:
        """
        Get row count for a table (with optional partition/sampling).

        Args:
            table: SQLAlchemy Table object
            partition_config: Partition configuration
            sampling_config: Sampling configuration

        Returns:
            Row count
        """
        from sqlalchemy import func, select

        if self.connector is None:
            raise RuntimeError("Connector is not initialized")
        with self.connector.engine.connect() as conn:
            # Build query with partition filtering
            query, _ = self.query_builder.build_profiling_query(
                table=table,
                partition_config=partition_config,
                sampling_config=None,  # Don't apply sampling for count
            )

            # Count rows
            count_query = select(func.count()).select_from(query.alias())
            result = conn.execute(count_query).scalar()
            return int(result) if result is not None else 0
