# preprocessing/core/pre_rules.py
import pandas as pd
from typing import List, Dict, Any, Optional, Tuple
import logging
from ..rules.rule_schemas import Rule
from ..rules.text_normalizers import TextNormalizers
from ..config.preprocess_config import ErrorMode, StageMetrics
from ..monitoring.audit_trail import AuditRecord
import time

log = logging.getLogger(__name__)


class RuleEngine:
    """Core rule application engine"""

    def __init__(self, error_mode: ErrorMode = ErrorMode.LENIENT):
        self.error_mode = error_mode
        self.normalizers = TextNormalizers()
        self._compiled_rules = {}

    def apply_rules(
        self,
        df: pd.DataFrame,
        rules: List[Dict[str, Any]],
        audit_trail: Optional[List[AuditRecord]] = None,
    ) -> Tuple[pd.DataFrame, StageMetrics]:
        """Apply preprocessing rules to dataframe"""

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

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

        # Validate and compile rules
        validated_rules = self._validate_rules(rules)

        # Apply each rule
        for rule_idx, rule in enumerate(validated_rules):
            try:
                df_result, rule_metrics = self._apply_single_rule(
                    df_result, rule, rule_idx, audit_trail
                )

                # Update metrics
                metrics.rules_applied += 1
                metrics.rows_modified += rule_metrics["rows_modified"]
                metrics.columns_modified.extend(rule_metrics["columns_modified"])

            except Exception as e:
                metrics.errors_count += 1
                self._handle_rule_error(rule, e)

        metrics.end_time = time.time()
        metrics.columns_modified = list(set(metrics.columns_modified))

        return df_result, metrics

    def _validate_rules(self, rules: List[Dict[str, Any]]) -> List[Rule]:
        """Validate and convert rules to schema objects"""
        validated = []

        for rule_dict in rules:
            try:
                validated_rule = Rule(**rule_dict)
                validated.append(validated_rule)
            except Exception as e:
                if self.error_mode == ErrorMode.STRICT:
                    raise ValueError(f"Invalid rule: {rule_dict}") from e
                else:
                    log.warning(f"Skipping invalid rule: {e}")

        return validated

    def _apply_single_rule(
        self,
        df: pd.DataFrame,
        rule: Rule,
        rule_idx: int,
        audit_trail: Optional[List[AuditRecord]],
    ) -> Tuple[pd.DataFrame, Dict[str, Any]]:
        """Apply a single rule and track changes"""

        # Resolve columns
        if rule.columns == "__ALL_STRINGS__":
            from .column_detector import ColumnDetector

            columns = ColumnDetector.detect_string_columns(df)
        else:
            columns = rule.columns

        rule_metrics = {"rows_modified": 0, "columns_modified": []}

        # Apply to each column
        for col in columns:
            if col not in df.columns:
                if not rule.skip_errors:
                    log.warning(f"Column '{col}' not found for rule {rule.action}")
                continue

            # Get before state for audit
            before_series = df[col].copy() if audit_trail is not None else None

            # Apply transformation
            transformed = self._apply_transformation(df[col], rule.action, rule.params)

            # Determine output column
            output_col = col
            if rule.output_suffix:
                output_col = f"{col}{rule.output_suffix}"

            # Update dataframe
            df[output_col] = transformed

            # Track changes
            if audit_trail is not None and before_series is not None:
                changes = (before_series != transformed).sum()
                if changes > 0:
                    audit_trail.append(
                        AuditRecord(
                            rule_id=f"rule_{rule_idx}",
                            action=rule.action,
                            column=col,
                            rows_affected=int(changes),
                            before_sample=str(before_series.iloc[0])
                            if len(before_series) > 0
                            else None,
                            after_sample=str(transformed.iloc[0])
                            if len(transformed) > 0
                            else None,
                        )
                    )
                    rule_metrics["rows_modified"] += changes

            rule_metrics["columns_modified"].append(output_col)

        return df, rule_metrics

    def _apply_transformation(
        self, series: pd.Series, action: str, params: Dict[str, Any]
    ) -> pd.Series:
        """Apply specific transformation to series"""

        # String operations
        if action == "lower":
            return series.str.lower()
        elif action == "upper":
            return series.str.upper()
        elif action == "strip":
            return series.str.strip()
        elif action == "collapse_ws":
            return series.str.replace(r"\s+", " ", regex=True).str.strip()

        # Advanced normalizations
        elif action == "normalize_unicode":
            return series.apply(
                lambda x: self.normalizers.normalize_unicode(str(x))
                if pd.notna(x)
                else x
            )
        elif action == "to_ascii":
            return series.apply(
                lambda x: self.normalizers.smart_to_ascii(str(x)) if pd.notna(x) else x
            )
        elif action == "remove_html":
            return series.apply(
                lambda x: self.normalizers.remove_html(str(x)) if pd.notna(x) else x
            )
        elif action == "normalize_company":
            lang = params.get("language", "en")
            return series.apply(
                lambda x: self.normalizers.normalize_company(str(x), lang)
                if pd.notna(x)
                else x
            )
        elif action == "normalize_phone":
            region = params.get("region", "US")
            return series.apply(
                lambda x: self.normalizers.normalize_phone(str(x), region)
                if pd.notna(x)
                else x
            )
        elif action == "normalize_url":
            return series.apply(
                lambda x: self.normalizers.normalize_url(str(x)) if pd.notna(x) else x
            )
        elif action == "extract_domain":
            return series.apply(
                lambda x: self.normalizers.extract_domain(str(x)) if pd.notna(x) else x
            )

        # Regex operations
        elif action == "regex_replace":
            return series.str.replace(
                params["pattern"], params["replacement"], regex=True
            )

        # Slice operations
        elif action == "slice":
            return series.str.slice(params.get("start"), params.get("stop"))

        else:
            raise ValueError(f"Unknown action: {action}")

    def _handle_rule_error(self, rule: Rule, error: Exception):
        """Handle errors based on error mode"""
        error_msg = f"Error applying rule {rule.action}: {error}"

        if self.error_mode == ErrorMode.STRICT:
            raise RuntimeError(error_msg) from error
        elif self.error_mode == ErrorMode.LENIENT:
            log.error(error_msg)
        # SILENT mode: do nothing
