import re
from typing import Iterator, Dict, Any, Optional, Set
from atomingest.normalization.steps import BaseNormalizerStep, register_step


@register_step("RemoveBOMStep")
class RemoveBOMStep(BaseNormalizerStep):
    """
    Normalization step to remove UTF-8 BOM from the beginning of files.
    Configuration:
        check_all_lines (bool): if True, checks for BOM at the start of every line.
                                (useful for concatenated files). Default is False.
                                (only checks the first line)
    """

    def process_lines(self, lines: Iterator[str]) -> Iterator[str]:
        check_all_lines = self.config.get("check_all_lines", False)

        for idx, line in enumerate(lines):
            if idx == 0 or check_all_lines:
                if line.startswith("\ufeff"):
                    self.report.add_warning("UTF-8 BOM detected and removed.")
                    line = line.lstrip("\ufeff")
            yield line


@register_step("SkipUntilHeaderStep")
class SkipUntilHeaderStep(BaseNormalizerStep):
    """
    Normalization step to drop metadata/garbage lines until a valid header is found
    The header is identified by matching a regex (`header_regex`) OR by checking
    if the line contains a set of expected column names (`expected_columns`)

    Configuration:
        header_regex (str): Regular expression to identify the header line.
        expected_columns (List[str]): List of expected column names to identify the header line.
        separator (str): The delimiter used to split the line into columns.
                        Default is comma (',').
        max_skip (int): Maximum number of lines to skip before giving up on finding a header.
                        Default is None (unlimited).
    """

    def process_lines(self, lines: Iterator[str]) -> Iterator[str]:
        header_regex = self.config.get("header_regex")
        expected_columns = self.config.get("expected_columns")
        separator = self.config.get("separator", ",")
        max_skip = self.config.get("max_skip")

        if not header_regex and not expected_columns:
            # No exit criteria defined; yield everything (or warn)
            yield from lines
            return

        # Pre-convert expected_columns to set for efficient lookup
        expected_columns_set = set(expected_columns) if expected_columns else None

        header_found = False
        skipped_count = 0

        for line in lines:
            if not header_found:
                is_match = False

                # Check header_regex first (takes priority)
                if header_regex and re.search(header_regex, line):
                    is_match = True
                elif expected_columns_set and not header_regex:
                    # Only check expected_columns if no header_regex provided
                    parts = [p.strip() for p in line.split(separator)]
                    if expected_columns_set.issubset(set(parts)):
                        is_match = True

                if is_match:
                    header_found = True
                    yield line
                else:
                    # Increment skip count and check against max_skip
                    skipped_count += 1
                    self.report.increment_metadata_skipped()

                    if max_skip is not None and skipped_count > max_skip:
                        msg = f"Header not found after skipping {skipped_count} lines (max_skip={max_skip})"
                        if self.strict:
                            raise ValueError(msg)
                        else:
                            self.report.add_warning(msg)
                            # In non-strict mode, stop skipping and yield remaining lines
                            yield line
                            yield from lines
                            return
            else:
                yield line

        if not header_found and skipped_count > 0:
            msg = "Header not found in file."
            if self.strict:
                raise ValueError(msg)
            else:
                self.report.add_warning(msg)


