import math
import os
import sys
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

# Add the directory containing this file to sys.path to allow importing longcat_image
# which is located in the same directory
current_file_path = Path(__file__).resolve()
node_dir = current_file_path.parent
if str(node_dir) not in sys.path:
    sys.path.insert(0, str(node_dir))

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoProcessor

# ComfyUI imports
import folder_paths
import comfy.sd
import comfy.utils
import comfy.model_management
import comfy.samplers
import comfy.model_patcher
import comfy.clip_vision

from longcat_image.models import LongCatImageTransformer2DModel
from longcat_image.pipelines import LongCatImageEditPipeline, LongCatImagePipeline
from longcat_image.pipelines.pipeline_longcat_image_edit import calculate_dimensions
from longcat_image.utils.model_utils import prepare_pos_ids, calculate_shift, retrieve_timesteps


# Helpers for device/dtype
def _select_device() -> torch.device:
    return comfy.model_management.get_torch_device()

def _select_dtype(device: torch.device) -> torch.dtype:
    # ComfyUI usually handles this, but we can default to fp16 for inference
    return torch.float16

def _torch_image_to_pil_list(image: torch.Tensor) -> List[Image.Image]:
    if image.dim() == 3:
        image = image.unsqueeze(0)
    image = image.clamp(0, 1).detach().cpu()
    pil_images = []
    for img in image:
        np_img = (img.numpy() * 255.0).round().astype(np.uint8)
        pil_images.append(Image.fromarray(np_img))
    return pil_images

def _pil_list_to_torch(images: List[Image.Image]) -> torch.Tensor:
    tensors = []
    for img in images:
        np_img = np.array(img, copy=False)
        if np_img.ndim == 2:
            np_img = np.repeat(np_img[..., None], 3, axis=2)
        tensor = torch.from_numpy(np_img.astype(np.float32) / 255.0)
        tensors.append(tensor)
    if not tensors:
        return torch.zeros((1, 512, 512, 3), dtype=torch.float32)
    return torch.stack(tensors, dim=0)

def _round_to_multiple(value: int, multiple: int = 16) -> int:
    return int(math.ceil(value / multiple) * multiple)


# --- Loaders & Wrappers ---

class LongCatCheckpointLoader:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "ckpt_name": (folder_paths.get_filename_list("checkpoints"),),
            }
        }

    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "load_checkpoint"
    CATEGORY = "longcat"

    def load_checkpoint(self, ckpt_name: str):
        ckpt_path = folder_paths.get_full_path("checkpoints", ckpt_name)
        
        # Load state dict
        sd = comfy.utils.load_torch_file(ckpt_path)
        
        # Load Model (Transformer)
        # We assume the checkpoint contains the transformer weights or is a diffusers dump
        # If it's a single file, we need to know the prefix or structure.
        # For now, let's assume it's a standard Diffusers-style checkpoint loaded as a dict.
        # If it's a folder, comfy.utils.load_torch_file might fail or return dict of files?
        # folder_paths usually returns file paths.
        
        # Initialize Transformer
        # We need config. If not in checkpoint, we might need default config.
        # Assuming standard LongCat config for now.
        transformer = LongCatImageTransformer2DModel(
            patch_size=2,
            in_channels=16,
            num_layers=28,
            num_single_layers=10, # Example default, adjust as needed
            attention_head_dim=72,
            num_attention_heads=16,
            joint_attention_dim=768, # Example default
            pooled_projection_dim=768, # Example default
            axes_dims_rope=[16, 56, 56],
        )
        
        # Load weights into transformer
        # We might need to filter keys if the checkpoint contains CLIP/VAE too.
        # This is a simplification. Real implementation needs robust key mapping.
        # transformer.load_state_dict(sd, strict=False) 
        
        # Wrap in ModelPatcher
        model = comfy.model_patcher.ModelPatcher(transformer, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
        
        # Load CLIP
        # Using Comfy's CLIP loader logic if possible, or creating one.
        # clip = comfy.sd.load_clip(ckpt_path) # This might work if it's a supported format
        
        # Load VAE
        # vae = comfy.sd.load_vae(ckpt_path)
        
        # Placeholder for actual loading logic which depends on file format
        # For this refactor, we assume we return standard objects.
        
        return (model, None, None) # Placeholder


# --- Encoders ---

class TextEncodeLongCatImage:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip": ("CLIP",),
                "prompt": ("STRING", {"multiline": True}),
                "enable_prompt_rewrite": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"
    CATEGORY = "longcat"

    def encode(self, clip: Any, prompt: str, enable_prompt_rewrite: bool):
        # clip is now a comfy.sd.CLIP object
        # We can use standard encoding or custom logic
        
        tokens = clip.tokenize(prompt)
        cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
        
        # Pack into CONDITIONING
        return ([[cond, {"pooled_output": pooled}]],)


class TextEncodeLongCatImageEdit:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "clip": ("CLIP",),
                "vae": ("VAE",),
                "image1": ("IMAGE",),
                "prompt": ("STRING", {"multiline": True}),
            },
            "optional": {
                "image2": ("IMAGE",),
                "image3": ("IMAGE",),
            }
        }

    RETURN_TYPES = ("CONDITIONING",)
    FUNCTION = "encode"
    CATEGORY = "longcat"

    def encode(self, clip: Any, vae: Any, image1: torch.Tensor, prompt: str, image2: Optional[torch.Tensor] = None, image3: Optional[torch.Tensor] = None):
        # Custom logic for edit encoding
        # We need to access the internal tokenizer/encoder of the CLIP object if possible
        # or use standard encoding and add image tokens?
        
        # Placeholder for complex logic
        return ([],)


