from typing import Callable, Any
from datasets import IterableDataset
from itertools import tee, islice

def huggingface_login(token: str):
    from huggingface_hub import login, whoami
    login(token=token)
    return whoami()

def print_trainer_summary(result):
    from .gpu import get_gpu_stats
    print(f"Time: {result.metrics['train_runtime']:.2f}")
    print(f"Samples/second: {result.metrics['train_samples_per_second']:.2f}")
    for key, value in get_gpu_stats().items(): print(f"{key}: {value:.4f}")

def dedupe_dataset(dataset, feature_key: str, hash_func: Callable[[Any], bytes] = None):
    seen_hashes = set()

    def default_hash_func(value):
        if hasattr(value, "tobytes"):
            return value.tobytes()  # e.g., NumPy arrays, PIL images
        elif isinstance(value, str):
            return value.encode("utf-8")
        else:
            import json
            return json.dumps(value, sort_keys=True).encode("utf-8")

    _hash_func = hash_func or default_hash_func

    def _filter(row):
        value = row[feature_key]
        hashed = _hash_func(value)
        if hashed in seen_hashes:
            return False
        seen_hashes.add(hashed)
        return True

    return dataset.filter(_filter)


def process_images(images, mode="color", quantize_colors=None, scale=1.0, crop_paddings=None, noise_factor=0.0, return_type="pt"):
    import numpy as np
    from PIL import Image
    from torchvision.transforms.functional import to_tensor

    processed = []
    for image in images:
        # Crop if needed
        if crop_paddings:
            width, height = image.size
            left = crop_paddings[3]
            top = crop_paddings[0]
            right = width - crop_paddings[1]
            bottom = height - crop_paddings[2]
            image = image.crop((left, top, right, bottom))


        # Apply color quantization (after conversion & crop, before tensor)
        if quantize_colors:
            image = image.quantize(colors=quantize_colors)

        # Resize
        if scale != 1.0:
            new_size = (int(image.width * scale), int(image.height * scale))
            image = image.resize(new_size, Image.LANCZOS)

        # Apply noise (directly on PIL image)
        if noise_factor:
            np_image = np.array(image).astype(np.float32) / 255.0
            noise = np.random.normal(0, noise_factor, np_image.shape)
            noisy_np = np.clip(np_image + noise, 0.0, 1.0) * 255.0
            image = Image.fromarray(noisy_np.astype(np.uint8), mode=image.mode)

        # Convert mode
        if mode == "color":
            image = image.convert("RGB")
        elif mode == "grayscale":
            image = image.convert("L")
        elif mode == "black_and_white":
            image = image.convert("1")
        else:
            raise ValueError(f"Unsupported color mode: {mode}")

        if return_type == "pt":
            import torch
            image_ts = to_tensor(image).to(dtype=torch.float16)
            processed.append(image_ts)
        elif return_type == "np":
            image_np = np.array(image).astype(np.float16) / 255.0
            processed.append(image_np)

    return processed

def hf_create_or_get_repo(repo_name: str, private: bool = True) -> str:
    from huggingface_hub import create_repo, HfApi
    from huggingface_hub.utils import RepositoryNotFoundError

    api = HfApi()
    
    try:
        # Check if repo exists
        _ = api.repo_info(repo_name)
        print(f"Repository '{repo_name}' already exists.")
        return f"https://huggingface.co/{repo_name}"
    except RepositoryNotFoundError:
        # Create if not exists
        repo_url = create_repo(repo_name, private=private)
        print(f"Repository '{repo_name}' created.")
        return repo_url

def hf_push_model_to_hub(repo_id: str, model, model_file_name="model.pt", private: bool = True):
    import torch
    from huggingface_hub import upload_file
    
    hf_create_or_get_repo(repo_id, private=private)

    torch.save(model.state_dict(), model_file_name)

    upload_file(
        path_or_fileobj=model_file_name,
        path_in_repo=model_file_name,
        repo_id=repo_id,
        repo_type="model"
    )

class ShiftedDataset(IterableDataset):
    def __init__(self, hf_dataset, shift=1):
        if shift < 1:
            raise ValueError("Shift must be >= 1")
        self.dataset = hf_dataset
        self.shift = shift

    def __iter__(self):
        it1, it2 = tee(self.dataset.__iter__(), 2)
        it2 = islice(it2, self.shift, None)

        for x, y in zip(it1, it2):
            yield {
                "x": x.get("x"),
                "y": y.get("y"),
            }

    def __len__(self):
        try:
            return len(self.dataset) - self.shift
        except TypeError:
            raise TypeError("Length is not defined for streaming datasets.")

    def __getitem__(self, idx):
        if idx < 0 or idx >= len(self):
            raise IndexError("Index out of range")
        x = self.dataset[idx]
        y = self.dataset[idx + self.shift]
        return {"x": x.get("x"), "y": y.get("y")}