@register_step("DropRepeatedHeaderRowsStep")
class DropRepeatedHeaderRowsStep(BaseNormalizerStep):
    """
    Row-step to detect repeated header rows mid-file (e.g. concatenated extracts) and drop them
    safely
    """

    def process_row(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        # Check if the row looks like a header (values equal to keys)
        # This occurs when a CSV parser reads a line that repeats the header.
        if row and all(k == v for k, v in row.items()):
            self.report.increment_duplicate_headers()
            return None
        return row


@register_step("TrimWhitespaceStep")
class TrimWhitespaceStep(BaseNormalizerStep):
    """
    Normalization step to trim leading and trailing whitespace from all string values in a row.
    Configuration:
        trim_quotes (bool): If True, also removes surrounding quotes (' or "). Default False.
        collapse_spaces (bool): If True, replaces multiple spaces with a single space.
    """

    def process_row(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        trim_quotes = self.config.get("trim_quotes", False)
        collapse_spaces = self.config.get("collapse_spaces", False)

        cleaned_row = {}
        for key, value in row.items():
            if isinstance(value, str):
                if collapse_spaces:
                    value = re.sub(r"\s+", " ", value)

                value = value.strip()

                if trim_quotes:
                    value = value.strip("'\"")
                    value = value.strip()
            cleaned_row[key] = value
        return cleaned_row


@register_step("NullTokenNormalizationStep")
class NullTokenNormalizationStep(BaseNormalizerStep):
    """
    Normalization step to convert configured null tokens to None
    Configuration:
        null_tokens (List[str]): List of string tokens to treat as null.
                                 Default ["", "NA", "NULL", "nan", "-"]
    """

    def process_row(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        default_tokens = ["", "NA", "NULL", "nan", "-"]
        null_tokens = set(self.config.get("null_tokens", default_tokens))

        cleaned_row: Dict[str, Any] = {}
        for key, value in row.items():
            if value in null_tokens:
                cleaned_row[key] = None
            else:
                cleaned_row[key] = value
        return cleaned_row


@register_step("ColumnNameNormalizationStep")
class ColumnNameNormalizationStep(BaseNormalizerStep):
    """
    Normalization step to standardize column names.
    Applies aliasing, then normalization (strip, case, space replacement, special char removal).

    Configuration:
        alias_map (Dict[str, str]): Mapping of input column names to desired output names.
                                    If provided, unknown columns trigger a warning.
        case (str): 'lower', 'upper', or None. Default 'lower'.
        replace_spaces (str): String to replace spaces with. Default '_'.
        remove_non_alphanumeric (bool): Remove characters other than a-z, 0-9, and _. Default True.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._warned_unknown_columns: Set[str] = set()

    def process_row(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        alias_map = self.config.get("alias_map")
        case_transform = self.config.get("case", "lower")
        replace_spaces = self.config.get("replace_spaces", "_")
        remove_non_alphanumeric = self.config.get("remove_non_alphanumeric", True)

        new_row = {}

        for col, val in row.items():
            final_col = col

            # 1. Alias Mapping
            if alias_map:
                if col in alias_map:
                    final_col = alias_map[col]
                    # We assume aliased names are already in the desired format
                    new_row[final_col] = val
                    continue
                else:
                    # Unknown column warning (one time per column)
                    if col not in self._warned_unknown_columns:
                        self.report.add_warning(f"Unknown column encountered: '{col}'")
                        self._warned_unknown_columns.add(col)

            # 2. Normalization (if not aliased)
            # Strip
            final_col = final_col.strip()

            # Case
            if case_transform == "lower":
                final_col = final_col.lower()
            elif case_transform == "upper":
                final_col = final_col.upper()

            # Replace spaces
            if replace_spaces is not None:
                final_col = final_col.replace(" ", replace_spaces)

            # Remove non-alphanumeric (keeping underscores)
            if remove_non_alphanumeric:
                final_col = re.sub(r"[^a-zA-Z0-9_]", "", final_col)

            new_row[final_col] = val

        return new_row


@register_step("BasicTypeCleanupStep")
class BasicTypeCleanupStep(BaseNormalizerStep):
    """
    Normalization step to clean and normalize basic data types:
    - Numeric values: remove commas and optionally convert to float/int
    - Boolean values: normalize various boolean representations
    - Date values: parse and convert to ISO 8601 format

    Configuration:
        numeric_columns (List[str]): List of column names to treat as numeric.
                                     Commas will be removed. Default: []
        convert_to_number (bool): If True, convert numeric strings to actual numbers.
                                 Default: False (keep as cleaned strings)
        boolean_columns (List[str]): List of column names to treat as boolean. Default: []
        true_values (List[str]): Values to treat as True. Default: ["true", "yes", "1", "y", "t"]
        false_values (List[str]): Values to treat as False. Default: ["false", "no", "0", "n", "f"]
        date_columns (List[str]): List of column names to treat as dates. Default: []
        date_formats (List[str]): List of date format strings to try parsing.
                                 Default: ["%Y-%m-%d", "%m/%d/%Y", "%d/%m/%Y", "%Y/%m/%d"]
        output_date_format (str): Format for output dates. Default: "%Y-%m-%d" (ISO 8601)
    """

    def process_row(self, row: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        from datetime import datetime

        # Get configuration
        numeric_columns = self.config.get("numeric_columns", [])
        convert_to_number = self.config.get("convert_to_number", False)
        boolean_columns = self.config.get("boolean_columns", [])
        true_values = set(
            v.lower()
            for v in self.config.get(
                "true_values", ["true", "yes", "1", "y", "t", "on"]
            )
        )
        false_values = set(
            v.lower()
            for v in self.config.get(
                "false_values", ["false", "no", "0", "n", "f", "off"]
            )
        )
        date_columns = self.config.get("date_columns", [])
        date_formats = self.config.get(
            "date_formats",
            ["%Y-%m-%d", "%m/%d/%Y", "%d/%m/%Y", "%Y/%m/%d", "%Y-%m-%d %H:%M:%S"],
        )
        output_date_format = self.config.get("output_date_format", "%Y-%m-%d")

        cleaned_row: Dict[str, Any] = {}

        for key, value in row.items():
            # Skip None values
            if value is None:
                cleaned_row[key] = value
                continue

            # Numeric cleanup
            if key in numeric_columns:
                try:
                    # Remove commas from numeric strings
                    if isinstance(value, str):
                        cleaned_value = value.replace(",", "")
                        if convert_to_number:
                            # Try to convert to int first, then float
                            if "." in cleaned_value:
                                cleaned_row[key] = float(cleaned_value)
                            else:
                                cleaned_row[key] = int(cleaned_value)
                        else:
                            cleaned_row[key] = cleaned_value
                    else:
                        cleaned_row[key] = value
                except (ValueError, AttributeError) as e:
                    msg = f"Failed to parse numeric value '{value}' in column '{key}': {e}"
                    self.report.add_warning(msg)
                    if self.strict:
                        raise ValueError(msg)
                    cleaned_row[key] = value
                continue

            # Boolean normalization
            if key in boolean_columns:
                if isinstance(value, str):
                    value_lower = value.lower().strip()
                    if value_lower in true_values:
                        cleaned_row[key] = True
                    elif value_lower in false_values:
                        cleaned_row[key] = False
                    else:
                        msg = (
                            f"Failed to parse boolean value '{value}' in column '{key}'"
                        )
                        self.report.add_warning(msg)
                        if self.strict:
                            raise ValueError(msg)
                        cleaned_row[key] = value
                elif isinstance(value, bool):
                    cleaned_row[key] = value
                else:
                    msg = (
                        f"Non-string/boolean value '{value}' in boolean column '{key}'"
                    )
                    self.report.add_warning(msg)
                    if self.strict:
                        raise ValueError(msg)
                    cleaned_row[key] = value
                continue

            # Date parsing
            if key in date_columns:
                if isinstance(value, str):
                    parsed_date = None
                    for fmt in date_formats:
                        try:
                            parsed_date = datetime.strptime(value.strip(), fmt)
                            break
                        except ValueError:
                            continue

                    if parsed_date:
                        if output_date_format is None:
                            # Default to ISO 8601
                            cleaned_row[key] = parsed_date.date()
                        else:
                            cleaned_row[key] = parsed_date.strftime(output_date_format)
                    else:
                        msg = (
                            f"Failed to parse date value '{value}' in column '{key}' "
                            f"with formats: {date_formats}"
                        )
                        self.report.add_warning(msg)
                        if self.strict:
                            raise ValueError(msg)
                        cleaned_row[key] = value
                else:
                    cleaned_row[key] = value
                continue

            # Keep other values unchanged
            cleaned_row[key] = value

        return cleaned_row


@register_step("DropTrailerLinesStep")
class DropTrailerLinesStep(BaseNormalizerStep):
    """
    Normalization step to remove trailer/footer lines based on regex/prefix.

    Configuration:
        trailer_regex (str): Regex pattern to identify trailer lines.
        row_count_regex (str): Optional regex to extract the row count from the trailer line.
                               Should contain a capturing group for the number.
        trailer_position (str): Where trailers are expected: 'anywhere' (default) or 'end-only'.
                                'anywhere': Checks every line in the file for trailers.
                                'end-only': Only checks the last max_trailer_lines for better performance.
        max_trailer_lines (int): When trailer_position='end-only', maximum number of lines to check
                                 from the end of the file. Default is 10.
    """

    def process_lines(self, lines: Iterator[str]) -> Iterator[str]:
        # Support "pattern" as an alias for "trailer_regex"
        trailer_regex = self.config.get("trailer_regex") or self.config.get("pattern")
        row_count_regex = self.config.get("row_count_regex")
        trailer_position = self.config.get("trailer_position", "anywhere")
        max_trailer_lines = self.config.get("max_trailer_lines", 10)

        if not trailer_regex:
            yield from lines
            return

        # Pre-compile regex patterns for better performance
        trailer_pattern = re.compile(trailer_regex)
        row_count_pattern = re.compile(row_count_regex) if row_count_regex else None

        # Track row count to detect multiple different values
        extracted_row_count = None

        if trailer_position == "end-only":
            # Buffer all lines to check only the end
            all_lines = list(lines)

            # Check last max_trailer_lines for trailers
            trailer_start_idx = len(all_lines)
            start_check_idx = max(0, len(all_lines) - max_trailer_lines)

            for idx in range(len(all_lines) - 1, start_check_idx - 1, -1):
                line = all_lines[idx]
                if trailer_pattern.search(line):
                    trailer_start_idx = idx
                    self.report.increment_trailer_skipped()

                    # Extract row count if pattern provided
                    if row_count_pattern:
                        match = row_count_pattern.search(line)
                        if match and match.groups():
                            try:
                                row_count = int(match.group(1))

                                if (
                                    extracted_row_count is not None
                                    and extracted_row_count != row_count
                                ):
                                    self.report.add_warning(
                                        f"Multiple different row counts found in trailers: "
                                        f"previous={extracted_row_count}, current={row_count}"
                                    )

                                extracted_row_count = row_count
                                self.report.set_trailer_row_count(row_count)
                            except ValueError:
                                self.report.add_warning(
                                    f"Failed to parse row count from trailer line: {line.strip()}"
                                )
                else:
                    # Stop checking once we hit a non-trailer line
                    break

            # Yield all lines before the trailers
            for line in all_lines[:trailer_start_idx]:
                yield line
        else:
            # Default 'anywhere' behavior: check each line as we go
            for line in lines:
                if trailer_pattern.search(line):
                    self.report.increment_trailer_skipped()
                    if row_count_pattern:
                        match = row_count_pattern.search(line)
                        if match and match.groups():
                            try:
                                row_count = int(match.group(1))

                                # Warn if multiple different row counts found
                                if (
                                    extracted_row_count is not None
                                    and extracted_row_count != row_count
                                ):
                                    self.report.add_warning(
                                        f"Multiple different row counts found in trailers: "
                                        f"previous={extracted_row_count}, current={row_count}"
                                    )

                                extracted_row_count = row_count
                                self.report.set_trailer_row_count(row_count)
                            except ValueError:
                                self.report.add_warning(
                                    f"Failed to parse row count from trailer line: {line.strip()}"
                                )
                else:
                    yield line
