import shutil
import numpy as np
import os

from pathlib import Path

from tqdm import tqdm

from semistaticsim.datawrangling.sssd import load_sssd, GeneratedSemiStaticData


def purge_dir(res_path, prefix):
    for item in os.listdir(res_path):
        item_path = os.path.join(res_path, item)
        if not os.path.isdir(item_path) or not item.startswith(prefix):
            if os.path.isdir(item_path):
                shutil.rmtree(item_path)
            else:
                os.remove(item_path)


def move_indexed_files(src_dir: Path, dst_dir: Path, indices, ext):
    dst_dir.mkdir(parents=True, exist_ok=True)
    for trg_idx, src_idx in enumerate(indices):
        shutil.move(src_dir / f"{src_idx:05d}.{ext}", dst_dir / f"{trg_idx:05d}.{ext}")


def save_run(mask: np.ndarray, sssd_data: GeneratedSemiStaticData, res_path: str, tmpdir: str, batch_size: int):
    # Delete existing results
    if os.path.exists(res_path):
        shutil.rmtree(res_path)
    os.makedirs(res_path, exist_ok=True)
    # Index df samples
    indices = np.nonzero(mask)[0]
    subset = sssd_data.take_by_global_timestep(stop_or_start=None, indices=indices)
    # --- Move Images ---
    images_src = tmpdir / "images"
    images_dst = res_path / "images"
    for subdir in images_src.iterdir():
        if subdir.is_dir():
            # detect extension from first file
            sample = next(subdir.iterdir())
            ext = sample.suffix.lstrip(".")

            move_indexed_files(src_dir=subdir, dst_dir=images_dst / subdir.name, indices=indices, ext=ext)

    # --- Move poses ---
    move_indexed_files(src_dir=tmpdir / "poses", dst_dir=res_path / "poses", indices=indices, ext="json")

    # Save parquet
    subset.dump_to_parquet(res_path, dump_leftover=True, verbose=False, batch_size=batch_size)

    # Copy id_to_color mapping if exists
    id_to_color_src = tmpdir / "id_to_color.json"
    if id_to_color_src.exists():
        shutil.copy(id_to_color_src, res_path / "id_to_color.json")


def format_results(res_path: str, tmpdir: str, batch_size: int) -> dict:
    sssd_data = load_sssd(tmpdir)
    # Get timestamps to split data folder
    timestamps = np.concatenate([data._timestamp for data in sssd_data.get_generator_of_selves()])
    floored_timestamps = np.floor(timestamps)
    unique_timestamps = np.unique(floored_timestamps)
    # Split data into multiple folders based on timestamps
    for t in tqdm(unique_timestamps, desc="Formatting results"):
        mask = floored_timestamps == t
        if mask.sum() > 1:
            print(f"Processing run_id: {int(t.item())} with {mask.sum()} samples")
            save_run(mask, sssd_data, res_path=Path(res_path) / f"run_{int(t.item())}", tmpdir=Path(tmpdir), batch_size=batch_size)