"""
ReporterSaveSchema - Output schema yaml files used by specified source modules
"""

import logging
from typing import Any

import yaml

from petsard.exceptions import ConfigError
from petsard.reporter.reporter_base import BaseReporter, RegexPatterns


class ReporterSaveSchema(BaseReporter):
    """
    Schema output reporter

    This reporter outputs schema information to YAML files based on the modules specified by the source parameter.
    Can be used to track and record data structure changes at each processing stage.
    """

    def __init__(self, config: dict):
        """
        Args:
            config (dict): Configuration dictionary.
                - method (str): Report method (must be 'SAVE_SCHEMA')
                - source (str | List[str]): Data source module name(s)
                    Supported modules: 'Loader', 'Splitter', 'Preprocessor', 'Synthesizer',
                               'Postprocessor', 'Constrainer'
                - output (str, optional): Output file name prefix, default 'petsard'
                - yaml_output (bool, optional): Whether to output YAML files, default False
                - properties (str | List[str], optional): Property names to output, default all
                    Supported properties: 'type', 'category', 'dtype', 'nullable', 'unique_count',
                               'precision', 'min', 'max', 'mean', 'std', 'categories'
                    Note:
                    - 'type': schema-defined type (e.g., 'str', 'int', 'float64')
                    - 'category': category marker in schema (True/False)
                    - 'dtype': actual pandas dtype (e.g., 'float64', 'object')
                    - 'precision': decimal places for numeric fields (from type_attr)

        Raises:
            ConfigError: If 'source' field is missing in config, or source/properties format is incorrect
        """
        super().__init__(config)

        # source should be string or list of strings: Union[str, List[str]]
        if "source" not in self.config:
            raise ConfigError("Configuration must include 'source' field")
        elif not isinstance(self.config["source"], str | list) or (
            isinstance(self.config["source"], list)
            and not all(isinstance(item, str) for item in self.config["source"])
        ):
            raise ConfigError("'source' must be a string or list of strings")

        # Convert source to list (if it's a string)
        if isinstance(self.config["source"], str):
            self.config["source"] = [self.config["source"]]

        # Process properties parameter
        if "properties" in self.config:
            if isinstance(self.config["properties"], str):
                self.config["properties"] = [self.config["properties"]]
            elif not isinstance(self.config["properties"], list) or not all(
                isinstance(item, str) for item in self.config["properties"]
            ):
                raise ConfigError("'properties' must be a string or list of strings")

        self._logger = logging.getLogger(f"PETsARD.{self.__class__.__name__}")

    def create(self, data: dict) -> dict[str, Any]:
        """
        Process data and extract schema information

        Args:
            data (dict): Data dictionary generated by ReporterOperator.set_input()
                Format reference: BaseReporter._verify_create_input()
                May contain 'metadata' key with Schema from each module

        Returns:
            dict[str, Any]: Processed schema data dictionary
                key: full experiment name
                value: Schema object
        """
        # Extract and store metadata (if available)
        if "metadata" in data:
            self._metadata_dict = data.pop("metadata")
            self._logger.debug(
                f"Received metadata for {len(self._metadata_dict)} modules"
            )
        else:
            self._metadata_dict = {}

        # Extract and store schema_history (if available)
        if "schema_history" in data:
            self._schema_history_dict = data.pop("schema_history")
            self._logger.debug(
                f"Received schema history for {len(self._schema_history_dict)} modules"
            )
            # Record number of snapshots for each module
            for module, history in self._schema_history_dict.items():
                self._logger.debug(f"  - {module}: {len(history)} snapshots")
        else:
            self._schema_history_dict = {}

        # Validate input data
        self._verify_create_input(data)

        processed_schemas = {}

        # Iterate through all data items
        for full_expt_tuple, df in data.items():
            if df is None:
                continue

            # Check if last module is in source list
            # full_expt_tuple format: ('Loader', 'default', 'Preprocessor', 'scaler')
            if len(full_expt_tuple) >= 2:
                last_module = full_expt_tuple[-2]
                last_expt_name = full_expt_tuple[-1]

                # Remove possible suffix "_[xxx]" to match source
                clean_expt_name = RegexPatterns.POSTFIX_REMOVAL.sub("", last_expt_name)

                # Check if module name or experiment name is in source
                if (
                    last_module in self.config["source"]
                    or clean_expt_name in self.config["source"]
                ):
                    # Generate full experiment name
                    full_expt_name = "_".join(
                        [
                            f"{full_expt_tuple[i]}[{full_expt_tuple[i + 1]}]"
                            for i in range(0, len(full_expt_tuple), 2)
                        ]
                    )
                    processed_schemas[full_expt_name] = df

        # Special handling: expand Preprocessor schema history
        # metadata_dict may contain keys in format "Preprocessor_step_name"
        for source in self.config["source"]:
            if source in ["Preprocessor", "Postprocessor"]:
                # Find all metadata entries starting with this module
                for meta_key in list(self._metadata_dict.keys()):
                    if meta_key.startswith(f"{source}_"):
                        # Extract step name
                        step_name = meta_key[
                            len(source) + 1 :
                        ]  # Remove "Preprocessor_" prefix

                        # Find corresponding data entry (if available)
                        # But for schema history, we mainly focus on metadata
                        # So process even without corresponding data

                        # Build a synthetic full_expt_name with step information
                        # Find original Preprocessor entry from processed_schemas
                        base_expt_name = None
                        for expt_name in processed_schemas.keys():
                            if f"{source}[" in expt_name:
                                base_expt_name = expt_name
                                break

                        if base_expt_name:
                            # Add step information after module name
                            # e.g.: "Loader[x]_Preprocessor[y]" -> "Loader[x]_Preprocessor[y_step_name]"
                            step_expt_name = (
                                base_expt_name.replace(
                                    f"{source}[", f"{source}[{step_name}_"
                                )
                                .replace(f"[{step_name}_", "[")
                                .replace("]", f"_{step_name}]")
                            )

                            # Use same dataframe as base (or empty)
                            processed_schemas[step_expt_name] = processed_schemas.get(
                                base_expt_name
                            )

                            self._logger.debug(
                                f"Added schema history entry: {step_expt_name}"
                            )

        self._logger.info(
            f"Processed schema information for {len(processed_schemas)} modules"
        )
        return processed_schemas

    def report(self, processed_data: dict[str, Any] | None = None) -> dict[str, Any]:
        """
        Generate and save schema report

        Default output is CSV format (flattened table), one row per source
        Optional YAML format output (yaml_output=True)

        Args:
            processed_data (dict[str, Any] | None): Processed data
                key: experiment name
                value: pandas DataFrame (for schema inference)

        Returns:
            dict[str, Any]: Saved schema data
        """
        if not processed_data:
            self._logger.warning(
                "No data to process, skipping schema report generation"
            )
            return {}

        saved_schemas = {}
        flattened_rows = []

        for expt_name, df in processed_data.items():
            if df is None:
                self._logger.debug(f"Skipping empty data: {expt_name}")
                continue

            try:
                # Try to get corresponding metadata
                metadata = self._get_metadata_for_expt(expt_name)

                # Infer schema from DataFrame and pass metadata
                schema_dict = self._infer_schema_from_dataframe(df, metadata)
                saved_schemas[expt_name] = schema_dict

                # Flatten entire source schema to single row
                row = self._flatten_source_schema(expt_name, schema_dict)
                flattened_rows.append(row)

                # Optional: save as YAML file
                if self.config.get("yaml_output", False):
                    output_filename = f"{self.config['output']}_schema_{expt_name}.yaml"
                    self._save_schema_to_yaml(schema_dict, output_filename)
                    self._logger.info(f"Saved YAML schema to {output_filename}")

            except Exception as e:
                self._logger.error(f"Error processing {expt_name}: {e}")
                continue

        # Default output: summary CSV (with source name in filename)
        if flattened_rows:
            import pandas as pd

            df_output = pd.DataFrame(flattened_rows)

            # Sort columns: source first, then other columns by field name
            # Format: [field_name]_property, so all properties of same field are grouped together
            columns = list(df_output.columns)

            # Separate source and other columns
            if "source" in columns:
                columns.remove("source")
                # Sort other columns
                sorted_columns = ["source"] + sorted(columns)
            else:
                sorted_columns = sorted(columns)

            # Reorder DataFrame
            df_output = df_output[sorted_columns]

            # Generate filename containing all source module names (similar to save_data approach)
            source_names = "-".join(self.config["source"])
            csv_filename = f"{self.config['output']}_schema_{source_names}_summary.csv"
            df_output.to_csv(csv_filename, index=False, encoding="utf-8")
            self._logger.info(f"Saved schema summary to {csv_filename}")

        return saved_schemas

    def _get_metadata_for_expt(self, expt_name: str):
        """
        Get metadata corresponding to experiment name

        Supported formats:
        - "Loader[default]_Preprocessor[v1]" (standard format)
        - "Loader[default]_Preprocessor[v1_after_encoder]" (schema history format)

        Args:
            expt_name: experiment name

        Returns:
            Schema or None
        """
        self._logger.debug(f"Looking up metadata for: {expt_name}")

        if not hasattr(self, "_metadata_dict"):
            self._logger.warning("No _metadata_dict attribute found")
            return None

        # Record all available metadata keys
        self._logger.debug(
            f"Available metadata keys: {list(self._metadata_dict.keys())}"
        )

        # Find last module part
        # e.g.: "Loader[default]_Preprocessor[v1_after_encoder]"
        # We need to extract "Preprocessor" and possibly "after_encoder"
        parts = expt_name.split("_")
        self._logger.debug(f"Experiment name parts: {parts}")

        # Find last part containing brackets
        last_module_part = None
        for part in reversed(parts):
            if "[" in part and "]" in part:
                last_module_part = part
                last_module_index = parts.index(part)
                break

        if not last_module_part:
            self._logger.warning(f"No module part with brackets found in: {expt_name}")
            return None

        # Extract module name (part before brackets)
        last_module = last_module_part.split("[")[0]
        self._logger.debug(f"Last module: {last_module}")

        # Check for step suffix
        # Bracket content format may be: "v1_after_encoder"
        bracket_content = last_module_part.split("[")[1].rstrip("]")
        self._logger.debug(f"Bracket content: {bracket_content}")

        # Check if there are more underscore-separated parts after brackets (these may be step names)
        # e.g.: "Loader[default]_Preprocessor[v1]_after_encoder"
        remaining_parts = parts[last_module_index + 1 :]
        self._logger.debug(f"Remaining parts after module: {remaining_parts}")

        if remaining_parts:
            # Has step suffix, e.g.: ["after", "encoder"]
            step_suffix = "_".join(remaining_parts)
            meta_key = f"{last_module}_{step_suffix}"
            self._logger.debug(f"Trying step metadata key: {meta_key}")
            metadata = self._metadata_dict.get(meta_key)

            if metadata:
                self._logger.debug(f"Found schema history metadata: {meta_key}")
                return metadata
            else:
                self._logger.debug(f"Step metadata key not found: {meta_key}")

        # No step suffix, use standard lookup
        self._logger.debug(f"Trying standard metadata key: {last_module}")
        metadata = self._metadata_dict.get(last_module)
        if metadata:
            self._logger.debug(f"Found standard metadata: {last_module}")
            return metadata
        else:
            self._logger.warning(f"Standard metadata key not found: {last_module}")

        return None

    def _infer_schema_from_dataframe(self, df, metadata=None) -> dict[str, Any]:
        """
        Infer schema structure from DataFrame

        Args:
            df: pandas DataFrame
            metadata: Schema metadata (optional)

        Returns:
            dict: schema dictionary containing field information
        """
        schema = {"columns": {}, "shape": {"rows": len(df), "columns": len(df.columns)}}

        # Get list of properties to output
        properties = self.config.get("properties", None)

        # Debug: record metadata information
        if metadata:
            self._logger.debug(
                f"Processing with metadata. Attributes count: {len(metadata.attributes) if hasattr(metadata, 'attributes') else 0}"
            )
        else:
            self._logger.warning("No metadata provided for schema inference")

        # Record information for each field
        for col in df.columns:
            col_info = {}

            # Get properties from metadata
            schema_type = None
            schema_category = None
            schema_nullable = None
            schema_precision = None

            if (
                metadata
                and hasattr(metadata, "attributes")
                and col in metadata.attributes
            ):
                attr = metadata.attributes[col]
                schema_type = attr.type

                # Read properties from type_attr
                if attr.type_attr:
                    schema_category = attr.type_attr.get("category", False)
                    schema_nullable = attr.type_attr.get("nullable", True)
                    schema_precision = attr.type_attr.get("precision")
                else:
                    schema_category = False
                    schema_nullable = True

                self._logger.debug(
                    f"Column {col}: type={schema_type}, category={schema_category}, nullable={schema_nullable}, precision={schema_precision}"
                )
            else:
                # When no metadata, infer from data
                dtype_str = str(df[col].dtype)

                # Infer category
                if dtype_str == "category":
                    schema_category = True
                elif df[col].dtype == "object" and len(df[col]) > 0:
                    unique_ratio = len(df[col].unique()) / len(df[col])
                    schema_category = unique_ratio < 0.05
                else:
                    schema_category = False

                # Infer nullable
                schema_nullable = (
                    bool(df[col].isna().any()) if len(df[col]) > 0 else True
                )

                self._logger.debug(
                    f"Column {col}: No metadata found, inferred category={schema_category}, nullable={schema_nullable}"
                )

            # Basic properties - both type and dtype should be output
            if properties is None or "type" in properties:
                # Prefer schema type, otherwise use inferred type
                if schema_type:
                    col_info["type"] = schema_type
                else:
                    # Infer type from dtype (simplified mapping)
                    dtype_str = str(df[col].dtype)
                    type_mapping = {
                        "int8": "int8",
                        "int16": "int16",
                        "int32": "int32",
                        "int64": "int64",
                        "float32": "float32",
                        "float64": "float64",
                        "bool": "boolean",
                        "object": "string",
                        "category": "string",
                    }
                    col_info["type"] = type_mapping.get(dtype_str, dtype_str)

            if properties is None or "category" in properties:
                # Output category (True/False), should never be None
                col_info["category"] = schema_category
                self._logger.debug(
                    f"Column {col}: Set category in col_info = {schema_category}"
                )

            if properties is None or "dtype" in properties:
                # dtype always output (actual pandas dtype)
                col_info["dtype"] = str(df[col].dtype)

            # Precision (using already obtained schema_precision)
            if properties is None or "precision" in properties:
                if schema_precision is not None:
                    col_info["precision"] = schema_precision

            if properties is None or "nullable" in properties:
                # Use already obtained schema_nullable
                col_info["nullable"] = (
                    schema_nullable if schema_nullable is not None else True
                )

            if properties is None or "unique_count" in properties:
                col_info["unique_count"] = int(df[col].nunique())

            # If numeric type, add statistical information
            if df[col].dtype.kind in "biufc":  # bool, int, unsigned int, float, complex
                # Check if any statistical properties are needed
                stats_needed = (
                    properties is None
                    or "min" in properties
                    or "max" in properties
                    or "mean" in properties
                    or "std" in properties
                )

                if stats_needed and not df[col].isna().all():
                    statistics = {}

                    # Decide precision of statistics based on data type
                    if df[col].dtype.kind in "biu":  # bool, int, unsigned int
                        # Integer type: round to integer
                        if properties is None or "min" in properties:
                            statistics["min"] = int(round(df[col].min()))
                        if properties is None or "max" in properties:
                            statistics["max"] = int(round(df[col].max()))
                        if properties is None or "mean" in properties:
                            statistics["mean"] = int(round(df[col].mean()))
                        if properties is None or "std" in properties:
                            statistics["std"] = int(round(df[col].std()))
                    else:  # float, complex
                        # Float type: detect data precision and limit decimal places
                        decimal_places = self._detect_decimal_places(df[col])
                        if properties is None or "min" in properties:
                            statistics["min"] = round(
                                float(df[col].min()), decimal_places
                            )
                        if properties is None or "max" in properties:
                            statistics["max"] = round(
                                float(df[col].max()), decimal_places
                            )
                        if properties is None or "mean" in properties:
                            statistics["mean"] = round(
                                float(df[col].mean()), decimal_places
                            )
                        if properties is None or "std" in properties:
                            statistics["std"] = round(
                                float(df[col].std()), decimal_places
                            )

                    if statistics:
                        col_info["statistics"] = statistics
                elif stats_needed:
                    # All NA case
                    statistics = {}
                    if properties is None or "min" in properties:
                        statistics["min"] = None
                    if properties is None or "max" in properties:
                        statistics["max"] = None
                    if properties is None or "mean" in properties:
                        statistics["mean"] = None
                    if properties is None or "std" in properties:
                        statistics["std"] = None
                    if statistics:
                        col_info["statistics"] = statistics

            # If object type (usually string), record sample values
            elif (df[col].dtype == "object" or df[col].dtype.name == "category") and (
                properties is None or "categories" in properties
            ):
                unique_values = df[col].dropna().unique()
                if len(unique_values) <= 10:  # Only record when few unique values
                    col_info["categories"] = [str(v) for v in unique_values]

            schema["columns"][str(col)] = col_info

        return schema

    def _detect_decimal_places(self, series, max_check: int = 100) -> int:
        """
        Detect decimal places of float field

        Args:
            series: pandas Series (float type)
            max_check: max number of data points to check (default 100)

        Returns:
            int: decimal places (max 6)
        """
        non_null = series.dropna()
        if len(non_null) == 0:
            return 2  # default 2 decimal places

        decimal_places = 0
        # Only check first max_check values for efficiency
        for val in non_null.head(max_check):
            # Convert to string and check decimal places
            val_str = f"{val:.10f}".rstrip("0").rstrip(".")
            if "." in val_str:
                decimal_places = max(decimal_places, len(val_str.split(".")[1]))

        # Limit to max 6 decimal places
        return min(decimal_places, 6) if decimal_places > 0 else 2

    def _save_schema_to_yaml(self, schema_dict: dict, filename: str) -> None:
        """
        Save schema dictionary as YAML file

        Args:
            schema_dict: schema dictionary
            filename: output filename
        """
        try:
            with open(filename, "w", encoding="utf-8") as f:
                yaml.dump(
                    schema_dict,
                    f,
                    default_flow_style=False,
                    allow_unicode=True,
                    sort_keys=False,
                )
            self._logger.debug(f"Schema written to file: {filename}")
        except Exception as e:
            self._logger.error(f"Error writing YAML file: {e}")
            raise

    def _flatten_source_schema(self, source: str, schema_dict: dict) -> dict:
        """
        Flatten entire source schema to single row

        One source becomes one row, all properties of all fields flattened into columns
        Use [field_name]_property format, so properties of same field are grouped when sorted alphabetically
        e.g.: source, [age]_dtype, [age]_nullable, [age]_min, [age]_max, [income]_dtype, ...

        Args:
            source: source experiment name
            schema_dict: complete schema dictionary

        Returns:
            dict: flattened row data, source is first field
        """
        row = {"source": source}

        # Iterate through all fields
        for column_name, column_info in schema_dict["columns"].items():
            # First level properties
            for key, value in column_info.items():
                if key == "statistics" and isinstance(value, dict):
                    # Second level: statistics (like min, max, mean)
                    for stat_key, stat_value in value.items():
                        row[f"[{column_name}]_{stat_key}"] = stat_value
                elif key == "categories" and isinstance(value, list):
                    # Category value list, convert to string
                    row[f"[{column_name}]_categories"] = "|".join(str(v) for v in value)
                elif not isinstance(value, (dict, list)):
                    # Other simple types (dtype, nullable, unique_count, category, type, etc.)
                    row[f"[{column_name}]_{key}"] = value

        return row
