"""File I/O utilities — data source reading and external object loading."""

from __future__ import annotations

import polars as pl

from haute._logging import get_logger
from haute._lru_cache import LRUCache

logger = get_logger(component="io")

_OBJECT_CACHE_MAX_SIZE = 32


def read_source(path: str) -> pl.LazyFrame:
    """Read a data file into a LazyFrame, dispatching on file extension.

    Centralises the csv/json/parquet dispatch that was previously duplicated
    across the executor, scorer, schema inference, and server modules.

    All formats except ``.json`` use Polars lazy scans, so a downstream
    ``.head(row_limit)`` pushes the limit into the I/O layer and avoids
    reading the full file.  Plain ``.json`` has no ``scan_json`` equivalent
    in Polars, so the entire file is read eagerly then wrapped as lazy.
    This is acceptable because API-input JSON files use the separate
    ``read_json_flat`` path which caches to parquet.

    Raises:
        ValueError: If the file extension is not supported.
    """
    if path.endswith(".csv"):
        return pl.scan_csv(path)
    if path.endswith(".json"):
        # No scan_json in Polars — read eagerly.  Callers should prefer
        # the JSON flatten/cache path (read_json_flat) for large files.
        return pl.read_json(path).lazy()
    if path.endswith(".jsonl"):
        return pl.scan_ndjson(path)
    if path.endswith(".parquet"):
        return pl.scan_parquet(path)
    suffix = path.rsplit(".", 1)[-1] if "." in path else ""
    logger.error("unsupported_file_type", path=path, suffix=suffix)
    raise ValueError(f"Unsupported file type: .{suffix}")


_object_cache: LRUCache[tuple[str, float, str, str], object] = LRUCache(
    max_size=_OBJECT_CACHE_MAX_SIZE,
)


def load_external_object(path: str, file_type: str, model_class: str = "classifier") -> object:
    """Load an external file (model, JSON, pickle, joblib) and return the object.

    Shared by the development executor and the deploy scoring engine.

    Results are cached by ``(path, mtime, file_type, model_class)`` so
    repeated calls (preview clicks, API scoring requests) skip disk I/O.
    The cache auto-invalidates when the file is modified on disk.
    Bounded to ``_OBJECT_CACHE_MAX_SIZE`` entries (LRU eviction).

    All paths are validated to be within the project root before loading.
    Pickle files are deserialized with a restricted unpickler that only
    allows known-safe classes.
    """
    import os

    from haute._sandbox import validate_project_path

    validate_project_path(path)

    try:
        mtime = os.path.getmtime(path)
    except OSError:
        mtime = 0.0
    key = (path, mtime, file_type, model_class)

    cached = _object_cache.get(key)
    if cached is not None:
        return cached

    obj = _load_external_object_uncached(path, file_type, model_class)
    _object_cache.put(key, obj)
    return obj


def _load_external_object_uncached(
    path: str, file_type: str, model_class: str,
) -> object:
    """Deserialize an external file from disk (no caching)."""
    if file_type == "json":
        import json as _json

        with open(path) as f:
            return _json.load(f)
    elif file_type == "joblib":
        from haute._sandbox import safe_joblib_load

        return safe_joblib_load(path)
    elif file_type == "catboost":
        if model_class == "regressor":
            from catboost import CatBoostRegressor

            m = CatBoostRegressor()
        else:
            from catboost import CatBoostClassifier

            m = CatBoostClassifier()
        m.load_model(path)
        return m
    else:  # pickle
        from haute._sandbox import safe_unpickle

        return safe_unpickle(path)
