app/services/dataset_metadata_service.py

"""
Production-ready asynchronous dataset metadata extraction service module.

Supports CSV, TSV, XLS, and XLSX files from local file objects or remote URLs.
Designed for high-concurrency production environments with memory-efficient processing.
"""

from __future__ import annotations

import asyncio
import csv
import io
import logging
import os
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import (
    Any,
    BinaryIO,
    Dict,
    List,
    Optional,
    Sequence,
    Tuple,
    Union,
)
from urllib.parse import unquote, urlparse

import aiohttp
import chardet
import numpy as np
import pandas as pd
import xlrd  # noqa: F401 – needed as engine for .xls
import openpyxl  # noqa: F401 – needed as engine for .xlsx

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_SAMPLE_BYTES_FOR_DETECTION: int = 65_536          # 64 KiB for encoding/delimiter sniffing
_DEFAULT_SAMPLE_ROWS: int = 5
_MAX_FILE_SIZE: int = 2 * 1024 * 1024 * 1024       # 2 GiB
_DOWNLOAD_CHUNK_SIZE: int = 256 * 1024              # 256 KiB streaming chunks
_DEFAULT_TIMEOUT_SECONDS: int = 120
_THREAD_POOL = ThreadPoolExecutor(
    max_workers=min(32, (os.cpu_count() or 4) + 4),
)

# Magic bytes for binary file-type detection
_XLSX_MAGIC: bytes = b"PK\x03\x04"
_XLS_MAGIC: bytes = b"\xd0\xcf\x11\xe0"


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class FileType(str, Enum):
    """Supported dataset file types."""

    CSV = "csv"
    TSV = "tsv"
    XLS = "xls"
    XLSX = "xlsx"


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------

class MetadataExtractionError(Exception):
    """Base exception for all metadata extraction failures."""


class UnsupportedFileTypeError(MetadataExtractionError):
    """File type cannot be determined or is not CSV/TSV/XLS/XLSX."""


class FileDownloadError(MetadataExtractionError):
    """Remote file cannot be downloaded."""


class FileTooLargeError(MetadataExtractionError):
    """File exceeds the configured maximum size."""


class EncodingDetectionError(MetadataExtractionError):
    """Text encoding cannot be reliably determined."""


class FileParsingError(MetadataExtractionError):
    """File content cannot be parsed into a tabular structure."""


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ExtractionConfig:
    """
    Immutable configuration for metadata extraction.

    Attributes
    ----------
    sample_rows : int
        Number of sample rows to include in the result.
    timeout_seconds : int
        HTTP download timeout for remote files.
    max_file_size_bytes : int
        Maximum allowed file size in bytes.
    detect_encoding : bool
        Whether to auto-detect text encoding for CSV/TSV files.
    fill_missing_sample_value : str
        Placeholder used to replace NaN in ``sample_data``.
    extra_na_values : Sequence[str]
        Additional strings treated as missing values during parsing.
    """

    sample_rows: int = _DEFAULT_SAMPLE_ROWS
    timeout_seconds: int = _DEFAULT_TIMEOUT_SECONDS
    max_file_size_bytes: int = _MAX_FILE_SIZE
    detect_encoding: bool = True
    fill_missing_sample_value: str = "Unknown"
    extra_na_values: Sequence[str] = field(
        default_factory=lambda: ("", "NA", "N/A", "null", "NULL", "None"),
    )

    def __post_init__(self) -> None:
        if self.sample_rows < 1:
            raise ValueError("sample_rows must be >= 1")
        if self.timeout_seconds < 1:
            raise ValueError("timeout_seconds must be >= 1")
        if self.max_file_size_bytes < 1:
            raise ValueError("max_file_size_bytes must be >= 1")


# ---------------------------------------------------------------------------
# Internal helpers — detection
# ---------------------------------------------------------------------------

def _detect_file_type_from_bytes(header: bytes) -> Optional[FileType]:
    """Detect file type by inspecting magic bytes."""
    if len(header) >= 4:
        if header[:4] == _XLSX_MAGIC:
            return FileType.XLSX
        if header[:4] == _XLS_MAGIC:
            return FileType.XLS
    return None


