structured.py

from typing import Literal, Optional, List, Union, Any, Dict
from functools import lru_cache
import json
import operator
import logging

from pydantic import BaseModel, Field, field_validator
import pandas as pd

__all__ = ["ALLOWED_COLUMNS", "CATEGORICAL_MAPPING", "DFQuery", "DFFilter", "safe_df_query"]

MAX_QUERY_LIMIT = 100
TAX_HISTORY_PARQUET_PATH = 'data/extracted/tax_data.parquet'

logger = logging.getLogger(__name__)

# ----------------------------------------------------------------------
# 1. CONSTRAINED LITERALS AND COLUMN DEFINITIONS
# ----------------------------------------------------------------------

# Define all allowed categorical values for robust filtering and validation
TAXPAYER_TYPES = Literal['Non-Profit', 'Partnership', 'Corporation', 'Individual', 'Trust']
INCOME_SOURCES = Literal['Investment', 'Royalties', 'Business Income', 'Capital Gains', 'Salary', 'Rental']
DEDUCTION_TYPES = Literal['Charitable Contributions', 'Mortgage Interest', 'Education Expenses', 'Business Expenses', 'Medical Expenses']
STATES = Literal['IL', 'PA', 'GA', 'TX', 'CA', 'FL', 'OH', 'NY', 'NC', 'MI']

# MASTER LIST OF ALL ALLOWED COLUMNS for validation
ALLOWED_COLUMNS = {
    "tax_year",
    "transaction_month", # added for filtering by month (numeric value 1-12)
    "taxpayer_type",
    "income_source",
    "deduction_type",
    "state",

    "income",
    "deductions",
    "taxable_income",
    "tax_rate",
    "tax_owed",
}

# Mapping for filter sanitization (allows LLM flexibility, ensures data integrity)
# Keys must be lowercased versions of the canonical values
CATEGORICAL_MAPPING = {
    "taxpayer_type": {s.lower().replace('-', '_'): s for s in TAXPAYER_TYPES.__args__},         # 'non-profit' -> 'Non-Profit'
    "income_source": {s.lower().replace(' ', '_'): s for s in INCOME_SOURCES.__args__},         # 'business income' -> 'Business Income'
    "deduction_type": {s.lower().replace(' ', '_'): s for s in DEDUCTION_TYPES.__args__},       # 'education expenses' -> 'Education Expenses'
    "state": {s.lower(): s for s in STATES.__args__},                                           # 'il' -> 'IL'
}

# Map the string operator to the actual Python function
OP_MAP = {
    "==": operator.eq,
    ">": operator.gt,
    "<": operator.lt,
    ">=": operator.ge,
    "<=": operator.le,
    "!=": operator.ne,
}

@lru_cache(maxsize=1)
def get_tax_history_df() -> pd.DataFrame:
    """
    Retrieves the tax history DataFrame from the database.
    """
    return pd.read_parquet(TAX_HISTORY_PARQUET_PATH)

# ----------------------------------------------------------------------
# 2. PYDANTIC MODELS
# ----------------------------------------------------------------------

class DFFilter(BaseModel):
    column: str
    value: Union[str, int, float]
    operator: Literal["==", ">", "<", ">=", "<=", "!="] = "=="

class DFQuery(BaseModel):
    operation: Literal["select", "filter", "aggregate", "filtered_aggregate"]

    # FIX: target_column now accepts a string OR a list of strings
    target_column: Optional[Union[str, List[str]]] = None
    filters: List[DFFilter] = Field(default_factory=list)
    agg: List[Literal["count", "sum", "mean", "min", "max"]] = Field(default_factory=list)
    group_by: Optional[Union[str, List[str]]] = None
    limit: int = Field(default=20, ge=1)

    # Validation for column names and clamping (as previously defined)
    def model_post_init(self, __context=None):
        self._validate_allowed_columns()

    def _validate_allowed_columns(self):
        cols = set()
        cols.update({f.column for f in self.filters})
        if self.group_by:
            if isinstance(self.group_by, str):
                cols.add(self.group_by)
            elif isinstance(self.group_by, list):
                cols.update(self.group_by)
        if self.target_column:
            if isinstance(self.target_column, str):
                cols.add(self.target_column)
            elif isinstance(self.target_column, list):
                cols.update(self.target_column)

        for col in cols:
            if col not in ALLOWED_COLUMNS:
                raise ValueError(f"Column '{col}' not allowed by the schema.")
        return self

    @field_validator('limit', mode='before')
    @classmethod
    def clamp_limit(cls, v: int) -> int:
        """Coerces the limit value to be no more than MAX_QUERY_LIMIT."""
        if v > MAX_QUERY_LIMIT:
            print(f"Warning: Requested limit of {v} clamped to {MAX_QUERY_LIMIT}.")
            return MAX_QUERY_LIMIT
        return v

# ----------------------------------------------------------------------
# 3. UTILITY FUNCTIONS
# ----------------------------------------------------------------------

