"""
File storage operations for datasets.

Handles reading and writing datasets in various formats (CSV, Excel, JSON, Parquet).
"""

import json
from pathlib import Path
from typing import Any

import pandas as pd
import polars as pl
import pyarrow.parquet as pq

from app.core.config import settings
from app.core.exceptions import (
    FileNotFoundError,
    FileWriteError,
    InvalidDatasetFormatError,
)
from app.core.logging import get_logger

logger = get_logger(__name__)


class FileStorage:
    """File storage operations for datasets."""

    def __init__(self) -> None:
        """Initialize file storage."""
        self.upload_dir = settings.upload_dir
        self.processed_dir = settings.processed_dir
        self.export_dir = settings.export_dir
        self.temp_dir = settings.temp_dir

    def save_uploaded_file(self, file_content: bytes, filename: str) -> Path:
        """
        Save uploaded file to storage.

        Args:
            file_content: File content as bytes
            filename: Original filename

        Returns:
            Path to saved file

        Raises:
            FileWriteError: If file write fails
        """
        try:
            file_path = self.upload_dir / filename
            file_path.write_bytes(file_content)
            logger.info("File saved", path=str(file_path), size=len(file_content))
            return file_path
        except Exception as e:
            logger.error("File save failed", filename=filename, error=str(e))
            raise FileWriteError(f"Failed to save file: {filename}", {"error": str(e)})

    def read_csv(self, file_path: Path, engine: str = "pandas") -> pd.DataFrame | pl.DataFrame:
        """
        Read CSV file.

        Args:
            file_path: Path to CSV file
            engine: Engine to use ('pandas' or 'polars')

        Returns:
            DataFrame

        Raises:
            FileNotFoundError: If file doesn't exist
            InvalidDatasetFormatError: If file is malformed
        """
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        try:
            if engine == "polars":
                return pl.read_csv(file_path)
            else:
                return pd.read_csv(file_path)
        except Exception as e:
            logger.error("CSV read failed", path=str(file_path), error=str(e))
            raise InvalidDatasetFormatError(f"Invalid CSV file: {file_path}", {"error": str(e)})

    def read_excel(self, file_path: Path, engine: str = "pandas") -> pd.DataFrame | pl.DataFrame:
        """
        Read Excel file.

        Args:
            file_path: Path to Excel file
            engine: Engine to use ('pandas' or 'polars')

        Returns:
            DataFrame

        Raises:
            FileNotFoundError: If file doesn't exist
            InvalidDatasetFormatError: If file is malformed
        """
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        try:
            if engine == "polars":
                # Polars doesn't have native Excel support, use pandas and convert
                df_pandas = pd.read_excel(file_path)
                return pl.from_pandas(df_pandas)
            else:
                return pd.read_excel(file_path)
        except Exception as e:
            logger.error("Excel read failed", path=str(file_path), error=str(e))
            raise InvalidDatasetFormatError(f"Invalid Excel file: {file_path}", {"error": str(e)})

    def read_json(self, file_path: Path, engine: str = "pandas") -> pd.DataFrame | pl.DataFrame:
        """
        Read JSON file.

        Args:
            file_path: Path to JSON file
            engine: Engine to use ('pandas' or 'polars')

        Returns:
            DataFrame

        Raises:
            FileNotFoundError: If file doesn't exist
            InvalidDatasetFormatError: If file is malformed
        """
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        try:
            if engine == "polars":
                return pl.read_json(file_path)
            else:
                return pd.read_json(file_path)
        except Exception as e:
            logger.error("JSON read failed", path=str(file_path), error=str(e))
            raise InvalidDatasetFormatError(f"Invalid JSON file: {file_path}", {"error": str(e)})

    def read_parquet(
        self, file_path: Path, engine: str = "pandas"
    ) -> pd.DataFrame | pl.DataFrame:
        """
        Read Parquet file.

        Args:
            file_path: Path to Parquet file
            engine: Engine to use ('pandas' or 'polars')

        Returns:
            DataFrame

        Raises:
            FileNotFoundError: If file doesn't exist
            InvalidDatasetFormatError: If file is malformed
        """
        if not file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        try:
            if engine == "polars":
                return pl.read_parquet(file_path)
            else:
                return pd.read_parquet(file_path)
        except Exception as e:
            logger.error("Parquet read failed", path=str(file_path), error=str(e))
            raise InvalidDatasetFormatError(
                f"Invalid Parquet file: {file_path}", {"error": str(e)}
            )

    def read_dataset(
        self, file_path: Path, file_format: str, engine: str = "pandas"
    ) -> pd.DataFrame | pl.DataFrame:
        """
        Read dataset in any supported format.

        Args:
            file_path: Path to file
            file_format: File format (csv, xlsx, json, parquet)
            engine: Engine to use ('pandas' or 'polars')

        Returns:
            DataFrame

        Raises:
            InvalidDatasetFormatError: If format is not supported
        """
        format_handlers = {
            "csv": self.read_csv,
            "xlsx": self.read_excel,
            "xls": self.read_excel,
            "json": self.read_json,
            "parquet": self.read_parquet,
        }

        handler = format_handlers.get(file_format.lower())
        if not handler:
            raise InvalidDatasetFormatError(
                f"Unsupported file format: {file_format}",
                {"supported": list(format_handlers.keys())},
            )

        return handler(file_path, engine=engine)

    def write_csv(self, df: pd.DataFrame | pl.DataFrame, file_path: Path) -> Path:
        """Write DataFrame to CSV."""
        try:
            if isinstance(df, pl.DataFrame):
                df.write_csv(file_path)
            else:
                df.to_csv(file_path, index=False)
            return file_path
        except Exception as e:
            raise FileWriteError(f"Failed to write CSV: {file_path}", {"error": str(e)})

    def write_excel(self, df: pd.DataFrame | pl.DataFrame, file_path: Path) -> Path:
        """Write DataFrame to Excel."""
        try:
            if isinstance(df, pl.DataFrame):
                df = df.to_pandas()
            df.to_excel(file_path, index=False)
            return file_path
        except Exception as e:
            raise FileWriteError(f"Failed to write Excel: {file_path}", {"error": str(e)})

    def write_parquet(self, df: pd.DataFrame | pl.DataFrame, file_path: Path) -> Path:
        """Write DataFrame to Parquet."""
        try:
            if isinstance(df, pl.DataFrame):
                df.write_parquet(file_path)
            else:
                df.to_parquet(file_path, index=False)
            return file_path
        except Exception as e:
            raise FileWriteError(f"Failed to write Parquet: {file_path}", {"error": str(e)})

    def write_json(self, df: pd.DataFrame | pl.DataFrame, file_path: Path) -> Path:
        """Write DataFrame to JSON."""
        try:
            if isinstance(df, pl.DataFrame):
                df.write_json(file_path)
            else:
                df.to_json(file_path, orient="records", indent=2)
            return file_path
        except Exception as e:
            raise FileWriteError(f"Failed to write JSON: {file_path}", {"error": str(e)})

    def write_dataset(
        self, df: pd.DataFrame | pl.DataFrame, file_path: Path, file_format: str
    ) -> Path:
        """
        Write dataset in specified format.

        Args:
            df: DataFrame to write
            file_path: Output path
            file_format: Format (csv, xlsx, json, parquet)

        Returns:
            Path to written file

        Raises:
            InvalidDatasetFormatError: If format not supported
        """
        format_handlers = {
            "csv": self.write_csv,
            "xlsx": self.write_excel,
            "xls": self.write_excel,
            "json": self.write_json,
            "parquet": self.write_parquet,
        }

        handler = format_handlers.get(file_format.lower())
        if not handler:
            raise InvalidDatasetFormatError(
                f"Unsupported output format: {file_format}",
                {"supported": list(format_handlers.keys())},
            )

        return handler(df, file_path)

    def delete_file(self, file_path: Path) -> bool:
        """Delete a file from storage."""
        try:
            if file_path.exists():
                file_path.unlink()
                logger.info("File deleted", path=str(file_path))
                return True
            return False
        except Exception as e:
            logger.error("File deletion failed", path=str(file_path), error=str(e))
            return False

    def get_file_size(self, file_path: Path) -> int:
        """Get file size in bytes."""
        if file_path.exists():
            return file_path.stat().st_size
        return 0
