import pandas as pd
from datetime import datetime, UTC
import os

RAW_DIR = "/opt/airflow/data/raw"
CLEAN_DIR = "/opt/airflow/data/processed"

def clean_data(asset: str):
    raw_file = os.path.join(RAW_DIR, f"{asset}.parquet")
    
    if not os.path.exists(raw_file):
        print(f"Raw file {raw_file} not found! Nothing to clean.")
        return

    df = pd.read_parquet(raw_file)

    if df.empty:
        print("Raw dataframe is empty! Nothing to clean.")
        return

    # Drop NaNs
    df.dropna(inplace=True)

    # Проверяем отрицательные цены
    for col in ["open", "high", "low", "close", "adj_close"]:
        if col in df.columns:
            df = df[df[col] >= 0]

    # Добавляем колонку день недели
    df["day_of_week"] = df["date"].dt.dayofweek

    df["cleaned_timestamp"] = datetime.now(UTC)

    os.makedirs(CLEAN_DIR, exist_ok=True)
    clean_file = os.path.join(CLEAN_DIR, f"{asset}_clean.parquet")
    df.to_parquet(clean_file, index=False)

    print(f"Cleaned data saved to {clean_file}")


# 🔹 Добавляем main
if __name__ == "__main__":
    clean_btc_data()