soup/soup_cli/data/loader.py

"""Data loading from local files and HuggingFace."""

from __future__ import annotations

import json
from pathlib import Path

from rich.console import Console

from soup_cli.config.schema import DataConfig
from soup_cli.data.formats import (
    detect_format,
    format_to_messages,
    is_audio_format,
    is_vision_format,
)

console = Console()

# File extensions we support
SUPPORTED_EXTENSIONS = {".jsonl", ".json", ".csv", ".parquet", ".txt"}


def load_raw_data(path: Path) -> list[dict]:
    """Load raw data from a file into list of dicts."""
    if not path.exists():
        raise FileNotFoundError(f"Data file not found: {path}")

    ext = path.suffix.lower()
    if ext not in SUPPORTED_EXTENSIONS:
        raise ValueError(f"Unsupported file format: {ext}. Supported: {SUPPORTED_EXTENSIONS}")

    if ext == ".jsonl":
        return _load_jsonl(path)
    elif ext == ".json":
        return _load_json(path)
    elif ext == ".csv":
        return _load_csv(path)
    elif ext == ".parquet":
        return _load_parquet(path)
    elif ext == ".txt":
        return _load_txt(path)

    raise ValueError(f"Unsupported format: {ext}")


def _load_jsonl(path: Path) -> list[dict]:
    data = []
    # v0.40.1 Part E — auto-strip UTF-8 BOM (Windows users overwhelmingly
    # write JSONL via PowerShell `Out-File -Encoding utf8` which adds BOM).
    # The ``utf-8-sig`` codec consumes the BOM transparently if present.
    with open(path, encoding="utf-8-sig") as f:
        for i, line in enumerate(f):
            line = line.strip()
            if not line:
                continue
            try:
                data.append(json.loads(line))
            except json.JSONDecodeError as e:
                console.print(f"[yellow]Warning: invalid JSON on line {i + 1}: {e}[/]")
    return data


def _load_json(path: Path) -> list[dict]:
    with open(path, encoding="utf-8") as f:
        raw = json.load(f)
    if isinstance(raw, list):
        return raw
    raise ValueError("JSON file must contain a list of objects")


def _load_csv(path: Path) -> list[dict]:
    import csv

    with open(path, encoding="utf-8") as f:
        reader = csv.DictReader(f)
        return list(reader)


def _load_parquet(path: Path) -> list[dict]:
    try:
        import pandas as pd
    except ImportError:
        raise ImportError("Install pandas to read parquet files: pip install pandas pyarrow")
    df = pd.read_parquet(path)
    return df.to_dict(orient="records")


def _load_txt(path: Path) -> list[dict]:
    """Load a plain text file as a list of {text: ...} dicts.

    Each non-empty line is treated as a separate document.
    Empty lines are skipped.
    """
    file_size = path.stat().st_size
    if file_size > 500 * 1024 * 1024:  # 500 MB
        console.print(
            f"[yellow]Warning: large text file ({file_size / 1024 / 1024:.0f} MB). "
            f"Consider splitting into smaller files or using JSONL format.[/]"
        )
    with open(path, encoding="utf-8") as f:
        content = f.read()

    # Split by double newline (paragraph/document separator) or treat each line as a doc
    lines = [line.strip() for line in content.split("\n") if line.strip()]
    if not lines:
        console.print(f"[yellow]Warning: empty text file: {path}[/]")
        return []

    return [{"text": line} for line in lines]