def _detect_file_type_from_extension(filename: str) -> Optional[FileType]:
    """Map a filename extension to a ``FileType``."""
    ext = Path(filename).suffix.lower().lstrip(".")
    mapping: Dict[str, FileType] = {
        "csv": FileType.CSV,
        "tsv": FileType.TSV,
        "xls": FileType.XLS,
        "xlsx": FileType.XLSX,
    }
    return mapping.get(ext)


def _detect_encoding(raw: bytes) -> str:
    """Detect text encoding via *chardet*, defaulting to ``utf-8``."""
    if not raw:
        return "utf-8"
    result = chardet.detect(raw)
    encoding = (result.get("encoding") or "utf-8").lower()
    confidence: float = result.get("confidence", 0.0)
    logger.debug("Encoding detection: %s (confidence=%.2f)", encoding, confidence)
    aliases: Dict[str, str] = {
        "ascii": "utf-8",
        "windows-1252": "latin-1",
        "iso-8859-1": "latin-1",
    }
    return aliases.get(encoding, encoding)


def _detect_delimiter(text_sample: str) -> str:
    """Detect CSV/TSV delimiter using ``csv.Sniffer``; falls back to comma."""
    try:
        dialect = csv.Sniffer().sniff(text_sample, delimiters=",\t;|")
        return dialect.delimiter
    except csv.Error:
        return "\t" if "\t" in text_sample else ","


# ---------------------------------------------------------------------------
# Internal helpers — DataFrame utilities
# ---------------------------------------------------------------------------

def _classify_columns(
    df: pd.DataFrame,
) -> Tuple[List[str], List[str], List[str], List[str]]:
    """Split columns into numeric, categorical, datetime and boolean lists."""
    numeric: List[str] = []
    categorical: List[str] = []
    datetime_cols: List[str] = []
    boolean_cols: List[str] = []

    for col in df.columns:
        dtype = df[col].dtype
        if pd.api.types.is_bool_dtype(dtype):
            boolean_cols.append(str(col))
        elif pd.api.types.is_numeric_dtype(dtype):
            numeric.append(str(col))
        elif pd.api.types.is_datetime64_any_dtype(dtype):
            datetime_cols.append(str(col))
        else:
            categorical.append(str(col))

    return numeric, categorical, datetime_cols, boolean_cols


def _compute_missing(df: pd.DataFrame) -> Dict[str, int]:
    """Return ``{column: count}`` for every column that has missing values."""
    missing = df.isnull().sum()
    return {str(col): int(cnt) for col, cnt in missing.items() if cnt > 0}


def _native_value(v: Any, fill: str) -> Any:
    """Convert numpy scalars to Python-native types for serialisation."""
    if isinstance(v, np.integer):
        return int(v)
    if isinstance(v, np.floating):
        return float(v)
    if isinstance(v, np.bool_):
        return bool(v)
    if isinstance(v, np.ndarray):
        return v.tolist()
    if pd.isna(v):
        return fill
    return v


def _build_response(
    *,
    file_type: FileType,
    encoding: Optional[str],
    delimiter: Optional[str],
    sheet_name: Optional[str],
    df_full: pd.DataFrame,
    df_sample: pd.DataFrame,
    fill_value: str,
    estimated_file_size: int,
) -> Dict[str, Any]:
    """Assemble the standardised metadata dictionary."""
    numeric, categorical, datetime_cols, boolean_cols = _classify_columns(df_full)
    missing = _compute_missing(df_full)

    sample_records: List[Dict[str, Any]] = [
        {k: _native_value(v, fill_value) for k, v in row.items()}
        for row in df_sample.fillna(fill_value).to_dict(orient="records")
    ]

    return {
        "success": True,
        "file_type": file_type.value,
        "encoding": encoding,
        "delimiter": delimiter,
        "sheet_name": sheet_name,
        "shape": {
            "rows": int(df_full.shape[0]),
            "columns": int(df_full.shape[1]),
        },
        "columns": [str(c) for c in df_full.columns],
        "sample_data": sample_records,
        "dtypes": {str(col): str(dtype) for col, dtype in df_full.dtypes.items()},
        "numeric_columns": numeric,
        "categorical_columns": categorical,
        "datetime_columns": datetime_cols,
        "boolean_columns": boolean_cols,
        "missing_values": missing,
        "memory_usage_bytes": int(df_full.memory_usage(deep=True).sum()),
        "estimated_file_size_bytes": estimated_file_size,
    }


