from __future__ import annotations

"""Runnable pipeline demo showing the current `DataTable` workflow.

This example is intentionally small but complete:
- `load_data()` loads a CSV into `self.datatable`
- `self.datatable` is used for all in-memory transforms
- `push_datatable()` writes into the pipeline-managed SQLite database
- `save_to_csv()` exports the current table
- `load_datatable()` reads the SQLite table back into `self.datatable`
"""

from pathlib import Path

from soigia.pipeline import Pipeline


DEMO_SOURCE_CSV = Path(__file__).with_name("demo_users.csv")
DEMO_UPSERT_CSV = Path(__file__).with_name("demo_upsert_users.csv")


class Shopee(Pipeline):
    _name = "Shopee"
    _inherit = ["sale", "purchase"]
    _data = ["product", "user", "customer"]
    cli_show_datatable_default = True

    stages = ["normalize", "save_users", "export_users", "reload_users"]

    def load_data(self):
        self.datatable = self.datatable.load_from_csv(DEMO_SOURCE_CSV)

    def normalize(self):
        self.datatable = self.datatable.copy()
        self.datatable["name"] = self.datatable["name"].astype(str).str.strip().str.title()
        self.datatable["status"] = self.datatable["status"].astype(str).str.lower()
        self.datatable["net_amount"] = (self.datatable["amount"] * 0.98).round(2)

    def save_users(self):
        self.push_datatable("users", self.datatable, mode="replace")

    def export_users(self):
        self.datatable.save_to_csv("data/demo/users.csv")

    def reload_users(self):
        self.datatable = self.load_datatable("users").sort_values("id").reset_index(drop=True)


class ShopeeUpsertDemo(Pipeline):
    stages = ["seed_users", "apply_changes", "reload_users", "export_vip_user"]

    def load_data(self):
        self.datatable = self.datatable.load_from_csv(DEMO_UPSERT_CSV)

    def seed_users(self):
        self.push_datatable("users", self.datatable, mode="replace")

    def apply_changes(self):
        # Upsert keeps id=2 and inserts id=3.
        self.upsert_rows(
            "users",
            [
                {"id": 2, "name": "bob", "status": "vip", "amount": 250.0},
                {"id": 3, "name": "carol", "status": "active", "amount": 300.0},
            ],
            key_columns=["id"],
            mode="upsert",
        )

    def reload_users(self):
        self.datatable = self.load_datatable("users").sort_values("id").reset_index(drop=True)

    def export_vip_user(self):
        # This is the chaining style you asked for:
        # filter(...) -> get(...) -> save_to_csv(...)
        self.datatable.filter(status="vip").get(id=2).save_to_csv("data/demo/user_2_vip.csv")


if __name__ == "__main__":
    Shopee.cli()