def load_dataset(data_config: DataConfig) -> dict:
    """Load dataset for training. Returns dict with 'train' and optionally 'val' keys.

    Supports:
    - Local files (.jsonl, .json, .csv, .parquet, .txt)
    - HuggingFace dataset names (auto-detected if no file extension)
    """
    train_path = data_config.train

    # Check if it's a HuggingFace dataset
    if not Path(train_path).suffix:
        return _load_hf_dataset(train_path, data_config)

    # Local file
    path = Path(train_path)
    raw_data = load_raw_data(path)

    # Detect or use specified format
    fmt = data_config.format
    if fmt == "auto":
        fmt = detect_format(raw_data)
        console.print(f"[dim]Auto-detected format: {fmt}[/]")

    # Convert to standard message format
    formatted = [format_to_messages(row, fmt) for row in raw_data]
    formatted = [r for r in formatted if r is not None]  # filter failed rows

    # Validate image paths for vision formats
    if is_vision_format(fmt):
        image_dir = Path(data_config.image_dir) if data_config.image_dir else path.parent
        formatted = _validate_vision_images(formatted, image_dir)

    # Validate audio paths for audio formats
    if is_audio_format(fmt):
        audio_dir = Path(data_config.audio_dir) if data_config.audio_dir else path.parent
        formatted = _validate_audio_files(formatted, audio_dir)

    # Split into train/val
    if data_config.val_split > 0:
        split_idx = int(len(formatted) * (1 - data_config.val_split))
        return {
            "train": formatted[:split_idx],
            "val": formatted[split_idx:],
        }

    return {"train": formatted}


def _validate_vision_images(data: list[dict], image_dir: Path) -> list[dict]:
    """Validate and resolve image paths in vision dataset rows.

    Each row must have an 'image' key with a filename or path.
    Resolves relative paths against image_dir.
    """
    valid = []
    missing = 0
    for row in data:
        if "image" not in row or not row["image"]:
            missing += 1
            continue
        image_path = Path(row["image"])
        if not image_path.is_absolute():
            image_path = image_dir / image_path
        row["image"] = str(image_path)
        valid.append(row)

    if missing > 0:
        console.print(f"[yellow]Warning: {missing} rows skipped (missing image path)[/]")
    return valid


def _validate_audio_files(data: list[dict], audio_dir: Path) -> list[dict]:
    """Validate and resolve audio file paths in audio dataset rows.

    Each row must have an 'audio' key with a filename or path.
    Resolves relative paths against audio_dir. Rejects path traversal.
    """
    valid = []
    missing = 0
    traversal = 0
    resolved_base = audio_dir.resolve()
    for row in data:
        if "audio" not in row or not row["audio"]:
            missing += 1
            continue
        audio_path = Path(row["audio"])
        if not audio_path.is_absolute():
            audio_path = audio_dir / audio_path
        # Path traversal protection: resolved path must stay under audio_dir
        resolved = audio_path.resolve()
        if not resolved.is_relative_to(resolved_base):
            traversal += 1
            continue
        valid.append({**row, "audio": str(resolved)})

    if missing > 0:
        console.print(f"[yellow]Warning: {missing} rows skipped (missing audio path)[/]")
    if traversal > 0:
        console.print(
            f"[red]Warning: {traversal} rows skipped (audio path traversal blocked)[/]"
        )
    return valid


def _load_hf_dataset(name: str, data_config: DataConfig) -> dict:
    """Load a dataset from HuggingFace Hub."""
    try:
        from datasets import load_dataset as hf_load
    except ImportError:
        raise ImportError("Install datasets: pip install datasets")

    console.print(f"[dim]Loading from HuggingFace: {name}[/]")
    ds = hf_load(name)

    if "train" not in ds:
        raise ValueError(f"Dataset {name} has no 'train' split")

    raw_data = [dict(row) for row in ds["train"]]
    fmt = data_config.format
    if fmt == "auto":
        fmt = detect_format(raw_data)

    formatted = [format_to_messages(row, fmt) for row in raw_data]
    formatted = [r for r in formatted if r is not None]

    if data_config.val_split > 0 and "validation" not in ds:
        split_idx = int(len(formatted) * (1 - data_config.val_split))
        return {"train": formatted[:split_idx], "val": formatted[split_idx:]}

    result = {"train": formatted}
    if "validation" in ds:
        val_data = [dict(row) for row in ds["validation"]]
        val_formatted = [format_to_messages(row, fmt) for row in val_data]
        result["val"] = [r for r in val_formatted if r is not None]

    return result