import logging
import asyncio
from dataclasses import dataclass
from typing import Union, Optional, List
from pathlib import Path

from atomsql import Database

from atomingest.core.loader import SchemaLoader
from atomingest.normalization.integration import normalize_then_validate
from atomingest.normalization.core import NormalizationConfig
from atomingest.pipeline.hooks import (
    HookRegistry,
    HookContext,
    HookType,
    create_registry_from_config,
)
from atomingest.pipeline.audit_metadata import (
    generate_run_id,
    MetadataInjectionHook,
)

logger = logging.getLogger(__name__)


@dataclass
class IngestionStats:
    """Statistics for the ingestion run"""

    processed: int = 0
    failed: int = 0
    total_rows: int = 0


class Ingester:
    """The main runner for data ingestion pipelines
    Ties together the schema loading, model generation, normalization, and persistence.
    """

    def __init__(self, db: Database, hook_registry: Optional[HookRegistry] = None):
        """Initialize the Ingester with a database connection.
        Args:
            db (Database): The database connection to use for ingestion.
            hook_registry (Optional[HookRegistry]): Optional pre-configured hook registry.
                If not provided, hooks will be loaded from the YAML config.
        """
        self.db = db
        self.schema_loader = SchemaLoader()
        self._hook_registry = hook_registry

    async def run_async(
        self, config: Union[Path, str], source: Union[Path, str]
    ) -> IngestionStats:
        """
        Execute the ingestion pipeline asynchronously.
        Wraps the synchronous run method in a thread to prevent blocking the event loop.

        Args:
            config: Path to the YAML configuration file.
            source: Path to the source data file or directory.

        Returns:
            IngestionStats: Statistics about the ingestion run.
        """
        logger.info("Starting async ingestion wrapper...")
        return await asyncio.to_thread(self.run, config, source)

    def run(self, config: Union[Path, str], source: Union[Path, str]) -> IngestionStats:
        """
        Execute the ingestion pipeline
        Args:
            config: Path to the YAML configuration file.
            source: Path to the source data file or directory.

        Returns:
            IngestionStats: Statistics about the ingestion run.
        """
        config_path = str(config)
        source_path = str(source)

        # Generate unique run ID for this ingestion
        run_id = generate_run_id()
        logger.info(
            f"Starting ingestion with config: {config_path} and source: {source_path} "
            f"[run_id: {run_id}]"
        )

        stats = IngestionStats()
        # Load the schema
        schema = self.schema_loader.load_from_yaml(config_path)
        # Generate the model
        ModelClass = self.schema_loader.create_model(schema)

        # schema["config"] is the full raw YAML dict; the user's pipeline settings
        # live under the nested "config" key within it.
        raw_config = schema.get("config", {})
        # pipeline_config: nested "config:" block in YAML (for batch_size, audit_metadata, etc.)
        pipeline_config = raw_config.get("config", {})
        batch_size = pipeline_config.get("batch_size", 100)
        # business_key is a top-level YAML key, not nested under config.config
        business_keys = raw_config.get("business_key", [])

        self.db.register(ModelClass)

        # Set up hook registry - use provided registry or create from config
        hook_registry = self._hook_registry
        if hook_registry is None:
            hooks_config = schema.get("hooks")
            if hooks_config:
                hook_registry = create_registry_from_config(hooks_config)
            else:
                hook_registry = HookRegistry()

        # Check if audit metadata injection is configured
        audit_config = pipeline_config.get("audit_metadata")
        if audit_config and audit_config.get("enabled", True):
            # Automatically add MetadataInjectionHook
            metadata_hook = MetadataInjectionHook(
                name="audit_metadata_injection", config=audit_config
            )
            hook_registry.register(HookType.PRE_SAVE, metadata_hook, order=-1000)
            logger.info("Audit metadata injection enabled")

        # The loader returns "normalization" with NormalizationConfig already parsed
        # If not available, create an empty config
        normalisation_config = schema.get("normalization")
        if normalisation_config is None:
            normalisation_config = NormalizationConfig.from_dict({})
        try:
            rows, report = normalize_then_validate(
                source=source_path, config=normalisation_config
            )
            stats.total_rows = report.rows_emitted
            logger.info(
                f"Normalization and validation completed. Rows emitted: {stats.total_rows}"
            )
        except Exception as e:
            logger.error(f"Normalization and validation failed: {e}")
            raise e

        logger.info(
            f"Beginning database insertion for table '{schema['target_table']}' "
            f"with batch size {batch_size}..."
        )

        batch: List[HookContext] = []

        def process_batch(current_batch: List[HookContext]):
            """
            Process a batch of rows.
            Attempts a bulk transaction first. If that fails, falls back to
            row-by-row processing to isolate errors.
            """
            if not current_batch:
                return

            try:
                # Optimized Path: Batched Transaction
                with self.db.transaction():
                    for ctx in current_batch:
                        self._persist_row(ctx, ModelClass, business_keys)
                        # Update stats and run post-save hooks only after successful save
                        stats.processed += 1
                        ctx.instance = getattr(ctx, "instance", None)
                        hook_registry.execute_hooks(HookType.POST_SAVE, ctx)

                logger.debug(f"Committed batch of {len(current_batch)} rows")

            except Exception as batch_error:
                logger.warning(
                    f"Batch transaction failed: {batch_error}. "
                    "Falling back to individual row processing."
                )
                # Fallback Path: Row-by-Row
                for ctx in current_batch:
                    try:
                        with self.db.transaction():
                            self._persist_row(ctx, ModelClass, business_keys)
                            stats.processed += 1
                            ctx.instance = getattr(ctx, "instance", None)
                            hook_registry.execute_hooks(HookType.POST_SAVE, ctx)
                    except Exception as row_error:
                        logger.error(
                            f"Failed to insert row {ctx.row_data}: Error {row_error}"
                        )
                        stats.failed += 1

                        # Execute on-error hooks (e.g., DLQ)
                        error_context = HookContext(
                            row_data=ctx.row_data,
                            error=row_error,
                            metadata={
                                "table": schema["target_table"],
                                "source": source_path,
                                "run_id": run_id,
                                "db": self.db,
                            },
                        )
                        hook_registry.execute_hooks(HookType.ON_ERROR, error_context)

        # Main Ingestion Loop
        for row_data in rows:
            try:
                # Create hook context for this row
                context = HookContext(
                    row_data=row_data,
                    metadata={
                        "table": schema["target_table"],
                        "source": source_path,
                        "run_id": run_id,
                        "db": self.db,
                    },
                )

                # Execute pre-save hooks
                context = hook_registry.execute_hooks(HookType.PRE_SAVE, context)

                # Check if hook requested to skip this row
                if context.skip_row:
                    logger.debug(f"Row skipped by pre-save hook: {row_data}")
                    continue

                # Check if hook requested to abort pipeline
                if context.abort:
                    logger.warning("Pipeline aborted by pre-save hook")
                    # Process pending batch before aborting?
                    # Usually abort means stop immediately.
                    break

                # Add to batch
                batch.append(context)

                # Process batch if full
                if len(batch) >= batch_size:
                    process_batch(batch)
                    batch = []

            except Exception as e:
                # This catches errors in pre-save hooks or context creation
                logger.error(f"Error preparing row {row_data}: {e}")
                stats.failed += 1
                error_context = HookContext(
                    row_data=row_data,
                    error=e,
                    metadata={
                        "table": schema["target_table"],
                        "source": source_path,
                        "run_id": run_id,
                        "db": self.db,
                    },
                )
                hook_registry.execute_hooks(HookType.ON_ERROR, error_context)

        # Process remaining rows in batch
        process_batch(batch)

        # Execute on-complete hooks
        complete_context = HookContext(
            row_data={},
            metadata={
                "table": schema["target_table"],
                "source": source_path,
                "run_id": run_id,
                "stats": {
                    "processed": stats.processed,
                    "failed": stats.failed,
                    "total_rows": stats.total_rows,
                },
            },
        )
        hook_registry.execute_hooks(HookType.ON_COMPLETE, complete_context)

        logger.info(
            f"Database insertion completed. Processed: {stats.processed}, Failed: {stats.failed}"
        )
        return stats

    def _persist_row(self, context, ModelClass, business_keys):
        """Helper to persist a single row, handling upserts if business keys exist."""
        if business_keys:
            # Idempotent Ingestion Logic
            try:
                lookup_kwargs = {k: context.row_data[k] for k in business_keys}
            except KeyError as missing_key:
                raise ValueError(
                    f"Row data missing required business key: {missing_key}"
                )

            instance, created = ModelClass.update_or_create(
                defaults=context.row_data, **lookup_kwargs
            )
            context.instance = instance
        else:
            # Standard append-only behavior
            instance = ModelClass(**context.row_data)
            instance.save()
            context.instance = instance
