import time
import gradio as gr
import torch
from einops import rearrange, repeat
from PIL import Image
import numpy as np

from flux.sampling import denoise, get_noise, get_schedule, prepare, rf_denoise, rf_inversion, unpack
from flux.util import (
    SamplingOptions,
    load_ae,
    load_clip,
    load_flow_model,
    load_t5,
)
from pulid.pipeline_flux import PuLIDPipeline
from pulid.utils import resize_numpy_image_long, seed_everything

# 간단한 인용 정보 추가
_CITE_ = """PuLID: Person-under-Language Image Diffusion Model"""

# GPU 사용 가능 여부 확인 및 장치 설정
def get_device():
    if torch.cuda.is_available():
        return torch.device('cuda')
    else:
        print("CUDA GPU를 찾을 수 없습니다. CPU를 사용합니다.")
        return torch.device('cpu')

def get_models(name: str, device, offload: bool):
    print(f"모델을 {device}에 로드합니다.")
    t5 = load_t5(device, max_length=128)
    clip_model = load_clip(device)
    model = load_flow_model(name, device="cpu" if offload else device)
    model.eval()
    ae = load_ae(name, device="cpu" if offload else device)
    return model, ae, t5, clip_model


class FluxGenerator:
    def __init__(self):
        # GPU 사용 가능 여부에 따라 장치 설정
        self.device = get_device()
        self.offload = False
        self.model_name = 'flux-dev'
        
        # 모델 로드 시도
        try:
            self.model, self.ae, self.t5, self.clip_model = get_models(
                self.model_name,
                device=self.device,
                offload=self.offload,
            )
            self.pulid_model = PuLIDPipeline(
                self.model, 
                'cuda' if torch.cuda.is_available() else 'cpu', 
                weight_dtype=torch.bfloat16 if self.device.type == 'cuda' else torch.float32
            )
            self.pulid_model.load_pretrain()
            self.initialized = True
        except Exception as e:
            print(f"모델 초기화 중 오류 발생: {e}")
            self.initialized = False

# 모델 초기화 시도
try:
    flux_generator = FluxGenerator()
    model_initialized = flux_generator.initialized
except Exception as e:
    print(f"FluxGenerator 초기화 중 오류 발생: {e}")
    model_initialized = False


