import numpy as np
import torch
from pathlib import Path
from huggingface_hub.constants import HF_HUB_CACHE
from PIL import Image
from spandrel import ImageModelDescriptor, ModelLoader

CLEAN_BUCKET = "HyperHail/CB"
CLEAN_FILE = "upscale_models/3_1x-Archivist_Soft.pth"
TILE_SIZE = 256
TILE_PAD = 32
PIXEL_MAX = 255


def download_clean():
    path = Path(HF_HUB_CACHE) / f"buckets--{CLEAN_BUCKET.replace('/', '--')}" / CLEAN_FILE
    if not path.is_file():
        path.parent.mkdir(parents=True, exist_ok=True)
        url = f"https://huggingface.co/buckets/{CLEAN_BUCKET}/resolve/{CLEAN_FILE}"
        torch.hub.download_url_to_file(url, str(path), progress=False)
    return str(path)


def load_clean(device):
    model = ModelLoader().load_from_file(download_clean())
    if not isinstance(model, ImageModelDescriptor) or model.scale != 1:
        raise ValueError("Clean after requires a 1x image model")
    if model.input_channels != 3 or model.output_channels != 3:
        raise ValueError("Clean after requires an RGB model")
    return model.to(device).eval()


@torch.inference_mode()
def clean_image(model, input_path, output_path):
    device = next(model.model.parameters()).device
    with Image.open(input_path) as source:
        image = source.convert("RGB")
        alpha = source.convert("RGBA").getchannel("A") if "A" in source.getbands() or "transparency" in source.info else None
    width, height = image.size
    result = Image.new("RGB", image.size)
    for y in range(0, height, TILE_SIZE):
        for x in range(0, width, TILE_SIZE):
            right, bottom = min(x + TILE_SIZE, width), min(y + TILE_SIZE, height)
            left, top = max(0, x - TILE_PAD), max(0, y - TILE_PAD)
            box = (left, top, min(width, right + TILE_PAD), min(height, bottom + TILE_PAD))
            pixels = np.array(image.crop(box), dtype=np.float32) / PIXEL_MAX
            tile = torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).to(device)
            output = model(tile)[0].clamp(0, 1).mul(PIXEL_MAX).round().byte()
            patch = Image.fromarray(output.permute(1, 2, 0).cpu().numpy())
            result.paste(patch.crop((x - left, y - top, right - left, bottom - top)), (x, y))
    if alpha is not None:
        result.putalpha(alpha)
    result.save(output_path, format="PNG")