# ---------------------------------------------------------------------------
# Synchronous readers (executed inside the thread pool)
# ---------------------------------------------------------------------------

def _read_csv_metadata(data: bytes, config: ExtractionConfig) -> Dict[str, Any]:
    """Parse CSV / TSV bytes and return standardised metadata."""
    sample_raw = data[:_SAMPLE_BYTES_FOR_DETECTION]
    encoding = _detect_encoding(sample_raw) if config.detect_encoding else "utf-8"

    try:
        text_sample = sample_raw.decode(encoding, errors="replace")
    except (UnicodeDecodeError, LookupError) as exc:
        raise EncodingDetectionError(
            f"Cannot decode sample with encoding '{encoding}': {exc}"
        ) from exc

    delimiter = _detect_delimiter(text_sample)
    file_type = FileType.TSV if delimiter == "\t" else FileType.CSV

    logger.info("CSV/TSV reader: encoding=%s, delimiter=%r", encoding, delimiter)

    try:
        df_full = pd.read_csv(
            io.BytesIO(data),
            encoding=encoding,
            delimiter=delimiter,
            low_memory=True,
            na_values=list(config.extra_na_values),
        )
    except Exception as exc:
        raise FileParsingError(f"Failed to parse CSV/TSV: {exc}") from exc

    return _build_response(
        file_type=file_type,
        encoding=encoding,
        delimiter=delimiter,
        sheet_name=None,
        df_full=df_full,
        df_sample=df_full.head(config.sample_rows),
        fill_value=config.fill_missing_sample_value,
        estimated_file_size=len(data),
    )


def _read_xlsx_metadata(data: bytes, config: ExtractionConfig) -> Dict[str, Any]:
    """Parse XLSX bytes and return standardised metadata."""
    try:
        xls = pd.ExcelFile(io.BytesIO(data), engine="openpyxl")
        sheet_name: str = xls.sheet_names[0]
        df_full = xls.parse(sheet_name=sheet_name, na_values=list(config.extra_na_values))
    except Exception as exc:
        raise FileParsingError(f"Failed to parse XLSX: {exc}") from exc

    return _build_response(
        file_type=FileType.XLSX,
        encoding=None,
        delimiter=None,
        sheet_name=sheet_name,
        df_full=df_full,
        df_sample=df_full.head(config.sample_rows),
        fill_value=config.fill_missing_sample_value,
        estimated_file_size=len(data),
    )


def _read_xls_metadata(data: bytes, config: ExtractionConfig) -> Dict[str, Any]:
    """Parse legacy XLS bytes and return standardised metadata."""
    try:
        xls = pd.ExcelFile(io.BytesIO(data), engine="xlrd")
        sheet_name: str = xls.sheet_names[0]
        df_full = xls.parse(sheet_name=sheet_name, na_values=list(config.extra_na_values))
    except Exception as exc:
        raise FileParsingError(f"Failed to parse XLS: {exc}") from exc

    return _build_response(
        file_type=FileType.XLS,
        encoding=None,
        delimiter=None,
        sheet_name=sheet_name,
        df_full=df_full,
        df_sample=df_full.head(config.sample_rows),
        fill_value=config.fill_missing_sample_value,
        estimated_file_size=len(data),
    )


# ---------------------------------------------------------------------------
# Source normalisation (async)
# ---------------------------------------------------------------------------

async def _read_from_file_like(file_obj: BinaryIO) -> Tuple[bytes, Optional[str]]:
    """Read all bytes from a synchronous file-like object off the event loop."""
    loop = asyncio.get_running_loop()
    data: bytes = await loop.run_in_executor(_THREAD_POOL, file_obj.read)
    filename: Optional[str] = getattr(file_obj, "name", None)
    return data, filename