@torch.inference_mode()
def generate_image(
    prompt: str,
    id_image = None,
    width: int = 512,
    height: int = 512,
    num_steps: int = 20,
    start_step: int = 0,
    guidance: float = 4.0,
    seed: int = -1,
    id_weight: float = 1.0,
    neg_prompt: str = "",
    true_cfg: float = 1.0,
    timestep_to_start_cfg: int = 1,
    max_sequence_length: int = 128,
    gamma: float = 0.5,
    eta: float = 0.7,
    s: float = 0,
    tau: float = 5,
):
    # 모델이 초기화되지 않았으면 오류 메시지 반환
    if not model_initialized:
        return None, "GPU 오류: CUDA GPU를 찾을 수 없어 모델을 초기화할 수 없습니다.", None
    
    # ID 이미지가 없으면 실행 불가
    if id_image is None:
        return None, "오류: ID 이미지가 필요합니다.", None
    
    try:
        flux_generator.t5.max_length = max_sequence_length

        # 시드 설정
        seed = int(seed)
        if seed == -1:
            seed = None

        opts = SamplingOptions(
            prompt=prompt,
            width=width,
            height=height,
            num_steps=num_steps,
            guidance=guidance,
            seed=seed,
        )

        if opts.seed is None:
            opts.seed = torch.Generator(device="cpu").seed()

        seed_everything(opts.seed)
        print(f"Generating prompt: '{opts.prompt}' (seed={opts.seed})...")
        t0 = time.perf_counter()

        use_true_cfg = abs(true_cfg - 1.0) > 1e-6

        # 1) 입력 노이즈 준비
        noise = get_noise(
            num_samples=1,
            height=opts.height,
            width=opts.width,
            device=flux_generator.device,
            dtype=torch.bfloat16 if flux_generator.device.type == 'cuda' else torch.float32,
            seed=opts.seed,
        )
        bs, c, h, w = noise.shape
        noise = rearrange(noise, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
        if noise.shape[0] == 1 and bs > 1:
            noise = repeat(noise, "1 ... -> bs ...", bs=bs)

        # ID 이미지 인코딩
        encode_t0 = time.perf_counter()
        id_image = id_image.resize((opts.width, opts.height), resample=Image.LANCZOS)
        x = torch.from_numpy(np.array(id_image).astype(np.float32))
        x = (x / 127.5) - 1.0
        x = rearrange(x, "h w c -> 1 c h w")
        x = x.to(flux_generator.device)
        
        dtype = torch.bfloat16 if flux_generator.device.type == 'cuda' else torch.float32
        with torch.autocast(device_type=flux_generator.device.type, dtype=dtype):
            x = flux_generator.ae.encode(x)
        x = x.to(dtype)

        encode_t1 = time.perf_counter()
        print(f"Encoded in {encode_t1 - encode_t0:.2f} seconds.")
        
        timesteps = get_schedule(opts.num_steps, x.shape[-1] * x.shape[-2] // 4, shift=False)

        # 2) 텍스트 임베딩 준비
        inp = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt=opts.prompt)
        inp_inversion = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt="")
        inp_neg = None
        if use_true_cfg:
            inp_neg = prepare(t5=flux_generator.t5, clip=flux_generator.clip_model, img=x, prompt=neg_prompt)

        # 3) ID 임베딩 생성
        id_embeddings = None
        uncond_id_embeddings = None
        if id_image is not None:
            id_image = np.array(id_image)
            id_image = resize_numpy_image_long(id_image, 1024)
            id_embeddings, uncond_id_embeddings = flux_generator.pulid_model.get_id_embedding(id_image, cal_uncond=use_true_cfg)

        y_0 = inp["img"].clone().detach()

        # 이미지 처리 과정
        inverted = rf_inversion(
            flux_generator.model,
            **inp_inversion,
            timesteps=timesteps,
            guidance=opts.guidance,
            id=id_embeddings,
            id_weight=id_weight,
            start_step=start_step,
            uncond_id=uncond_id_embeddings,
            true_cfg=true_cfg,
            timestep_to_start_cfg=timestep_to_start_cfg,
            neg_txt=inp_neg["txt"] if use_true_cfg else None,
            neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None,
            neg_vec=inp_neg["vec"] if use_true_cfg else None,
            aggressive_offload=False,
            y_1=noise,
            gamma=gamma
        )
        
        inp["img"] = inverted
        inp_inversion["img"] = inverted

        edited = rf_denoise(
            flux_generator.model,
            **inp,
            timesteps=timesteps,
            guidance=opts.guidance,
            id=id_embeddings,
            id_weight=id_weight,
            start_step=start_step,
            uncond_id=uncond_id_embeddings,
            true_cfg=true_cfg,
            timestep_to_start_cfg=timestep_to_start_cfg,
            neg_txt=inp_neg["txt"] if use_true_cfg else None,
            neg_txt_ids=inp_neg["txt_ids"] if use_true_cfg else None,
            neg_vec=inp_neg["vec"] if use_true_cfg else None,
            aggressive_offload=False,
            y_0=y_0,
            eta=eta,
            s=s,
            tau=tau,
        )

        # 결과 이미지 디코딩
        edited = unpack(edited.float(), opts.height, opts.width)
        with torch.autocast(device_type=flux_generator.device.type, dtype=dtype):
            edited = flux_generator.ae.decode(edited)

        t1 = time.perf_counter()
        print(f"Done in {t1 - t0:.2f} seconds.")

        # PIL 이미지로 변환
        edited = edited.clamp(-1, 1)
        edited = rearrange(edited[0], "c h w -> h w c")
        edited = Image.fromarray((127.5 * (edited + 1.0)).cpu().byte().numpy())

        return edited, str(opts.seed), flux_generator.pulid_model.debug_img_list
    
    except Exception as e:
        import traceback
        error_msg = f"이미지 생성 중 오류 발생: {str(e)}\n{traceback.format_exc()}"
        print(error_msg)
        return None, error_msg, None


def create_demo():
    with gr.Blocks() as demo:
        gr.Markdown("# PuLID: 인물 이미지 변환 도구")
        
        if not model_initialized:
            gr.Markdown("## ⚠️ 오류: CUDA GPU를 찾을 수 없습니다")
            gr.Markdown("이 응용 프로그램은 CUDA 지원 GPU가 필요합니다. CPU에서는 실행할 수 없습니다.")
            return demo
            
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="프롬프트", value="portrait, color, cinematic")
                id_image = gr.Image(label="ID 이미지", type="pil")
                id_weight = gr.Slider(0.0, 1.0, 0.4, step=0.05, label="ID 가중치")
                num_steps = gr.Slider(1, 24, 16, step=1, label="단계 수")
                guidance = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="가이던스")

                with gr.Accordion("고급 옵션", open=False):
                    neg_prompt = gr.Textbox(label="네거티브 프롬프트", value="")
                    true_cfg = gr.Slider(1.0, 10.0, 3.5, step=0.1, label="CFG 스케일")
                    seed = gr.Textbox(-1, label="시드 (-1: 랜덤)")
                    gr.Markdown("### 기타 옵션")
                    gamma = gr.Slider(0.0, 1.0, 0.5, step=0.1, label="감마")
                    eta = gr.Slider(0.0, 1.0, 0.8, step=0.1, label="에타")

                generate_btn = gr.Button("이미지 생성")

            with gr.Column():
                output_image = gr.Image(label="생성된 이미지")
                seed_output = gr.Textbox(label="결과/오류 메시지")
                gr.Markdown(_CITE_)

        generate_btn.click(
            fn=generate_image,
            inputs=[prompt, id_image, 512, 512, num_steps, 0, guidance, seed, id_weight, neg_prompt,
                    true_cfg, 1, 128, gamma, eta, 0, 5],
            outputs=[output_image, seed_output],
        )

    return demo


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="PuLID for FLUX.1-dev")
    parser.add_argument('--version', type=str, default='v0.9.1')
    parser.add_argument("--name", type=str, default="flux-dev")
    parser.add_argument("--port", type=int, default=8080)
    args = parser.parse_args()

    demo = create_demo()
    demo.launch(ssr_mode=False)