def _sanitize_categorical_filter(column: str, value: Union[str, int, float]) -> Union[str, int, float]:
    """Cleans and validates filter values for categorical columns."""
    if column in CATEGORICAL_MAPPING and isinstance(value, str):
        # Normalize common LLM/user variants:
        # - case insensitivity
        # - spaces / hyphens / underscores treated equivalently
        normalized_input = value.lower().strip()
        normalized_input = normalized_input.replace("-", "_").replace(" ", "_")

        if normalized_input in CATEGORICAL_MAPPING[column]:
            # Return the canonical, correctly cased value
            return CATEGORICAL_MAPPING[column][normalized_input]
        else:
            raise ValueError(
                f"Invalid value for categorical column '{column}': '{value}'. "
                f"Allowed values are: {', '.join(CATEGORICAL_MAPPING[column].values())}"
            )

    return value

def safe_df_query(df: pd.DataFrame, q: DFQuery) -> Dict[str, Any] | List[Dict[str, Any]]:
    """
    Executes a Pydantic-validated query against the DataFrame.
    """

    q._validate_allowed_columns()

    # Observability: log the validated query right before execution.
    try:
        logger.info("Executing DFQuery: %s", q.model_dump())
    except Exception:
        logger.info("Executing DFQuery (repr): %r", q)

    working_df = df.copy()

    # --- Apply Filters ---
    if q.filters:
        combined_mask = pd.Series([True] * len(working_df))

        for f in q.filters:
            if f.column not in working_df.columns:
                raise ValueError(f"Filter column {f.column} not found in DataFrame.")

            if f.operator == "==":
                sanitized_value = _sanitize_categorical_filter(f.column, f.value)
            else:
                sanitized_value = f.value

            op_func = OP_MAP.get(f.operator)
            if op_func is None:
                raise ValueError(f"Unsupported operator: {f.operator}")

            # Apply the dynamic comparison using the operator function
            combined_mask &= op_func(working_df[f.column], sanitized_value)

        working_df = working_df.loc[combined_mask]

    # --- Aggregation Operations ---
    if q.operation in ("aggregate", "filtered_aggregate"):
        if not q.target_column or not q.agg:
            raise ValueError("target_column and agg required for aggregation.")

        # 1. Determine the structure of the aggregation request
        target_cols_list = [q.target_column] if isinstance(q.target_column, str) else q.target_column
        # Deduplicate agg functions to prevent duplicate column names
        agg_funcs = list(dict.fromkeys(q.agg))  # Preserve order while removing duplicates

        # 2. Construct the Aggregation Map for Pandas: {'amount': ['max', 'min', 'mean'], 'tax_owed': ['max', 'min', 'mean']}
        # This structure is needed for multi-column aggregation.
        pandas_agg_map = {col: agg_funcs for col in target_cols_list}

        if q.group_by:
            # Multi-column/Grouped Aggregation
            result = working_df.groupby(q.group_by).agg(pandas_agg_map)

            # Create unique column names before flattening to prevent duplicates
            original_columns = result.columns  # MultiIndex like [('deductions', 'mean'), ('tax_owed', 'mean')]

            # Generate unique names based on column and aggregation
            unique_names = []
            seen = set()
            for col_name, agg_name in original_columns:
                candidate_name = f"{agg_name}_{col_name}"
                if candidate_name not in seen:
                    unique_names.append(candidate_name)
                    seen.add(candidate_name)
                else:
                    # If still duplicate, add counter
                    counter = 1
                    while f"{candidate_name}_{counter}" in seen:
                        counter += 1
                    unique_name = f"{candidate_name}_{counter}"
                    unique_names.append(unique_name)
                    seen.add(unique_name)

            # Apply the unique names
            result.columns = unique_names

            # Use .to_json() for safe grouped serialization
            return json.loads(result.reset_index().rename_axis(None, axis=1).to_json(orient='records'))

        else:
            # Multi-column/Scalar Aggregation (No GroupBy)
            result_df = working_df.agg(pandas_agg_map)

            # Flatten the resulting DataFrame (it's a single-row DataFrame)
            results_dict = {}
            for col in target_cols_list:
                for agg in agg_funcs:
                    key = f"{agg}_{col}"
                    # Safely extract scalar value and ensure JSON serializability
                    agg_result = result_df.loc[agg, col]

                    if hasattr(agg_result, 'item') and not isinstance(agg_result, (str, bytes)):
                        results_dict[key] = agg_result.item()
                    else:
                        results_dict[key] = agg_result

            return results_dict

    # --- Select/Filter Operations (FINAL FIX: Replaced .sample() with .head()) ---

    if q.operation == "select":
        if not q.target_column:
            raise ValueError("target_column required for select.")

        target_cols_list = [q.target_column] if isinstance(q.target_column, str) else q.target_column

        # 1. Select the relevant columns
        result_df = working_df[target_cols_list].copy()

        # 2. FIX: ALWAYS drop duplicates on any select operation.
        # This ensures the output only contains unique rows, whether one column
        # (like 'state') or multiple columns (like ['state', 'taxpayer_type']).
        result_df = result_df.drop_duplicates()

        # 3. Use .head() for deterministic limiting
        return result_df.head(q.limit).to_dict('records')

    if q.operation == "filter":
        return working_df.head(q.limit).to_dict('records')

    raise ValueError(f"Unsupported operation: {q.operation}")