import os
import pandas as pd

class DataStorage:
    def __init__(self, base_path: str = "data"):
        self.base_path = base_path

    def _ensure_dir(self, path: str):
        os.makedirs(path, exist_ok=True)

    def save_data(self, df: pd.DataFrame, market: str, symbol: str, statement_type: str):
        if df.empty:
            print(f"Warning: No data to save for {market} {symbol} {statement_type}")
            return

        # Prepare directory: data/{market}/{symbol}
        # Sanitize symbol for filename (e.g., remove special chars if any, though usually fine)
        dir_path = os.path.join(self.base_path, market, symbol)
        self._ensure_dir(dir_path)
        
        file_path = os.path.join(dir_path, f"{statement_type}.csv")
        df.to_csv(file_path, index=False)
        print(f"Saved {file_path}")

    def load_data(self, market: str, symbol: str, statement_type: str) -> pd.DataFrame:
        file_path = os.path.join(self.base_path, market, symbol, f"{statement_type}.csv")
        if not os.path.exists(file_path):
            print(f"File not found: {file_path}")
            return pd.DataFrame()
        return pd.read_csv(file_path)
