import pandas as pd
import json
import csv
from pathlib import Path
from typing import Union

from emissary_ml_toolkit.utils.exceptions import DatasetError


def load_csv(file_path: str) -> pd.DataFrame:
    try:
        return pd.read_csv(file_path)
    except Exception:
        with open(file_path) as f:
            sample = f.read(10240)
            sniffer = csv.Sniffer()
            delimiter = sniffer.sniff(sample).delimiter
        return pd.read_csv(file_path, delimiter=delimiter)


def load_json(file_path: str) -> pd.DataFrame:
    data = []
    with open(file_path, 'r') as f:
        try:
            f.seek(0)
            data = json.load(f)
        except json.JSONDecodeError:
            f.seek(0)
            for line in f:
                if line.strip():
                    data.append(json.loads(line))

    return pd.DataFrame(data)


def load_parquet(file_path: str) -> pd.DataFrame:
    return pd.read_parquet(file_path)


def load_file(file_path: Union[str, Path]) -> pd.DataFrame:
    file_path = Path(file_path)
    ext = file_path.suffix.lower()

    loaders = {
        '.csv': load_csv,
        '.json': load_json,
        '.jsonl': load_json,
        '.parquet': load_parquet,
    }

    if ext not in loaders:
        raise DatasetError(f"Unsupported file format: {ext}")

    return loaders[ext](str(file_path))
