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.
"""

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 file extensions for JSON extraction
SUPPORTED_EXTENSIONS = {'.csv', '.xls', '.xlsx'}

# Resource limits to prevent abuse
from app.core.config import settings as _app_settings
MAX_FILE_SIZE_BYTES = _app_settings.MAX_FILE_SIZE_BYTES
MAX_CSV_ROWS = 100000
MAX_EXCEL_ROWS = 50000
MAX_MEMORY_ROWS = 100000


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]:
    approx_mb = (rows * cols * 50) / (1024 * 1024)
    if rows * cols > MAX_MEMORY_ROWS * 20:
        return f"Data size too large (approx {approx_mb:.1f} MB). Too many cells: {rows}x{cols}"
    return None


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 we have raw data
    if file_data is not None:
        size_error = _validate_file_size(len(file_data))
        if size_error:
            return {"error": size_error, "file_type": ext}
    elif Path(file_path).exists():
        size_error = _validate_file_size(Path(file_path).stat().st_size)
        if size_error:
            return {"error": size_error, "file_type": ext}

    try:
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", UserWarning)

            if ext == '.csv':
                if file_data:
                    stream = io.BytesIO(file_data)
                    # Peek for encoding detection
                    sample = stream.read(1024)
                    stream.seek(0)
                    df = pd.read_csv(stream, nrows=MAX_CSV_ROWS + 1, low_memory=False)
                else:
                    df = pd.read_csv(file_path, nrows=MAX_CSV_ROWS + 1, low_memory=False)
            else:
                if file_data:
                    df = pd.read_excel(io.BytesIO(file_data), engine='openpyxl' if ext == '.xlsx' else 'xlrd')
                else:
                    df = pd.read_excel(file_path, engine='openpyxl' if ext == '.xlsx' else 'xlrd')

        # Enforce row limits to prevent memory exhaustion
        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)
            }

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

        # Efficient conversion to JSON-serializable format
        result = {
            "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()}
            }
        }

        logger.info("Extracted JSON from %s: %d rows, %d cols", ext, len(df), len(df.columns))
        return result

    except pd.errors.EmptyDataError:
        return {
            "error": f"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: {str(exc)}",
            "file_type": ext,
            "exception_type": type(exc).__name__
        }





def is_supported_file_type(file_path: Union[str, Path]) -> bool:
    """
    Check if file type is supported for JSON extraction.
    
    Parameters
    ----------
    file_path : Union[str, Path]
        Path to the file or filename with extension
        
    Returns
    -------
    bool
        True if supported, False otherwise
    """
    extension = Path(file_path).suffix.lower()
    return extension in SUPPORTED_EXTENSIONS


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