import pandas as pd
import zipfile
import os

from abc import ABC, abstractmethod
from pathlib import Path
from typing import Union, List, Optional
from emissary_ml_toolkit.config.storage import StorageConfig
from io import IOBase, StringIO, BytesIO
# from emissary_ml_toolkit.utils.exceptions import StorageError


class BaseStorage(ABC):
    def __init__(self, config: StorageConfig):
        self.config = config
        self.connection = None
        self._setup_client()

    @abstractmethod
    def _setup_client(self) -> None:
        """Initialize the storage client"""
        pass

    @abstractmethod
    def download_file(self, storage_path: str, local_path: Optional[Union[str, Path]] = None) -> Path:
        """
        Download a file from storage to local path

        Args:
            storage_path: Full path in storage (e.g., "bucket/folder/file.txt" for S3,
                          "container/blob" for Azure, "bucket/object" for GCS)
            local_path: Local destination path
        """
        pass

    @abstractmethod
    def upload_file(self, source: Union[str, Path, IOBase, bytes], storage_path: str) -> None:
        """
        Upload a local file to storage

        Args:
            source: Local source
            storage_path: Full destination path in storage
        """
        pass

    @abstractmethod
    def list_objects(self, prefix: str = "", bucket: Optional[str] = None) -> List[str]:
        """List objects in storage with optional prefix filter"""
        pass

    @abstractmethod
    def exists(self, storage_path: str) -> bool:
        """Check if object exists in storage"""
        pass

    # Utility methods that work across all storage types
    def download_directory(self, prefix: str, local_dir: Union[str, Path]) -> Path:
        """Download all files with given prefix to local directory"""
        local_dir = Path(local_dir)
        local_dir.mkdir(parents=True, exist_ok=True)

        objects = self.list_objects(prefix)
        for obj in objects:
            # Create relative path structure
            rel_path = Path(obj.key).relative_to(prefix) if prefix else Path(obj.key)
            local_file_path = local_dir / rel_path
            local_file_path.parent.mkdir(parents=True, exist_ok=True)
            self.download_file(obj.key, local_file_path)

        return local_dir

    def upload_dataframe(self, df: pd.DataFrame, storage_path: str, **to_format_kwargs):
        file_format = storage_path.split(".")[-1]
        if file_format in ['parquet']:
            buffer = BytesIO()
            getattr(df, f'to_{file_format}')(buffer, **to_format_kwargs)
        else:
            buffer = StringIO()
            if file_format == 'jsonl':
                df.to_json(buffer, orient='records', lines=True, **to_format_kwargs)
            else:
                getattr(df, f'to_{file_format}')(buffer, **to_format_kwargs)

        buffer.seek(0)
        self.upload_file(buffer, storage_path)
        return

    def zip_directory(self, source: Union[str, Path], zip_path: Union[str, Path]):
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zip_file:
            for root, dirs, files in os.walk(source):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, source)
                    zip_file.write(file_path, arcname)

        return zip_path