async def _read_from_path(path: Path) -> Tuple[bytes, str]:
    """Read all bytes from a local ``Path`` off the event loop."""
    if not path.exists():
        raise FileNotFoundError(f"File not found: {path}")
    if not path.is_file():
        raise IsADirectoryError(f"Path is not a regular file: {path}")
    loop = asyncio.get_running_loop()
    data: bytes = await loop.run_in_executor(_THREAD_POOL, path.read_bytes)
    return data, path.name


async def _download_from_url(
    url: str,
    config: ExtractionConfig,
) -> Tuple[bytes, Optional[str]]:
    """Stream-download a remote file with size guard and timeout."""
    timeout = aiohttp.ClientTimeout(total=config.timeout_seconds)
    try:
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(url) as resp:
                if resp.status != 200:
                    raise FileDownloadError(
                        f"HTTP {resp.status} when fetching {url}"
                    )

                content_length = resp.content_length
                if content_length and content_length > config.max_file_size_bytes:
                    raise FileTooLargeError(
                        f"Remote file advertises {content_length} bytes, "
                        f"limit is {config.max_file_size_bytes}"
                    )

                chunks: List[bytes] = []
                total = 0
                async for chunk in resp.content.iter_chunked(_DOWNLOAD_CHUNK_SIZE):
                    total += len(chunk)
                    if total > config.max_file_size_bytes:
                        raise FileTooLargeError(
                            f"Download exceeded {config.max_file_size_bytes} bytes"
                        )
                    chunks.append(chunk)

                data = b"".join(chunks)
    except (aiohttp.ClientError, asyncio.TimeoutError) as exc:
        raise FileDownloadError(f"Download failed for {url}: {exc}") from exc

    parsed = urlparse(url)
    filename = unquote(Path(parsed.path).name) if parsed.path else None
    return data, filename


def _is_url(source: Any) -> bool:
    """Return ``True`` if *source* looks like an HTTP/HTTPS URL."""
    return isinstance(source, str) and source.lower().startswith(("http://", "https://"))


# ---------------------------------------------------------------------------
# File-type resolution
# ---------------------------------------------------------------------------

def _resolve_file_type(data: bytes, filename: Optional[str]) -> FileType:
    """
    Determine ``FileType`` using a cascading strategy:

    1. Magic bytes
    2. File extension
    3. Text heuristic (attempt decode → sniff delimiter)
    """
    # 1. Magic bytes
    ft = _detect_file_type_from_bytes(data[:32])
    if ft is not None:
        return ft

    # 2. Extension
    if filename:
        ft = _detect_file_type_from_extension(filename)
        if ft is not None:
            return ft

    # 3. Text heuristic
    sample = data[:_SAMPLE_BYTES_FOR_DETECTION]
    try:
        text = sample.decode("utf-8", errors="strict")
    except UnicodeDecodeError:
        enc = _detect_encoding(sample)
        try:
            text = sample.decode(enc, errors="strict")
        except (UnicodeDecodeError, LookupError):
            raise UnsupportedFileTypeError(
                "Cannot determine file type from content or filename."
            )

    delimiter = _detect_delimiter(text)
    return FileType.TSV if delimiter == "\t" else FileType.CSV


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

# Accepted source types
FileSource = Union[str, Path, bytes, BinaryIO]


