import os
import pandas as pd
import yfinance as yf
from datetime import datetime


ASSET = "BTC-USD"
START_DATE = "2015-01-01"
DATA_DIR = "/opt/airflow/data/raw"


def download_data(asset: str, start_date: str) -> pd.DataFrame:
    print(f"Downloading data for {asset} from {start_date}...")

    df = yf.download(asset, start=start_date, progress=False)

    if df.empty:
        raise ValueError("Downloaded dataframe is empty!")

    df.reset_index(inplace=True)

    # Standardize column names# Flatten columns if MultiIndex
    if isinstance(df.columns, pd.MultiIndex):
        df.columns = df.columns.get_level_values(0)

    df.columns = [str(col).lower().replace(" ", "_") for col in df.columns]

    # Add ingestion timestamp
    df["ingestion_timestamp"] = datetime.utcnow()

    print(f"Downloaded {len(df)} rows.")

    return df


def save_to_parquet(df: pd.DataFrame, asset: str):
    os.makedirs(DATA_DIR, exist_ok=True)

    file_path = os.path.join(DATA_DIR, f"{asset}.parquet")

    df.to_parquet(file_path, index=False)

    print(f"Data saved to {file_path}")


if __name__ == "__main__":
    df = download_data(ASSET, START_DATE)
    save_to_parquet(df, ASSET)