import polars as pl
import duckdb
from tqdm import tqdm
from core_eda import TextEDA
from rich import print


class PipelineText:
    def __init__(self, mode: str = ''):
        self.mode = mode
        print(f'[Text Cleaning] {mode}')

    @staticmethod
    def clean_text(data: pl.DataFrame, col: str = 'item_name') -> pl.DataFrame:
        lst = [TextEDA.clean_text_pipeline(str(x)) for x in tqdm(data[col].to_list(), desc='Clean Text')]
        return data.with_columns(pl.Series(name=f'{col}_clean', values=lst))

    def run(self, data, key_col: list = None):
        # load data
        query = f"""select * from data"""
        df = duckdb.sql(query).pl()
        print(f'-> Base Data {self.mode}: {df.shape}')

        df = (
            df
            .pipe(PipelineText.clean_text)
            .drop_nulls(subset=key_col)
            .select(pl.all().name.prefix(f'{self.mode}_'))
        )
        print(f'-> Join Data {self.mode}: {df.shape}')
        return df
