# preprocessing/core/pre_orchestrator.py - enhanced
from typing import List
import pandas as pd
import logging
from ..config.preprocess_config import (
    PreprocessConfig,
    PreprocessResult,
    PreprocessMetrics,
)
from .pre_rules import RuleEngine
from .pre_domains import DomainFilter
from .pre_blocking import BlockingKeyGenerator
from ..monitoring.audit_trail import AuditTrail
from ..monitoring.metrics import StageMetrics
import time

log = logging.getLogger(__name__)


class PreprocessOrchestrator:
    """
    Main preprocessing pipeline orchestrator.

    Rule Execution Order:
    1. Profile rules (lowest priority) - base configurations
    2. YAML ruleset rules (medium priority) - project-specific
    3. GUI rules (highest priority) - user overrides

    Within each category, rules execute in order defined.
    Duplicate column transformations are prevented unless
    output_suffix differs.
    """

    def __init__(self):
        self.rule_processor = RuleEngine()
        self.domain_filter = DomainFilter()
        self.blocking_generator = BlockingKeyGenerator()
        self.audit_trail = AuditTrail()

    def process_dataframe(
        self, df: pd.DataFrame, config: PreprocessConfig
    ) -> PreprocessResult:
        """Main entry point for preprocessing"""
        log.info(f"Starting preprocessing pipeline with {len(df)} rows")

        # Initialize result tracking
        result_df = df.copy()
        stage_metrics = []
        total_start = time.time()

        # Stage 1: Rule Processing
        if config.gui_rules or config.yaml_ruleset:
            log.info("Stage 1: Applying preprocessing rules")
            start_time = time.time()
            result_df, rule_metrics = self.rule_processor.apply_rules(
                result_df,
                config.gui_rules
                or [],  # RuleEngine expects a list, not separate params
                self.audit_trail.records,  # Pass audit trail as third parameter
            )
            # rule_metrics is already a StageMetrics object
            stage_metrics.append(rule_metrics)

        # Stage 2: B2C Domain Filtering
        if config.b2c_enabled:
            log.info("Stage 2: Applying B2C domain filtering")

            result_df, domain_metrics = self.domain_filter.apply_domain_filtering(
                result_df, config, self.audit_trail.records
            )
            stage_metrics.append(domain_metrics)

        # Stage 3: Blocking Key Generation
        if config.blocking_enabled and config.block_from_col:
            log.info("Stage 3: Generating blocking keys")
            start_time = time.time()

            result_df, blocking_metrics = (
                self.blocking_generator.generate_blocking_keys(
                    result_df, config.block_from_col, config.block_key_length
                )
            )

            stage_metrics.append(
                StageMetrics(
                    stage_name="Blocking Key Generation",
                    start_time=start_time,
                    end_time=time.time(),
                    rows_modified=len(result_df),
                    columns_modified=["_ACTIVE_BLOCK_KEY_"],
                    rules_applied=1,
                    errors_count=0,
                )
            )

        # Create final metrics object
        final_metrics = PreprocessMetrics(
            start_time=total_start,
            end_time=time.time(),
            total_rows=len(df),
            total_columns=len(df.columns),
            stage_metrics={m.stage_name: m for m in stage_metrics},
        )

        # Create result object
        return PreprocessResult(
            data=result_df, metrics=final_metrics, audit_trail=self.audit_trail.records
        )

    def _merge_rules(
        self, profile_rules: List, yaml_rules: List, gui_rules: List
    ) -> List:
        """Merge rules with proper precedence and deduplication"""
        merged = []
        seen_transforms = set()  # (action, column) pairs

        # Process in priority order
        for rules in [profile_rules, yaml_rules, gui_rules]:
            for rule in rules:
                # Create unique key for this transformation
                for col in rule.get("columns", []):
                    transform_key = (
                        rule["action"],
                        col,
                        rule.get("output_suffix", ""),  # Include suffix in key
                    )

                    if transform_key not in seen_transforms:
                        merged.append(rule)
                        seen_transforms.add(transform_key)
                    else:
                        log.debug(f"Skipping duplicate transform: {transform_key}")

        return merged
