app/services/preprocessing.py

"""
Data preprocessing pipeline.
Handles loading, cleaning, and city extraction from the places CSV dataset.
"""

import logging
import json
from pathlib import Path
from typing import Optional

import pandas as pd

from app.core.constants import EGYPT_CITIES

logger = logging.getLogger(__name__)


def extract_city(address: Optional[str]) -> Optional[str]:
    """
    Extract an Egyptian city name from a raw address string.

    Args:
        address: Raw address text, may be None / NaN.

    Returns:
        Matched city name or None.
    """
    if pd.isna(address) or not address:
        return None

    for city in EGYPT_CITIES:
        if city.lower() in address.lower():
            return city

    return None


def load_and_preprocess(csv_path: str) -> pd.DataFrame:
    """
    Load the places CSV and add a 'city' column derived from addresses.

    Args:
        csv_path: Path to the raw CSV file.

    Returns:
        Preprocessed DataFrame with a 'city' column.
    """
    logger.info("Loading dataset from %s", csv_path)
    data = pd.read_csv(csv_path)
    logger.info("Loaded %d rows, columns: %s", len(data), list(data.columns))

    data["city"] = data["address"].apply(extract_city)
    logger.info("City extraction complete. Distribution:\n%s", data["city"].value_counts())

    return data


def filter_by_city(data: pd.DataFrame, city: str) -> pd.DataFrame:
    """
    Filter the DataFrame to rows belonging to a specific city.

    Args:
        data: Preprocessed DataFrame.
        city: City name to filter on (case-insensitive match against 'city' column).

    Returns:
        Filtered DataFrame.
    """
    filtered = data[data["city"].str.lower() == city.lower()]
    logger.info("Filtered to %d rows for city '%s'", len(filtered), city)
    return filtered


def dataframe_to_json_string(df: pd.DataFrame) -> str:
    """
    Serialize a DataFrame to a compact JSON string (records orientation).

    Args:
        df: DataFrame to serialize.

    Returns:
        JSON string.
    """
    return df.to_json(orient="records", force_ascii=False)


def save_city_json(df: pd.DataFrame, output_path: str) -> None:
    """
    Persist a city-filtered DataFrame as a JSON file.

    Args:
        df: DataFrame for one city.
        output_path: Destination file path.
    """
    path = Path(output_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    json_str = dataframe_to_json_string(df)
    with open(path, "w", encoding="utf-8") as f:
        f.write(json_str)
    logger.info("Saved city JSON to %s", path)


def load_city_json(json_path: str) -> pd.DataFrame:
    """
    Load a previously saved city JSON back into a DataFrame.

    Args:
        json_path: Path to the JSON file.

    Returns:
        DataFrame.
    """
    logger.info("Loading city JSON from %s", json_path)
    return pd.read_json(json_path)