async def extract_metadata(
    source: FileSource,
    *,
    config: Optional[ExtractionConfig] = None,
) -> Dict[str, Any]:
    """
    Extract standardised metadata from a dataset file.

    Parameters
    ----------
    source : str | Path | bytes | BinaryIO
        One of the following:

        * **str** — an ``http://`` or ``https://`` URL to a remote file.
        * **pathlib.Path** — path to a local file on disk.
        * **bytes** — raw file content already in memory.
        * **BinaryIO** — any readable binary file-like object
          (including ``io.BytesIO``, ``open(..., 'rb')``, or a FastAPI
          ``UploadFile.file`` / ``SpooledTemporaryFile``).

    config : ExtractionConfig, optional
        Extraction options (timeouts, sample size, size limits, etc.).
        Defaults to ``ExtractionConfig()`` with sensible production values.

    Returns
    -------
    dict[str, Any]
        A dictionary with the following top-level keys:

        ``success``, ``file_type``, ``encoding``, ``delimiter``,
        ``sheet_name``, ``shape``, ``columns``, ``sample_data``,
        ``dtypes``, ``numeric_columns``, ``categorical_columns``,
        ``datetime_columns``, ``boolean_columns``, ``missing_values``,
        ``memory_usage_bytes``, ``estimated_file_size_bytes``.

    Raises
    ------
    TypeError
        If *source* is not one of the accepted types.
    FileNotFoundError
        If a ``Path`` source does not exist.
    UnsupportedFileTypeError
        If the file type is not CSV, TSV, XLS, or XLSX.
    FileDownloadError
        If a remote URL cannot be fetched.
    FileTooLargeError
        If the file exceeds ``config.max_file_size_bytes``.
    EncodingDetectionError
        If encoding detection fails for a text-based file.
    FileParsingError
        If the file content cannot be parsed into a DataFrame.

    Examples
    --------
    >>> import asyncio
    >>> from pathlib import Path
    >>> result = asyncio.run(extract_metadata(Path("titanic.csv")))
    >>> result["shape"]
    {'rows': 891, 'columns': 12}
    """
    if config is None:
        config = ExtractionConfig()

    # ---- Normalise source → (bytes, Optional[filename]) ----
    data: bytes
    filename: Optional[str] = None

    if _is_url(source):
        assert isinstance(source, str)
        logger.info("Downloading remote file: %s", source)
        data, filename = await _download_from_url(source, config)

    elif isinstance(source, Path):
        logger.info("Reading local path: %s", source)
        data, filename = await _read_from_path(source)

    elif isinstance(source, bytes):
        logger.info("Received raw bytes (%d B)", len(source))
        data = source

    elif hasattr(source, "read"):
        logger.info("Reading from file-like object")
        data, filename = await _read_from_file_like(source)  # type: ignore[arg-type]

    else:
        raise TypeError(
            f"Unsupported source type: {type(source).__name__}. "
            "Expected str (URL), Path, bytes, or a binary file-like object."
        )

    # ---- Validate ----
    if not data:
        raise FileParsingError("Source contains no data (0 bytes).")

    if len(data) > config.max_file_size_bytes:
        raise FileTooLargeError(
            f"File size ({len(data)} bytes) exceeds the configured "
            f"limit ({config.max_file_size_bytes} bytes)."
        )

    # ---- Detect file type ----
    file_type = _resolve_file_type(data, filename)
    logger.info("Resolved file type: %s (%d bytes)", file_type.value, len(data))

    # ---- Dispatch to the appropriate synchronous reader ----
    readers = {
        FileType.CSV: _read_csv_metadata,
        FileType.TSV: _read_csv_metadata,
        FileType.XLSX: _read_xlsx_metadata,
        FileType.XLS: _read_xls_metadata,
    }

    reader = readers[file_type]
    loop = asyncio.get_running_loop()
    result: Dict[str, Any] = await loop.run_in_executor(
        _THREAD_POOL,
        reader,
        data,
        config,
    )

    logger.info(
        "Extraction complete — %s, shape=(%d, %d)",
        file_type.value,
        result["shape"]["rows"],
        result["shape"]["columns"],
    )
    return result


# ---------------------------------------------------------------------------
# CLI convenience
# ---------------------------------------------------------------------------

# if __name__ == "__main__":
#     import json
#     import sys

#     async def _main() -> None:
#         if len(sys.argv) < 2:
#             print("Usage: python solution.py <file_path_or_url>")
#             sys.exit(1)

#         arg = sys.argv[1]
#         src: FileSource = arg if _is_url(arg) else Path(arg)  # type: ignore[assignment]

#         result = await extract_metadata(src)
#         print(json.dumps(result, indent=2, default=str))

#     asyncio.run(_main())