"""File, HDF5, JSONL I/O operations."""

import json
import os
import logging
from pathlib import Path
from typing import Dict, List, Optional, Callable, Union, Any

import h5py
import torch
from torch import Tensor
from tqdm import tqdm
from datasets import Dataset

from pipeline.utils import error_handler
from pipeline.processors import BaseProcessor
from pipeline.packing import SequencePacker

logger = logging.getLogger(__name__)


class IOHandler:
    """File and HDF5 read/write operations."""

    @staticmethod
    def fetch_files(directory: str, suffix: Optional[str] = None) -> List[str]:
        files = [
            os.path.join(root, f)
            for root, _, files in os.walk(directory)
            for f in files
        ]
        if suffix:
            files = [f for f in files if f.endswith(suffix)]
        return sorted(files)

    @staticmethod
    def fetch_folders(
        root_dir: str, filter_func: Optional[Callable[[str], bool]] = None
    ) -> List[str]:
        folders = []
        for root, dirs, _ in os.walk(root_dir):
            for dir_name in dirs:
                folder_path = os.path.join(root, dir_name)
                if filter_func is None or filter_func(folder_path):
                    folders.append(folder_path)
        return folders

    @staticmethod
    @error_handler()
    def save_h5(
        output_dir: str, file_name: str, tensor_group: Dict[str, List[Tensor]]
    ) -> None:
        os.makedirs(output_dir, exist_ok=True)
        full_path = os.path.join(output_dir, f"{file_name}.h5")

        with h5py.File(full_path, "w") as f:
            for key, tensors in tensor_group.items():
                grp = f.create_group(key)
                for idx, tensor in enumerate(tensors):
                    grp.create_dataset(f"data_{idx}", data=tensor.cpu().numpy())

    @staticmethod
    @error_handler()
    def load_h5(file_path: str, share_memory: bool = True) -> Dict[str, List[Tensor]]:
        tensor_group: Dict[str, List[Tensor]] = {}

        root_path = Path(file_path)
        h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))

        for h5_file in h5_files:
            with h5py.File(h5_file, "r") as f:
                for key in f.keys():
                    grp = f[key]
                    dsets = []
                    for dset_name in grp.keys():
                        dset = grp[dset_name]
                        tensor = torch.from_numpy(dset[:])
                        if share_memory:
                            tensor = tensor.share_memory_()
                        dsets.append(tensor)

                    if tensor_group.get(key) is None:
                        tensor_group[key] = []
                    tensor_group[key].extend(dsets)

        return tensor_group


# ── Stage 1: Export HuggingFace Dataset to JSONL ──────────────────────────


@error_handler()
def export_dataset(
    dataset: Dataset,
    output_dir: str,
    output_prefix: str,
    *,
    chunk_size: int = 1_000_000,
    max_chunks: Optional[int] = None,
    process_func: Optional[
        Callable[[Dict[str, Any]], Union[Dict[str, Any], List[Dict[str, Any]]]]
    ] = None,
    column: str = "text",
) -> List[str]:
    """
    Export HuggingFace Dataset to JSONL files in chunks.

    Args:
        dataset: HuggingFace Dataset object
        output_dir: Output directory
        output_prefix: Output file name prefix, e.g., "chinese-c4-pretrain"
        chunk_size: Maximum number of samples per file
        max_chunks: Maximum number of chunks to process (for debugging)
        process_func: Single sample transformation function (dict) -> dict | list[dict]
        column: Default text column name (only used when process_func is None)

    Returns:
        List of generated file paths
    """
    os.makedirs(output_dir, exist_ok=True)
    total = len(dataset)
    num_chunks = (total + chunk_size - 1) // chunk_size
    lim = min(max_chunks, num_chunks) if max_chunks else num_chunks

    output_files: List[str] = []
    for i in range(lim):
        start = i * chunk_size
        end = min(start + chunk_size, total)
        chunk = dataset.select(range(start, end))

        path = os.path.join(output_dir, f"{output_prefix}_chunk_{i}.jsonl")
        try:
            with open(path, "w", encoding="utf-8") as f:
                for example in chunk:
                    processed = (
                        process_func(example)
                        if process_func
                        else {column: example[column]}
                    )
                    items = processed if isinstance(processed, list) else [processed]
                    for item in items:
                        f.write(json.dumps(item, ensure_ascii=False) + "\n")
            output_files.append(path)
            logger.info(f"[{i + 1}/{lim}] Saved {path}")
        except (OSError, IOError) as e:
            logger.error(f"Failed to write chunk {i} to {path}: {e}")

    return output_files


# ── Stage 2: Tokenize JSONL and cache to HDF5 ────────────────────────────


@error_handler()
def cache_jsonl(
    files: List[str],
    output_dir: str,
    processor: BaseProcessor,
    *,
    pack_size: int = -1,
    pad_value: int = 1,
) -> List[str]:
    """
    Tokenize JSONL files and pack them into HDF5 storage.

    Args:
        files: List of JSONL file paths
        output_dir: H5 output directory
        processor: Initialized Processor instance
        pack_size: Packing length, <=0 means no packing
        pad_value: Padding value

    Returns:
        List of generated H5 file paths
    """
    os.makedirs(output_dir, exist_ok=True)
    output_files: List[str] = []
    output_keys = processor.output_keys

    for file_path in files:
        file_name = Path(file_path).stem

        arrows: Dict[str, List] = {key: [] for key in output_keys}

        with open(file_path, "r", encoding="utf-8") as f:
            for line_num, line in enumerate(
                tqdm(f, desc=f"Processing {file_name}", leave=False), start=1
            ):
                try:
                    result = processor.process(json.loads(line))
                    if result is not None:
                        for key in output_keys:
                            arrows[key].append(result[key])
                except json.JSONDecodeError as e:
                    logger.warning(
                        f"JSON decode error in {file_path} line {line_num}: {e}. Skipping line."
                    )
                    continue
                except Exception as e:
                    logger.warning(
                        f"Unexpected error processing line {line_num} in {file_path}: {e}. Skipping line."
                    )
                    continue

        if pack_size > 0:
            output = {}
            for key in output_keys:
                packer = SequencePacker(pack_size, pad_value)
                output[key] = packer.pack(arrows[key])
        else:
            output = arrows

        IOHandler.save_h5(output_dir, file_name, output)
        h5_path = os.path.join(output_dir, f"{file_name}.h5")
        output_files.append(h5_path)
        logger.info(f"Saved {h5_path}")

    return output_files