from io import StringIO
import pandas as pd
import numpy as np
from typing import List, Dict, Union, Optional, Any
from pydantic import BaseModel, ValidationError as PydanticValidationError, ConfigDict, Field
import logging

from aiecs.tools.base_tool import BaseTool
from aiecs.tools import register_tool

# Custom exceptions
class PandasToolError(Exception):
    """Base exception for PandasTool errors."""
    pass

class InputValidationError(PandasToolError):
    """Input validation error."""
    pass

class DataFrameError(PandasToolError):
    """DataFrame operation error."""
    pass

class SecurityError(PandasToolError):
    """Security-related error."""
    pass

class ValidationError(PandasToolError):
    """Validation error."""
    pass

@register_tool("pandas")
class PandasTool(BaseTool):
    """
    Tool encapsulating pandas functionality for data processing, supporting 30+ operations including:
      - Data reading/writing (CSV, JSON, Excel).
      - Descriptive statistics (summary, describe, value_counts).
      - Filtering and selection (filter, select_columns, drop_columns).
      - Grouping and aggregation (groupby, pivot_table).
      - Merging and concatenation (merge, concat).
      - Data transformation (sort_values, rename_columns, replace_values, fill_na, astype, apply).
      - Data reshaping (melt, pivot, stack, unstack).
      - Data cleaning (strip_strings, to_numeric, to_datetime).
      - Statistical computations (mean, sum, count, min, max).
      - Window functions (rolling).
      - Sampling and viewing (head, tail, sample).

    Inherits from BaseTool to leverage ToolExecutor for caching, concurrency, and error handling.
    """
    
    # Configuration schema
    class Config(BaseModel):
        """Configuration for the pandas tool"""
        model_config = ConfigDict(env_prefix="PANDAS_TOOL_")
        
        csv_delimiter: str = Field(
            default=",",
            description="Delimiter for CSV files"
        )
        encoding: str = Field(
            default="utf-8",
            description="Encoding for file operations"
        )
        default_agg: Dict[str, str] = Field(
            default={"numeric": "mean", "object": "count"},
            description="Default aggregation functions"
        )
        chunk_size: int = Field(
            default=10000,
            description="Chunk size for large file processing"
        )
        max_csv_size: int = Field(
            default=1000000,
            description="Threshold for chunked CSV processing"
        )
        allowed_file_extensions: List[str] = Field(
            default=['.csv', '.xlsx', '.json'],
            description="Allowed file extensions"
        )
    
    def __init__(self, config: Optional[Dict] = None):
        """
        Initialize PandasTool with configuration.

        Args:
            config (Dict, optional): Configuration overrides for PandasTool.

        Raises:
            ValueError: If config is invalid.
        """
        super().__init__(config)
        
        # Parse configuration
        self.config = self.Config(**(config or {}))
        
        self.logger = logging.getLogger(__name__)
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            handler.setFormatter(logging.Formatter('%(asctime)s %(levelname)s %(message)s'))
            self.logger.addHandler(handler)
        self.logger.setLevel(logging.INFO)

    def _validate_df(self, records: List[Dict]) -> pd.DataFrame:
        """
        Convert records to a DataFrame and validate.

        Args:
            records (List[Dict]): List of records to convert.

        Returns:
            pd.DataFrame: Validated DataFrame.

        Raises:
            InputValidationError: If records are empty or invalid.
        """
        if not records:
            raise InputValidationError("Records list is empty")
        try:
            df = pd.DataFrame(records)
            if df.empty:
                raise InputValidationError("DataFrame is empty")
            return df
        except Exception as e:
            raise InputValidationError(f"Failed to create DataFrame: {e}")

    def _validate_columns(self, df: pd.DataFrame, columns: List[str]) -> None:
        """
        Validate column names exist in DataFrame.

        Args:
            df (pd.DataFrame): DataFrame to validate.
            columns (List[str]): Columns to check.

        Raises:
            InputValidationError: If columns are not found.
        """
        if not columns:
            return
        available_columns = set(df.columns)
        missing = [col for col in columns if col not in available_columns]
        if missing:
            raise InputValidationError(
                f"Columns not found: {missing}. Available columns: {list(available_columns)}"
            )

    def _to_json_serializable(self, result: Union[pd.DataFrame, pd.Series, Dict]) -> Union[List[Dict], Dict]:
        """
        Convert result to JSON-serializable format.

        Args:
            result (Union[pd.DataFrame, pd.Series, Dict]): Result to convert.

        Returns:
            Union[List[Dict], Dict]: JSON-serializable result.
        """
        if isinstance(result, pd.DataFrame):
            for col in result.select_dtypes(include=['datetime64']).columns:
                result[col] = result[col].dt.strftime('%Y-%m-%d %H:%M:%S')
            return result.to_dict(orient="records")
        elif isinstance(result, pd.Series):
            if pd.api.types.is_datetime64_any_dtype(result):
                result = result.dt.strftime('%Y-%m-%d %H:%M:%S')
            return result.to_dict()
        elif isinstance(result, dict):
            def convert_value(v):
                if isinstance(v, (np.floating, np.integer)):
                    return float(v)
                elif isinstance(v, np.bool_):
                    return bool(v)
                elif isinstance(v, (pd.Timestamp, np.datetime64)):
                    return str(v)
                elif isinstance(v, np.ndarray):
                    return v.tolist()
                elif pd.isna(v):
                    return None
                return v
            return {k: convert_value(v) for k, v in result.items()}
        return result

    def read_csv(self, csv_str: str) -> List[Dict]:
        """Read CSV string into a DataFrame."""
        try:
            if len(csv_str) > self.config.max_csv_size:
                chunks = []
                for chunk in pd.read_csv(
                    StringIO(csv_str),
                    sep=self.config.csv_delimiter,
                    encoding=self.config.encoding,
                    chunksize=self.config.chunk_size
                ):
                    chunks.append(chunk)
                df = pd.concat(chunks)
            else:
                df = pd.read_csv(StringIO(csv_str), sep=self.config.csv_delimiter, encoding=self.config.encoding)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Failed to read CSV: {e}")

    def read_json(self, json_str: str) -> List[Dict]:
        """Read JSON string into a DataFrame."""
        try:
            df = pd.read_json(StringIO(json_str))
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Failed to read JSON: {e}")

    def read_file(self, file_path: str, file_type: str = "csv") -> List[Dict]:
        """Read data from a file (CSV, Excel, JSON)."""
        try:
            if file_type == "csv":
                file_size = sum(1 for _ in open(file_path, 'r', encoding=self.config.encoding))
                if file_size > self.config.chunk_size:
                    chunks = []
                    for chunk in pd.read_csv(
                        file_path,
                        sep=self.config.csv_delimiter,
                        encoding=self.config.encoding,
                        chunksize=self.config.chunk_size
                    ):
                        chunks.append(chunk)
                    df = pd.concat(chunks)
                else:
                    df = pd.read_csv(file_path, sep=self.config.csv_delimiter, encoding=self.config.encoding)
            elif file_type == "excel":
                df = pd.read_excel(file_path)
            elif file_type == "json":
                df = pd.read_json(file_path)
            else:
                raise ValidationError(f"Unsupported file type: {file_type}")
            return self._to_json_serializable(df)
        except ValidationError:
            raise
        except Exception as e:
            raise DataFrameError(f"Failed to read file: {e}")

    def write_file(self, records: List[Dict], file_path: str, file_type: str = "csv") -> Dict:
        """Write DataFrame to a file."""
        df = self._validate_df(records)
        try:
            if file_type == "csv":
                df.to_csv(file_path, index=False, sep=self.config.csv_delimiter, encoding=self.config.encoding)
            elif file_type == "excel":
                df.to_excel(file_path, index=False)
            elif file_type == "json":
                df.to_json(file_path, orient="records")
            else:
                raise ValidationError(f"Unsupported file type: {file_type}")
            return {"success": True, "file_path": file_path, "rows": len(df)}
        except Exception as e:
            raise DataFrameError(f"Failed to write file: {e}")

    def summary(self, records: List[Dict]) -> Dict:
        """Compute summary statistics for DataFrame."""
        df = self._validate_df(records)
        desc = df.describe(include="all").to_dict()
        return self._to_json_serializable(desc)

    def describe(self, records: List[Dict], columns: Optional[List[str]] = None) -> Dict:
        """Compute descriptive statistics for specified columns."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        desc = df.describe().to_dict()
        return self._to_json_serializable(desc)

    def value_counts(self, records: List[Dict], columns: List[str]) -> Dict:
        """Compute value counts for specified columns."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        result = {col: df[col].value_counts().to_dict() for col in columns}
        return self._to_json_serializable(result)

    def filter(self, records: List[Dict], condition: str) -> List[Dict]:
        """Filter DataFrame based on a condition."""
        df = self._validate_df(records)
        try:
            df = df.query(condition, engine="python")
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Invalid query condition: {e}")

    def select_columns(self, records: List[Dict], columns: List[str]) -> List[Dict]:
        """Select specified columns from DataFrame."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        return self._to_json_serializable(df[columns])

    def drop_columns(self, records: List[Dict], columns: List[str]) -> List[Dict]:
        """Drop specified columns from DataFrame."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        return self._to_json_serializable(df.drop(columns=columns))

    def drop_duplicates(self, records: List[Dict], columns: Optional[List[str]] = None) -> List[Dict]:
        """Drop duplicate rows based on specified columns."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
        return self._to_json_serializable(df.drop_duplicates(subset=columns))

    def dropna(self, records: List[Dict], axis: int = 0, how: str = "any") -> List[Dict]:
        """Drop rows or columns with missing values."""
        df = self._validate_df(records)
        if how not in ["any", "all"]:
            raise ValidationError("how must be 'any' or 'all'")
        return self._to_json_serializable(df.dropna(axis=axis, how=how))

    def groupby(self, records: List[Dict], by: List[str], agg: Dict[str, str]) -> List[Dict]:
        """Group DataFrame and apply aggregations."""
        df = self._validate_df(records)
        self._validate_columns(df, by + list(agg.keys()))
        try:
            df = df.groupby(by).agg(agg).reset_index()
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Groupby failed: {e}")

    def pivot_table(self, records: List[Dict], values: List[str], index: List[str], columns: List[str], aggfunc: str = "mean") -> List[Dict]:
        """Create a pivot table from DataFrame."""
        df = self._validate_df(records)
        self._validate_columns(df, values + index + columns)
        try:
            df = pd.pivot_table(df, values=values, index=index, columns=columns, aggfunc=aggfunc)
            return self._to_json_serializable(df.reset_index())
        except Exception as e:
            raise DataFrameError(f"Pivot table failed: {e}")

    def merge(self, records: List[Dict], records_right: List[Dict], on: Union[str, List[str]], join_type: str = "inner") -> List[Dict]:
        """Merge two DataFrames."""
        df_left = self._validate_df(records)
        df_right = self._validate_df(records_right)
        if join_type not in ["inner", "left", "right", "outer"]:
            raise ValidationError("join_type must be one of: inner, left, right, outer")
        self._validate_columns(df_left, [on] if isinstance(on, str) else on)
        self._validate_columns(df_right, [on] if isinstance(on, str) else on)
        try:
            df = df_left.merge(df_right, on=on, how=join_type)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Merge failed: {e}")

    def concat(self, records_list: List[List[Dict]], axis: int = 0) -> List[Dict]:
        """Concatenate multiple DataFrames."""
        if not records_list or not all(records_list):
            raise ValidationError("Records list is empty")
        dfs = [self._validate_df(records) for records in records_list]
        try:
            df = pd.concat(dfs, axis=axis, ignore_index=True)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Concat failed: {e}")

    def sort_values(self, records: List[Dict], sort_by: List[str], ascending: Union[bool, List[bool]] = True) -> List[Dict]:
        """Sort DataFrame by specified columns."""
        df = self._validate_df(records)
        self._validate_columns(df, sort_by)
        try:
            df = df.sort_values(by=sort_by, ascending=ascending)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Sort failed: {e}")

    def rename_columns(self, records: List[Dict], mapping: Dict[str, str]) -> List[Dict]:
        """Rename DataFrame columns."""
        df = self._validate_df(records)
        self._validate_columns(df, list(mapping.keys()))
        return self._to_json_serializable(df.rename(columns=mapping))

    def replace_values(self, records: List[Dict], to_replace: Dict, columns: Optional[List[str]] = None) -> List[Dict]:
        """Replace values in DataFrame."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        return self._to_json_serializable(df.replace(to_replace))

    def fill_na(self, records: List[Dict], value: Union[str, int, float], columns: Optional[List[str]] = None) -> List[Dict]:
        """Fill missing values in DataFrame."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df[columns] = df[columns].fillna(value)
        else:
            df = df.fillna(value)
        return self._to_json_serializable(df)

    def astype(self, records: List[Dict], dtypes: Dict[str, str]) -> List[Dict]:
        """Convert column types in DataFrame."""
        df = self._validate_df(records)
        self._validate_columns(df, list(dtypes.keys()))
        try:
            df = df.astype(dtypes)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Type conversion failed: {e}")

    def apply(self, records: List[Dict], func: str, columns: List[str], axis: int = 0) -> List[Dict]:
        """Apply a function to specified columns or rows."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        allowed_funcs = {
            'upper': lambda x: x.upper() if isinstance(x, str) else x,
            'lower': lambda x: x.lower() if isinstance(x, str) else x,
            'strip': lambda x: x.strip() if isinstance(x, str) else x,
            'capitalize': lambda x: x.capitalize() if isinstance(x, str) else x,
            'title': lambda x: x.title() if isinstance(x, str) else x,
            'len': lambda x: len(str(x)) if pd.notna(x) else 0,
            'abs': lambda x: abs(float(x)) if pd.notna(x) and not isinstance(x, str) else x,
            'round': lambda x: round(float(x)) if pd.notna(x) and not isinstance(x, str) else x,
            'ceil': lambda x: np.ceil(float(x)) if pd.notna(x) and not isinstance(x, str) else x,
            'floor': lambda x: np.floor(float(x)) if pd.notna(x) and not isinstance(x, str) else x,
            'int': lambda x: int(float(x)) if pd.notna(x) and not isinstance(x, str) else None,
            'float': lambda x: float(x) if pd.notna(x) and not isinstance(x, str) else None,
            'str': lambda x: str(x) if pd.notna(x) else "",
            'bool': lambda x: bool(x) if pd.notna(x) else False,
            'date_only': lambda x: x.date() if isinstance(x, pd.Timestamp) else x,
            'year': lambda x: x.year if isinstance(x, pd.Timestamp) else None,
            'month': lambda x: x.month if isinstance(x, pd.Timestamp) else None,
            'day': lambda x: x.day if isinstance(x, pd.Timestamp) else None,
        }
        try:
            if axis == 0:
                for col in columns:
                    df[col] = df[col].apply(allowed_funcs[func])
            else:
                df[columns] = df[columns].apply(allowed_funcs[func], axis=1)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Apply failed: {e}")

    def melt(self, records: List[Dict], id_vars: List[str], value_vars: List[str]) -> List[Dict]:
        """Melt DataFrame to long format."""
        df = self._validate_df(records)
        self._validate_columns(df, id_vars + value_vars)
        try:
            df = pd.melt(df, id_vars=id_vars, value_vars=value_vars)
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Melt failed: {e}")

    def pivot(self, records: List[Dict], index: str, columns: str, values: str) -> List[Dict]:
        """Pivot DataFrame to wide format."""
        df = self._validate_df(records)
        self._validate_columns(df, [index, columns, values])
        try:
            df = df.pivot(index=index, columns=columns, values=values)
            return self._to_json_serializable(df.reset_index())
        except Exception as e:
            raise DataFrameError(f"Pivot failed: {e}")

    def stack(self, records: List[Dict]) -> List[Dict]:
        """Stack DataFrame columns into rows."""
        df = self._validate_df(records)
        try:
            df = df.stack().reset_index()
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Stack failed: {e}")

    def unstack(self, records: List[Dict], level: Union[int, str] = -1) -> List[Dict]:
        """Unstack DataFrame rows into columns."""
        df = self._validate_df(records)
        try:
            df = df.unstack(level=level).reset_index()
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Unstack failed: {e}")

    def strip_strings(self, records: List[Dict], columns: List[str]) -> List[Dict]:
        """Strip whitespace from string columns."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        for col in columns:
            if df[col].dtype == "object":
                df[col] = df[col].str.strip()
        return self._to_json_serializable(df)

    def to_numeric(self, records: List[Dict], columns: List[str]) -> List[Dict]:
        """Convert columns to numeric type."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        try:
            for col in columns:
                df[col] = pd.to_numeric(df[col], errors="coerce")
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"To numeric failed: {e}")

    def to_datetime(self, records: List[Dict], columns: List[str], format: Optional[str] = None) -> List[Dict]:
        """Convert columns to datetime type."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        try:
            for col in columns:
                df[col] = pd.to_datetime(df[col], format=format, errors="coerce")
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"To datetime failed: {e}")

    def mean(self, records: List[Dict], columns: Optional[List[str]] = None) -> Dict:
        """Compute mean of numeric columns."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        return self._to_json_serializable(df.select_dtypes(include=np.number).mean())

    def sum(self, records: List[Dict], columns: Optional[List[str]] = None) -> Dict:
        """Compute sum of numeric columns."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        return self._to_json_serializable(df.select_dtypes(include=np.number).sum())

    def count(self, records: List[Dict], columns: Optional[List[str]] = None) -> Dict:
        """Compute count of non-null values."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        return self._to_json_serializable(df.count())

    def min(self, records: List[Dict], columns: Optional[List[str]] = None) -> Dict:
        """Compute minimum values."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        return self._to_json_serializable(df.min())

    def max(self, records: List[Dict], columns: Optional[List[str]] = None) -> Dict:
        """Compute maximum values."""
        df = self._validate_df(records)
        if columns:
            self._validate_columns(df, columns)
            df = df[columns]
        return self._to_json_serializable(df.max())

    def rolling(self, records: List[Dict], columns: List[str], window: int, function: str = "mean") -> List[Dict]:
        """Apply rolling window function to columns."""
        df = self._validate_df(records)
        self._validate_columns(df, columns)
        allowed_funcs = ["mean", "sum", "min", "max", "std", "count", "median"]
        if function not in allowed_funcs:
            raise ValidationError(f"Function '{function}' not allowed. Available: {allowed_funcs}")
        try:
            for col in columns:
                if pd.api.types.is_numeric_dtype(df[col]):
                    df[f"{col}_{function}_{window}"] = getattr(df[col].rolling(window), function)()
            return self._to_json_serializable(df)
        except Exception as e:
            raise DataFrameError(f"Rolling operation failed: {e}")

    def head(self, records: List[Dict], n: int = 5) -> List[Dict]:
        """Return first n rows of DataFrame."""
        df = self._validate_df(records)
        return self._to_json_serializable(df.head(n))

    def tail(self, records: List[Dict], n: int = 5) -> List[Dict]:
        """Return last n rows of DataFrame."""
        df = self._validate_df(records)
        return self._to_json_serializable(df.tail(n))

    def sample(self, records: List[Dict], n: int = 5, random_state: Optional[int] = None) -> List[Dict]:
        """Return random sample of n rows from DataFrame."""
        df = self._validate_df(records)
        return self._to_json_serializable(df.sample(n=min(n, len(df)), random_state=random_state))
