"""Utility functions for the mlfastflow package."""

import os
import time
import warnings
from functools import wraps
from pathlib import Path
from typing import List, Optional, Union, Any

try:
    import polars as pl
except ImportError:
    pl = None

try:
    import pandas as pd
except ImportError:
    pd = None


def timer_decorator(func):
    """Decorator that logs the execution time of the decorated function.

    Time is displayed in a human-friendly format: milliseconds for fast
    calls, seconds for moderate ones, and minutes for long-running ones.

    Args:
        func: The function to be decorated.

    Returns:
        The wrapped function.
    """
    import logging

    logger = logging.getLogger(__name__)

    @wraps(func)
    def wrapper_timer(*args, **kwargs):
        start_time = time.perf_counter()
        try:
            return func(*args, **kwargs)
        finally:
            elapsed = time.perf_counter() - start_time
            if elapsed < 1:
                time_str = f"{elapsed * 1000:.1f} ms"
            elif elapsed < 60:
                time_str = f"{elapsed:.2f} s"
            else:
                minutes, secs = divmod(elapsed, 60)
                time_str = f"{int(minutes)}m {secs:.1f}s"
            logger.info("Finished %r in %s", func.__name__, time_str)

    return wrapper_timer


def concat_files(
    folder_path: Union[str, Path],
    file_type: str = "csv",
    add_source_column: bool = False,
    output_file: Optional[Union[str, Path]] = None,
    how: str = "diagonal_relaxed",
) -> str:
    """Concatenate all files of a specific type in a folder and its subfolders.

    Uses Polars lazy scanning where possible so that files are streamed
    through memory rather than loaded all at once.

    Args:
        folder_path: Path to the folder containing files to concatenate.
        file_type: File extension to look for (``'csv'`` or ``'parquet'``).
        add_source_column: If ``True``, adds a ``SOURCE`` column with the
            source filename.
        output_file: Path for the combined output file. If ``None``, saves
            as ``<folder_name>_combined.<file_type>`` next to *folder_path*.
        how: Strategy for handling schema mismatches between files.
            ``'diagonal_relaxed'`` (default) fills missing columns with
            nulls and casts types. ``'vertical'`` requires identical schemas.

    Returns:
        Path to the created output file.

    Raises:
        ImportError: If Polars is not installed.
        FileNotFoundError: If *folder_path* does not exist.
        ValueError: If *folder_path* is not a directory or *file_type* is
            unsupported.
    """
    import logging

    logger = logging.getLogger(__name__)

    if pl is None:
        raise ImportError(
            "polars is required for concat_files. Install with: pip install polars"
        )

    folder_path = Path(folder_path).resolve()

    if not folder_path.exists():
        raise FileNotFoundError(f"Folder path does not exist: {folder_path}")
    if not folder_path.is_dir():
        raise ValueError(f"Path is not a directory: {folder_path}")

    file_type = file_type.lower()
    supported_types = {"csv", "parquet"}
    if file_type not in supported_types:
        raise ValueError(
            f"Unsupported file type: '{file_type}'. Supported: {supported_types}"
        )

    # Resolve output path
    if output_file is not None:
        out_path = Path(output_file).resolve()
    else:
        out_path = folder_path.parent / f"{folder_path.name}_combined.{file_type}"

    # Discover files
    all_files = sorted(folder_path.rglob(f"*.{file_type}"))

    if not all_files:
        raise FileNotFoundError(
            f"No .{file_type} files found in {folder_path}"
        )

    logger.info("Found %d .%s file(s) to combine.", len(all_files), file_type)

    # Read each file into a lazy/eager frame
    frames: List[Any] = []
    failed: List[str] = []

    for fpath in all_files:
        try:
            if file_type == "csv":
                df = pl.scan_csv(fpath)
            else:
                df = pl.scan_parquet(fpath)

            if add_source_column:
                df = df.with_columns(pl.lit(fpath.name).alias("SOURCE"))

            frames.append(df)
        except Exception as e:
            logger.error("Failed to read %s: %s", fpath, e)
            failed.append(str(fpath))

    if not frames:
        raise RuntimeError(
            f"All {len(all_files)} files failed to read. "
            f"Check logs for details."
        )

    if failed:
        logger.warning(
            "%d of %d files failed to read and were skipped.",
            len(failed),
            len(all_files),
        )

    # Concatenate and write
    combined_lf = pl.concat(frames, how=how)

    out_path.parent.mkdir(parents=True, exist_ok=True)

    if file_type == "csv":
        combined_lf.collect().write_csv(out_path)
    else:
        if hasattr(combined_lf, "sink_parquet"):
            combined_lf.sink_parquet(out_path)
        else:
            combined_lf.collect().write_parquet(out_path)

    size_mb = out_path.stat().st_size / (1024 * 1024)
    logger.info(
        "Combined %d files → %s (%.2f MB).",
        len(frames),
        out_path,
        size_mb,
    )
    return str(out_path)