# --- VAE ---

class VAEEncodeLongCat:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vae": ("VAE",),
                "image": ("IMAGE",),
            }
        }
    
    RETURN_TYPES = ("LATENT",)
    FUNCTION = "encode"
    CATEGORY = "longcat"
    
    def encode(self, vae: Any, image: torch.Tensor):
        return (vae.encode(image[:,:,:,:3]),)


class VAEDecodeLongCat:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "vae": ("VAE",),
                "samples": ("LATENT",),
            }
        }
    
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "decode"
    CATEGORY = "longcat"
    
    def decode(self, vae: Any, samples: Dict):
        return (vae.decode(samples["samples"]),)


# --- Sampler ---

class LongCatSampler:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("MODEL",),
                "positive": ("CONDITIONING",),
                "negative": ("CONDITIONING",),
                "latent_image": ("LATENT",),
                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
                "steps": ("INT", {"default": 30, "min": 1, "max": 1000}),
                "cfg": ("FLOAT", {"default": 4.5, "min": 0.0, "max": 100.0}),
                "sampler_name": (comfy.samplers.KSampler.SAMPLERS,),
                "scheduler": (comfy.samplers.KSampler.SCHEDULERS,),
                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
            }
        }

    RETURN_TYPES = ("LATENT",)
    FUNCTION = "sample"
    CATEGORY = "longcat"

    def sample(self, model: Any, positive: List, negative: List, latent_image: Dict, seed: int, steps: int, cfg: float, sampler_name: str, scheduler: str, denoise: float):
        # Use ComfyUI's common_ksampler logic
        # This handles the sampling loop using the provided sampler/scheduler
        
        return comfy.samplers.common_ksampler(
            model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise
        )


class LongCatImageSizeScale:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "target_area": ("INT", {"default": 1024 * 1024, "min": 256 * 256, "max": 4096 * 4096}),
                "upscale_method": (["nearest-exact", "bilinear", "area", "bicubic", "lanczos"],),
            }
        }

    RETURN_TYPES = ("IMAGE", "INT", "INT")
    RETURN_NAMES = ("image", "width", "height")
    FUNCTION = "scale"
    CATEGORY = "longcat"

    def scale(self, image: torch.Tensor, target_area: int, upscale_method: str):
        if image.dim() == 3:
            image = image.unsqueeze(0)
        b, h, w, c = image.shape
        ratio = w / float(h)
        width = math.sqrt(target_area * ratio)
        height = width / ratio
        width = _round_to_multiple(int(round(width)), 16)
        height = _round_to_multiple(int(round(height)), 16)
        
        # Use ComfyUI's common_upscale
        scaled = comfy.utils.common_upscale(image.permute(0, 3, 1, 2), width, height, upscale_method, "center")
        scaled = scaled.permute(0, 2, 3, 1)
        
        return (scaled, int(width), int(height))