import os
import torch
import numpy as np
import onnxruntime as ort
from PIL import Image

# 获取当前文件所在目录，用于定位模型
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(CURRENT_DIR, "onnx_models")

class PixelArtConverter:
    def __init__(self):
        self.sessions = {}

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),  # ComfyUI 输入格式 [Batch, H, W, C]
                "resolution": (["128", "256"], {"default": "128"}),
                "temperature": ("FLOAT", {"default": 0.01, "min": 0.0001, "max": 1.0, "step": 0.01}),
                "noise_amp": ("FLOAT", {"default": 0.1, "min": 0.0, "max": 10.0, "step": 0.01}),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
            }
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("image",)
    FUNCTION = "process"
    CATEGORY = "PixelArt"

    def _get_session(self, model_key):
        model_filename = f"pixelart_{model_key}.onnx"
        model_path = os.path.join(MODELS_DIR, model_filename)

        if not os.path.exists(model_path):
            raise FileNotFoundError(f"ONNX Model not found at: {model_path}")

        if model_key not in self.sessions:
            print(f"Loading ONNX model: {model_path}...")
            # 优先尝试 CUDA，如果没有则 CPU
            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
            self.sessions[model_key] = ort.InferenceSession(model_path, providers=providers)
        
        return self.sessions[model_key]

    def process(self, image, resolution, temperature, noise_amp, seed):
        # 1. 准备输入
        # ComfyUI 图片是 Tensor [Batch, H, W, C]，范围 0-1
        # 我们这里处理 Batch 中的每一张图
        
        output_images = []
        
        for i in range(image.shape[0]):
            img_tensor = image[i] # [H, W, C]
            
            # Tensor -> Numpy -> PIL (为了保持和原项目一致的 Resize 逻辑)
            img_np = (img_tensor.cpu().numpy() * 255.0).clip(0, 255).astype(np.uint8)
            img_pil = Image.fromarray(img_np)
            
            # 处理 Alpha (如果是 RGBA 转为白底 RGB)
            if img_pil.mode == 'RGBA':
                bg = Image.new('RGBA', img_pil.size, (255, 255, 255, 255))
                combined = Image.alpha_composite(bg, img_pil)
                img_pil = combined.convert('RGB')
            else:
                img_pil = img_pil.convert('RGB')

            # 2. 预处理 (Pre-processing)
            w, h = img_pil.size
            INPUT_SIZE = 1024
            
            scale = INPUT_SIZE / max(h, w)
            new_w = int(round(w * scale))
            new_h = int(round(h * scale))
            
            # Resize (LANCZOS)
            img_resized = img_pil.resize((new_w, new_h), Image.Resampling.LANCZOS)
            
            # 填充到 1024x1024 白底
            canvas = Image.new("RGB", (INPUT_SIZE, INPUT_SIZE), (255, 255, 255))
            pad_left = (INPUT_SIZE - new_w) // 2
            pad_top = (INPUT_SIZE - new_h) // 2
            canvas.paste(img_resized, (pad_left, pad_top))
            
            # 归一化 (-1 到 1)
            img_input_np = np.array(canvas).astype(np.float32)
            img_input_np = (img_input_np / 255.0) * 2.0 - 1.0
            # [H, W, C] -> [1, C, H, W]
            img_input = img_input_np.transpose(2, 0, 1)[None, :, :, :]

            # 3. 推理 (Inference)
            sess = self._get_session(resolution)
            
            # 随机种子处理
            rng = np.random.default_rng(seed)
            noise_input = rng.standard_normal((1, 32, 256)).astype(np.float32) * noise_amp
            temp_input = np.array([temperature], dtype=np.float32)
            
            input_feed = {
                'input_image': img_input,
                'input_noise': noise_input,
                'temperature': temp_input
            }
            
            outputs = sess.run(None, input_feed)
            pred_img = outputs[0][0] # [3, H, W]
            
            # 4. 后处理 (Post-processing)
            pred_img = pred_img.transpose(1, 2, 0) # [H, W, 3]
            pred_img = np.clip(pred_img, 0.0, 1.0)
            
            # 5. 裁剪逻辑 (Crop back)
            # 计算裁剪区域
            out_h_total, out_w_total = pred_img.shape[:2]
            ratio = out_h_total / float(INPUT_SIZE)
            
            crop_top = int(pad_top * ratio)
            crop_left = int(pad_left * ratio)
            crop_h = int(new_h * ratio)
            crop_w = int(new_w * ratio)
            
            # 确保至少有 1px
            crop_h = max(1, crop_h)
            crop_w = max(1, crop_w)
            
            final_img_np = pred_img[crop_top : crop_top+crop_h, crop_left : crop_left+crop_w]
            
            # 转回 Tensor [H, W, C]
            output_tensor = torch.from_numpy(final_img_np)
            output_images.append(output_tensor)

        # 堆叠 Batch: [Batch, H, W, C]
        # 注意：如果 Batch 内图片裁剪后尺寸不一致，torch.stack 会报错。
        # 通常 ComfyUI 单次处理通常尺寸一致，这里为了稳健，如果尺寸不一致可以在这里做 Padding，
        # 但最简单的做法是假设用户 Batch 输入尺寸长宽比一致。
        if len(output_images) > 1:
            try:
                result = torch.stack(output_images, dim=0)
            except RuntimeError:
                # 如果尺寸不一致，无法堆叠，只能抛出第一张或者报错 (这是ComfyUI Tensor的限制)
                print("Warning: Batch images have different aspect ratios after pixel art conversion. Returning first image only.")
                result = output_images[0].unsqueeze(0)
        else:
            result = output_images[0].unsqueeze(0)

        return (result,)