#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Public API for HDF5 to Zarr migration.

Internal helpers in _h5_helpers.py, error compat in _compat.py.
"""

import os
import warnings
from pathlib import Path
from typing import Any, List, Optional, Tuple, Union

import h5py
import zarr
from tqdm import tqdm

from ._compat import (
    FileFormatError,
    PathNotFoundError,
    SciTeXIOError,
    check_file_exists,
    check_path,
)
from ._h5_helpers import (
    copy_h5_attributes,
    get_zarr_compressor,
    migrate_group,
    validate_migration,
)


def migrate_h5_to_zarr(
    h5_path: Union[str, Path],
    zarr_path: Optional[Union[str, Path]] = None,
    compressor: Optional[Union[str, Any]] = "zstd",
    chunks: Optional[Union[bool, Tuple[int, ...]]] = True,
    overwrite: bool = False,
    show_progress: bool = True,
    validate: bool = True,
) -> str:
    """
    Migrate HDF5 file to Zarr format.

    Parameters
    ----------
    h5_path : str or Path
        Path to input HDF5 file
    zarr_path : str or Path, optional
        Path for output Zarr store. If None, uses h5_path with .zarr extension
    compressor : str or compressor object, optional
        Compression to use: 'zstd', 'lz4', 'gzip', 'blosc', or None
    chunks : bool or tuple, optional
        Chunking strategy. True for auto, False for no chunks, or specific shape
    overwrite : bool, optional
        Whether to overwrite existing Zarr store
    show_progress : bool, optional
        Whether to show migration progress
    validate : bool, optional
        Whether to validate the migration by comparing shapes

    Returns
    -------
    str
        Path to created Zarr store
    """
    h5_path = Path(h5_path)
    if not h5_path.is_absolute():
        check_file_exists(str(h5_path))
    else:
        if not h5_path.exists():
            raise PathNotFoundError(str(h5_path))

    if zarr_path is None:
        zarr_path = h5_path.with_suffix(".zarr")
    else:
        zarr_path = Path(zarr_path)
        if not zarr_path.is_absolute():
            check_path(str(zarr_path))

    if zarr_path.exists() and not overwrite:
        raise SciTeXIOError(
            f"Zarr store already exists: {zarr_path}",
            suggestion="Use overwrite=True to replace existing store",
        )

    compressor_obj = get_zarr_compressor(compressor)

    if show_progress:
        print(f"Migrating HDF5 to Zarr:")
        print(f"  Source: {h5_path}")
        print(f"  Target: {zarr_path}")
        print(f"  Compressor: {compressor}")

    try:
        with h5py.File(str(h5_path), "r") as h5_file:
            if zarr_path.exists() and overwrite:
                import shutil

                shutil.rmtree(zarr_path)

            zarr_store = zarr.open(str(zarr_path), mode="w")
            copy_h5_attributes(h5_file, zarr_store)
            migrate_group(h5_file, zarr_store, compressor_obj, chunks, show_progress)

            if show_progress:
                print("Migration complete!")

            if validate:
                if show_progress:
                    print("Validating migration...")
                validate_migration(h5_file, zarr_store, show_progress)

    except OSError as e:
        if "Unable to open file" in str(e) or "bad symbol table" in str(e):
            warnings.warn(f"HDF5 file appears to be corrupted: {h5_path}")
            raise FileFormatError(
                str(h5_path),
                expected_format="HDF5",
                actual_format="corrupted HDF5",
            )
        else:
            raise SciTeXIOError(
                f"Failed to open HDF5 file: {h5_path}",
                context={"error": str(e)},
            )
    except Exception as e:
        raise SciTeXIOError(
            f"Migration failed: {str(e)}",
            context={"h5_path": str(h5_path), "zarr_path": str(zarr_path)},
            suggestion="Check file permissions and disk space",
        )

    return str(zarr_path)


def migrate_h5_to_zarr_batch(
    h5_paths: List[Union[str, Path]],
    output_dir: Optional[Union[str, Path]] = None,
    compressor: Optional[Union[str, Any]] = "zstd",
    chunks: Optional[Union[bool, Tuple[int, ...]]] = True,
    overwrite: bool = False,
    parallel: bool = False,
    n_workers: Optional[int] = None,
) -> List[str]:
    """
    Migrate multiple HDF5 files to Zarr format.

    Parameters
    ----------
    h5_paths : list of str or Path
        List of HDF5 files to migrate
    output_dir : str or Path, optional
        Directory for output Zarr stores
    compressor : str or compressor object, optional
        Compression to use
    chunks : bool or tuple, optional
        Chunking strategy
    overwrite : bool, optional
        Whether to overwrite existing Zarr stores
    parallel : bool, optional
        Whether to process files in parallel
    n_workers : int, optional
        Number of parallel workers

    Returns
    -------
    list of str
        Paths to created Zarr stores
    """
    h5_paths = [Path(p) for p in h5_paths]

    zarr_paths = []
    for h5_path in h5_paths:
        if output_dir is None:
            zarr_path = h5_path.with_suffix(".zarr")
        else:
            output_dir_path = Path(output_dir)
            output_dir_path.mkdir(parents=True, exist_ok=True)
            zarr_path = output_dir_path / h5_path.with_suffix(".zarr").name
        zarr_paths.append(zarr_path)

    print(f"Migrating {len(h5_paths)} HDF5 files to Zarr format...")

    if parallel and len(h5_paths) > 1:
        migrated_paths = _migrate_parallel(
            h5_paths, zarr_paths, compressor, chunks, overwrite, n_workers
        )
    else:
        migrated_paths = _migrate_sequential(
            h5_paths, zarr_paths, compressor, chunks, overwrite
        )

    print(f"\nSuccessfully migrated {len(migrated_paths)}/{len(h5_paths)} files")
    return migrated_paths


def _migrate_parallel(h5_paths, zarr_paths, compressor, chunks, overwrite, n_workers):
    """Run migrations in parallel."""
    import functools
    from concurrent.futures import ProcessPoolExecutor, as_completed

    if n_workers is None:
        n_workers = min(os.cpu_count() or 4, len(h5_paths))

    print(f"Using {n_workers} parallel workers...")

    migrate_func = functools.partial(
        migrate_h5_to_zarr,
        compressor=compressor,
        chunks=chunks,
        overwrite=overwrite,
        show_progress=False,
        validate=True,
    )

    with ProcessPoolExecutor(max_workers=n_workers) as executor:
        futures = {
            executor.submit(migrate_func, h5_path, zarr_path): i
            for i, (h5_path, zarr_path) in enumerate(zip(h5_paths, zarr_paths))
        }

        results = []
        with tqdm(total=len(h5_paths), desc="Migrating") as pbar:
            for future in as_completed(futures):
                idx = futures[future]
                try:
                    result = future.result()
                    results.append((idx, result))
                except Exception as e:
                    print(f"\nError migrating {h5_paths[idx]}: {e}")
                    results.append((idx, None))
                pbar.update(1)

        results.sort(key=lambda x: x[0])
        return [r[1] for r in results if r[1] is not None]


def _migrate_sequential(h5_paths, zarr_paths, compressor, chunks, overwrite):
    """Run migrations sequentially."""
    migrated_paths = []
    for h5_path, zarr_path in tqdm(
        zip(h5_paths, zarr_paths), total=len(h5_paths), desc="Migrating"
    ):
        try:
            result = migrate_h5_to_zarr(
                h5_path,
                zarr_path,
                compressor=compressor,
                chunks=chunks,
                overwrite=overwrite,
                show_progress=False,
                validate=True,
            )
            migrated_paths.append(result)
        except Exception as e:
            print(f"\nError migrating {h5_path}: {e}")

    return migrated_paths
