"""Clean, type-cast, and enrich Olist DataFrames before loading to staging."""

import pandas as pd

# Brazil bounding box for coordinate clipping
_LAT_MIN, _LAT_MAX = -33.75, 5.27
_LNG_MIN, _LNG_MAX = -73.99, -34.79

TIMESTAMP_COLS = {
    "orders": [
        "order_purchase_timestamp",
        "order_approved_at",
        "order_delivered_carrier_date",
        "order_delivered_customer_date",
        "order_estimated_delivery_date",
    ],
    "order_reviews": ["review_creation_date", "review_answer_timestamp"],
    "order_items": ["shipping_limit_date"],
}

FLOAT_COLS = {
    "order_items": ["price", "freight_value"],
    "order_payments": ["payment_value"],
    "geolocation": ["geolocation_lat", "geolocation_lng"],
}

INT_COLS = {
    "order_payments": ["payment_sequential", "payment_installments"],
    "order_reviews": ["review_score"],
    "order_items": ["order_item_id"],
}


class OlistTransformer:
    def transform_orders(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        for col in TIMESTAMP_COLS["orders"]:
            df[col] = pd.to_datetime(df[col], errors="coerce")
            df[col] = df[col].dt.strftime("%Y-%m-%d %H:%M:%S").where(df[col].notna(), None)
        df = df.drop_duplicates(subset=["order_id"])
        return df

    def transform_order_items(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        for col in FLOAT_COLS["order_items"]:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        df["order_item_id"] = pd.to_numeric(df["order_item_id"], errors="coerce")
        df["shipping_limit_date"] = pd.to_datetime(
            df["shipping_limit_date"], errors="coerce"
        ).dt.strftime("%Y-%m-%d %H:%M:%S")
        return df

    def transform_order_payments(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        for col in ["payment_sequential", "payment_installments"]:
            df[col] = pd.to_numeric(df[col], errors="coerce")
        df["payment_value"] = pd.to_numeric(df["payment_value"], errors="coerce")
        return df

    def transform_order_reviews(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df["review_score"] = pd.to_numeric(df["review_score"], errors="coerce")
        for col in TIMESTAMP_COLS["order_reviews"]:
            df[col] = pd.to_datetime(df[col], errors="coerce")
            df[col] = df[col].dt.strftime("%Y-%m-%d %H:%M:%S").where(df[col].notna(), None)
        df = df.drop_duplicates(subset=["review_id"])
        return df

    def transform_customers(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df = df.drop_duplicates(subset=["customer_id"])
        return df

    def transform_sellers(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df = df.drop_duplicates(subset=["seller_id"])
        return df

    def transform_products(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        numeric_dim_cols = [
            "product_weight_g",
            "product_length_cm",
            "product_height_cm",
            "product_width_cm",
            "product_photos_qty",
            "product_name_lenght",
            "product_description_lenght",
        ]
        for col in numeric_dim_cols:
            df[col] = pd.to_numeric(df[col], errors="coerce")
            median = df[col].median()
            df[col] = df[col].fillna(median)
        df = df.drop_duplicates(subset=["product_id"])
        return df

    def transform_category_translations(self, df: pd.DataFrame) -> pd.DataFrame:
        return df.drop_duplicates(subset=["product_category_name"])

    def transform_geolocation(self, df: pd.DataFrame) -> pd.DataFrame:
        df = df.copy()
        df["geolocation_lat"] = pd.to_numeric(df["geolocation_lat"], errors="coerce")
        df["geolocation_lng"] = pd.to_numeric(df["geolocation_lng"], errors="coerce")
        # Clip to Brazil bounding box
        df["geolocation_lat"] = df["geolocation_lat"].clip(_LAT_MIN, _LAT_MAX)
        df["geolocation_lng"] = df["geolocation_lng"].clip(_LNG_MIN, _LNG_MAX)
        # Dedup zip codes: keep median lat/lng per prefix
        df = (
            df.groupby("geolocation_zip_code_prefix", as_index=False)
            .agg(
                geolocation_lat=("geolocation_lat", "median"),
                geolocation_lng=("geolocation_lng", "median"),
                geolocation_city=("geolocation_city", "first"),
                geolocation_state=("geolocation_state", "first"),
            )
        )
        return df

    def transform_all(
        self, raw: dict[str, pd.DataFrame]
    ) -> dict[str, pd.DataFrame]:
        return {
            "orders": self.transform_orders(raw["orders"]),
            "order_items": self.transform_order_items(raw["order_items"]),
            "order_payments": self.transform_order_payments(raw["order_payments"]),
            "order_reviews": self.transform_order_reviews(raw["order_reviews"]),
            "customers": self.transform_customers(raw["customers"]),
            "sellers": self.transform_sellers(raw["sellers"]),
            "products": self.transform_products(raw["products"]),
            "category_translations": self.transform_category_translations(
                raw["category_translations"]
            ),
            "geolocation": self.transform_geolocation(raw["geolocation"]),
        }