import asyncio
from pathlib import Path
from typing import Dict

import pandas as pd
from loguru import logger

from bb_integrations_lib.protocols.pipelines import Step
from bb_integrations_lib.shared.model import FileReference, FileType, RawData
from pandas import DataFrame


class LoadFileToDataFrameStep(Step):
    def __init__(self, sheet_name: str | int = 0, file_type: FileType | None = None, save_to: str | None = None, has_header: bool = True, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.sheet_name = sheet_name
        self.file_type = file_type
        self.save_to = save_to
        self.has_header = has_header

    def describe(self) -> str:
        return "Load file into dataframe"

    async def execute(self, i: FileReference | RawData) -> DataFrame:
        if isinstance(i, FileReference):
            if i.file_type == FileType.excel:
                sheet = i.sheet_name if i.sheet_name is not None else 0
                df = pd.read_excel(i.file_path, sheet_name=sheet)
            elif i.file_type == FileType.csv:
                df = pd.read_csv(i.file_path, header=None if not self.has_header else 0)
            else:
                raise NotImplementedError()
        elif isinstance(i, RawData):
            if self.file_type == FileType.csv or i.file_name.endswith("csv"):
                df = pd.read_csv(i.data, header=None if not self.has_header else 0)
            elif i.file_name.endswith("xlsx") or i.file_name.endswith("xls"):
                df = pd.read_excel(i.data, sheet_name=0)
            else:
                raise NotImplementedError()

        if self.save_to is not None:
            if isinstance(i, RawData):
                file_name = i.file_name
            else:
                file_name = Path(i.file_path).name
            save_dir = Path.cwd() / self.save_to
            save_dir.mkdir(parents=True, exist_ok=True)
            save_path = save_dir / file_name
            df.to_csv(save_path, index=False, header=self.has_header)
            logger.info(f"Saved file to {save_path}")

        return df

if __name__ == "__main__":
    async def main():
        input = FileReference("/home/ben-allen/Downloads/herdrich-payroll-test.xlsx", FileType.excel, "payroll")
        output = await LoadFileToDataFrameStep().execute(input)
        print(output.head(5))

    asyncio.run(main())