def profile(
    df: Any,
    title: str = "Data Profiling Report",
    output_path: Optional[Union[str, Path]] = None,
    minimal: bool = True,
) -> Any:
    """Generate a data profiling report for a DataFrame.

    Creates an HTML report using ``ydata-profiling`` and saves it to disk.
    Accepts both Pandas and Polars DataFrames (Polars is converted to
    Pandas internally, since ydata-profiling requires Pandas).

    Args:
        df: A Pandas or Polars DataFrame.
        title: Title of the report (also used to derive the filename).
        output_path: Directory where the HTML report will be saved.
            If ``None``, saves to the current working directory.
        minimal: If ``True`` (default), generate a faster, lighter report.
            Set to ``False`` for a comprehensive report with correlations.

    Returns:
        The ``ProfileReport`` object, which can be used for further
        inspection or exported to other formats.

    Raises:
        ImportError: If ``ydata-profiling`` or ``pandas`` is not installed.
        TypeError: If *df* is not a recognised DataFrame type.
        ValueError: If the DataFrame is empty.
    """
    import logging
    import re

    logger = logging.getLogger(__name__)

    if pd is None:
        raise ImportError(
            "pandas is required for profile. Install with: pip install pandas"
        )

    # Validate input type and convert to Pandas before checking optional deps
    if isinstance(df, pd.DataFrame):
        pandas_df = df
    elif pl is not None and isinstance(df, (pl.DataFrame, pl.LazyFrame)):
        pandas_df = df.collect().to_pandas() if isinstance(df, pl.LazyFrame) else df.to_pandas()
    else:
        raise TypeError(
            f"Expected a Pandas or Polars DataFrame, got {type(df).__name__}"
        )

    if pandas_df.empty:
        raise ValueError("Cannot profile an empty DataFrame.")

    try:
        from ydata_profiling import ProfileReport
    except ImportError:
        raise ImportError(
            "ydata-profiling is required for profile. "
            "Install with: pip install ydata-profiling"
        )

    # Derive a safe filename from the title
    safe_title = re.sub(r"[^\w\s-]", "", title).strip().replace(" ", "_")
    if not safe_title:
        safe_title = "profiling_report"
    filename = f"{safe_title}.html"

    if output_path is not None:
        output_dir = Path(output_path)
        output_dir.mkdir(parents=True, exist_ok=True)
        full_path = output_dir / filename
    else:
        full_path = Path(filename)

    # Suppress irrelevant warnings only during report generation
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", message=".*IProgress not found.*")

        # Disable promotional banners
        prev_banner = os.environ.get("YDATA_PROFILING_DISABLE_PREMIUM_BANNER")
        os.environ["YDATA_PROFILING_DISABLE_PREMIUM_BANNER"] = "1"

        try:
            report = ProfileReport(pandas_df, title=title, minimal=minimal)
            report.to_file(full_path)
        finally:
            # Restore previous environment state
            if prev_banner is None:
                os.environ.pop("YDATA_PROFILING_DISABLE_PREMIUM_BANNER", None)
            else:
                os.environ["YDATA_PROFILING_DISABLE_PREMIUM_BANNER"] = prev_banner

    logger.info("Profile report saved to %s", full_path)
    return report



