# preprocessing/core/pre_domains.py - REVISED
from typing import Tuple, Optional, List
import pandas as pd
from ..config.preprocess_config import PreprocessConfig, StageMetrics
from ..monitoring.audit_trail import AuditRecord
import time
import logging

log = logging.getLogger(__name__)


class DomainFilter:
    """Integrates with EXISTING B2C domain filtering"""

    def __init__(self):
        # Import the existing B2C module - respecting current implementation
        try:
            from fmatch.core.b2c_engine_integration import cleanse_domain_field

            self.cleanse_domain_field = cleanse_domain_field
            self._available = True
        except ImportError:
            log.warning("B2C engine integration not available")
            self.cleanse_domain_field = None
            self._available = False

    def apply_domain_filtering(
        self,
        df: pd.DataFrame,
        config: PreprocessConfig,
        audit_trail: Optional[List[AuditRecord]] = None,
    ) -> Tuple[pd.DataFrame, StageMetrics]:
        """Apply B2C domain filtering, tracking stats externally"""

        start_time = time.time()
        metrics = StageMetrics(
            stage_name="domain_filtering",
            start_time=start_time,
            end_time=0,
            rows_modified=0,
            columns_modified=[],
            rules_applied=0,
            errors_count=0,
        )

        if not config.b2c_enabled or not self._available:
            metrics.end_time = time.time()
            return df, metrics

        # Detect domain/email columns using SAME keywords as existing implementation
        domain_keywords = ["domain", "website", "url", "site", "web", "email", "mail"]
        domain_columns = [
            col
            for col in df.columns
            if any(keyword in col.lower() for keyword in domain_keywords)
        ]

        if not domain_columns:
            log.info("No domain/email columns detected for B2C filtering")
            metrics.end_time = time.time()
            return df, metrics

        # Find company column for fallback (matching existing logic)
        company_keywords = ["company", "organization", "account", "customer", "name"]
        company_columns = [
            col
            for col in df.columns
            if any(keyword in col.lower() for keyword in company_keywords)
        ]
        company_col = company_columns[0] if company_columns else None

        # Work on copy if tracking changes
        df_result = df.copy() if audit_trail is not None else df

        # Track total stats across all columns
        total_b2c_stripped = 0
        total_fallbacks = 0

        for col in domain_columns:
            if col not in df_result.columns:
                continue

            try:
                # Count non-empty values before
                before_non_empty = df_result[col].notna().sum()
                got_fallback = 0
                before_series = df_result[col].copy() if audit_trail else None

                # Apply existing B2C filtering - returns Series directly!
                df_result[col] = self.cleanse_domain_field(
                    df_result[col],
                    strip_b2c=True,
                    extract_from_email=True,
                    company_name_fallback=df_result[company_col]
                    if company_col
                    else None,
                )

                # Count non-empty values after
                after_non_empty = df_result[col].notna().sum()

                # Calculate stats externally (since function doesn't return them)
                # FIX: Ensure positive number for domains removed
                newly_empty = max(0, before_non_empty - after_non_empty)
                if newly_empty > 0:
                    total_b2c_stripped += newly_empty

                    # Check how many got company fallback
                    if company_col and before_series is not None:
                        # Where B2C was stripped but we have a value now
                        got_fallback = (
                            (before_series.notna())
                            & (before_series != df_result[col])
                            & (df_result[col].notna())
                        ).sum()
                        total_fallbacks += got_fallback

                # Update metrics
                metrics.columns_modified.append(col)

                # Audit trail
                if audit_trail and before_series is not None:
                    changes = (before_series != df_result[col]).sum()
                    if changes > 0:
                        audit_trail.append(
                            AuditRecord(
                                rule_id="b2c_filter",
                                action="strip_b2c_domains",
                                column=col,
                                rows_affected=changes,
                                metadata={
                                    "b2c_stripped": newly_empty,
                                    "fallback_used": got_fallback,
                                },
                            )
                        )

                # FIX: Ensure positive display even if calculation had issues
                domains_removed = max(0, newly_empty)
                log.info(
                    f"B2C filtering on '{col}': {domains_removed} domains stripped"
                )

            except Exception as e:
                metrics.errors_count += 1
                log.warning(f"B2C filtering failed for {col}: {e}")
                # Maintain existing pattern - don't raise, just log

        # Cap total to avoid misleading numbers
        # Total stripped can't exceed the number of non-empty domain cells originally
        original_non_empty_total = sum(
            df[col].notna().sum() for col in domain_columns if col in df.columns
        )
        total_b2c_stripped_capped = min(total_b2c_stripped, original_non_empty_total)

        # Log if capping was applied
        if total_b2c_stripped > total_b2c_stripped_capped:
            log.debug(
                f"B2C metrics capped: raw={total_b2c_stripped}, capped={total_b2c_stripped_capped}"
            )

        # Update final metrics
        metrics.rows_modified = total_b2c_stripped_capped
        metrics.end_time = time.time()

        # Store B2C-specific metrics in parent metrics object
        if hasattr(metrics, "_b2c_stats"):
            metrics._b2c_stats = {
                "domains_filtered": total_b2c_stripped_capped,
                "fallbacks_used": total_fallbacks,
            }

        return df_result, metrics
