extraction/json_extractor.py

"""
JSON extractor for structured data files.

Supports only CSV, XLS, XLSX file types for JSON extraction.
Returns error for unsupported file types.
"""

from __future__ import annotations

import io
import warnings
from pathlib import Path
from typing import Any, Dict, Optional, Union

import pandas as pd

from logger import get_logger

logger = get_logger(__name__)

SUPPORTED_EXTENSIONS: frozenset[str] = frozenset({".csv", ".xls", ".xlsx"})

MAX_CSV_ROWS = 100_000
MAX_EXCEL_ROWS = 50_000
MAX_CELL_COUNT = 2_000_000

try:
    from app.core.config import settings as _app_settings
    MAX_FILE_SIZE_BYTES: int = _app_settings.MAX_FILE_SIZE_BYTES
except Exception:
    MAX_FILE_SIZE_BYTES = 50 * 1024 * 1024


def _validate_file_size(size: int) -> Optional[str]:
    if size > MAX_FILE_SIZE_BYTES:
        return f"File size {size} bytes exceeds limit of {MAX_FILE_SIZE_BYTES} bytes"
    return None


def _check_memory_usage(rows: int, cols: int) -> Optional[str]:
    cell_count = rows * cols
    if cell_count > MAX_CELL_COUNT:
        approx_mb = (cell_count * 50) / (1024 * 1024)
        return f"Data size too large (approx {approx_mb:.1f} MB). Too many cells: {rows}x{cols}"
    return None


def _read_dataframe(ext: str, file_path: Union[str, Path], file_data: Optional[bytes]) -> "pd.DataFrame":
    """Read a CSV/Excel file from either an in-memory buffer or a path."""
    if ext == ".csv":
        if file_data:
            return pd.read_csv(io.BytesIO(file_data), nrows=MAX_CSV_ROWS + 1, low_memory=False)
        return pd.read_csv(file_path, nrows=MAX_CSV_ROWS + 1, low_memory=False)

    engine = "openpyxl" if ext == ".xlsx" else "xlrd"
    if file_data:
        return pd.read_excel(io.BytesIO(file_data), engine=engine)
    return pd.read_excel(file_path, engine=engine)


def extract_json_from_file(
    file_path: Union[str, Path],
    file_data: Optional[bytes] = None,
) -> Dict[str, Any]:
    """
    Extract JSON data from structured files (CSV, XLS, XLSX).

    Parameters
    ----------
    file_path : Union[str, Path]
        Path to the file or filename with extension
    file_data : Optional[bytes]
        Raw file data (for stream processing)

    Returns
    -------
    Dict[str, Any]
        Extracted data or error information
    """
    ext = Path(file_path).suffix.lower()

    if ext not in SUPPORTED_EXTENSIONS:
        return {
            "error": f"Unsupported file type: {ext}. Supported: {', '.join(sorted(SUPPORTED_EXTENSIONS))}",
            "file_type": ext,
        }

    # Validate file size
    if file_data is not None:
        source_bytes = len(file_data)
    else:
        path = Path(file_path)
        source_bytes = path.stat().st_size if path.exists() else 0
    size_error = _validate_file_size(source_bytes)
    if size_error:
        return {"error": size_error, "file_type": ext}

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", UserWarning)
            df = _read_dataframe(ext, file_path, file_data)
    except pd.errors.EmptyDataError:
        return {"error": "File is empty or has no data", "file_type": ext}
    except MemoryError:
        return {"error": "Out of memory processing file", "file_type": ext}
    except Exception as exc:
        logger.exception("JSON extraction failed for %s", ext)
        return {
            "error": f"Processing failed: {exc}",
            "file_type": ext,
            "exception_type": type(exc).__name__,
        }

    max_rows = MAX_EXCEL_ROWS if ext != ".csv" else MAX_CSV_ROWS
    if len(df) > max_rows:
        return {
            "error": f"File contains {len(df)} rows, exceeds limit of {max_rows}",
            "file_type": ext,
            "row_count": len(df),
        }

    mem_error = _check_memory_usage(len(df), len(df.columns))
    if mem_error:
        return {"error": mem_error, "file_type": ext}

    logger.info("Extracted JSON from %s: %d rows, %d cols", ext, len(df), len(df.columns))
    return {
        "success": True,
        "file_type": ext,
        "data": {
            "columns": list(df.columns),
            "rows": df.where(pd.notnull(df), None).to_dict(orient="records"),
            "shape": [len(df), len(df.columns)],
            "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
        },
    }


def is_supported_file_type(file_path: Union[str, Path]) -> bool:
    """Check if file type is supported for JSON extraction."""
    return Path(file_path).suffix.lower() in SUPPORTED_EXTENSIONS


def get_supported_extensions() -> list[str]:
    """Get sorted list of supported file extensions for JSON extraction."""
    return sorted(SUPPORTED_EXTENSIONS)