def csv2parquet(
    input_path: Union[str, Path],
    output_dir: Optional[Union[str, Path]] = None,
    sub_folders: bool = False,
    schema_overrides: Optional[dict] = None,
    compression: str = "snappy",
    overwrite: bool = False,
) -> List[str]:
    """Convert CSV file(s) to Parquet format using Polars LazyFrame.

    Uses ``sink_parquet`` for streaming conversion when possible, which
    processes data in batches without loading the entire file into memory.
    Falls back to ``collect().write_parquet()`` for older Polars versions.

    Args:
        input_path: Path to a CSV file or directory containing CSV files.
        output_dir: Directory to save the Parquet files. If ``None``, saves
            alongside the original CSV files.
        sub_folders: If ``True`` and *input_path* is a directory, recurse
            into subdirectories.
        schema_overrides: Column name → Polars dtype overrides for CSV parsing.
        compression: Parquet compression algorithm. Options: ``'snappy'``,
            ``'zstd'``, ``'lz4'``, ``'gzip'``, ``'uncompressed'``.
        overwrite: If ``False`` (default), skip files that already have a
            corresponding ``.parquet`` file.

    Returns:
        List of paths to the created Parquet files.

    Raises:
        ImportError: If Polars is not installed.
        FileNotFoundError: If *input_path* does not exist.
        ValueError: If *input_path* is neither a file nor a directory.
    """
    import logging

    logger = logging.getLogger(__name__)

    if pl is None:
        raise ImportError(
            "polars is required for csv2parquet. Install with: pip install polars"
        )

    input_path = Path(input_path).resolve()

    if not input_path.exists():
        raise FileNotFoundError(f"Input path does not exist: {input_path}")

    # Resolve output directory
    if output_dir:
        output_base = Path(output_dir).resolve()
        output_base.mkdir(parents=True, exist_ok=True)
    else:
        output_base = None

    # Collect CSV files to process
    if input_path.is_file():
        if input_path.suffix.lower() != ".csv":
            logger.warning("Not a CSV file, skipping: %s", input_path)
            return []
        csv_files = [input_path]
        base_dir = input_path.parent
    elif input_path.is_dir():
        pattern = "*.csv"
        csv_files = sorted(
            input_path.rglob(pattern) if sub_folders else input_path.glob(pattern)
        )
        base_dir = input_path
        if not csv_files:
            logger.info("No CSV files found in %s", input_path)
            return []
        logger.info("Found %d CSV file(s) to convert.", len(csv_files))
    else:
        raise ValueError(f"Input path is neither a file nor a directory: {input_path}")

    created_files: List[str] = []
    skipped = 0
    failed: List[str] = []

    for csv_file in csv_files:
        # Determine output path
        if output_base:
            rel = csv_file.relative_to(base_dir)
            out_file = output_base / rel.with_suffix(".parquet")
            out_file.parent.mkdir(parents=True, exist_ok=True)
        else:
            out_file = csv_file.with_suffix(".parquet")

        # Skip if already exists and overwrite is off
        if out_file.exists() and not overwrite:
            logger.info("Skipping (already exists): %s", out_file)
            skipped += 1
            continue

        try:
            lf = pl.scan_csv(csv_file, schema_overrides=schema_overrides)

            # Prefer sink_parquet (streaming, constant memory) when available
            if hasattr(lf, "sink_parquet"):
                lf.sink_parquet(out_file, compression=compression)
            else:
                lf.collect().write_parquet(out_file, compression=compression)

            size_mb = out_file.stat().st_size / (1024 * 1024)
            logger.info(
                "Converted %s → %s (%.2f MB)", csv_file.name, out_file, size_mb
            )
            created_files.append(str(out_file))

        except Exception as e:
            logger.error("Failed to convert %s: %s", csv_file, e)
            failed.append(str(csv_file))

    # Summary
    parts = []
    if created_files:
        parts.append(f"{len(created_files)} converted")
    if skipped:
        parts.append(f"{skipped} skipped")
    if failed:
        parts.append(f"{len(failed)} failed")
    if parts:
        logger.info("csv2parquet summary: %s.", ", ".join(parts))
    else:
        logger.info("No files were processed.")

    return created_files
