import tempfile
import time
from collections.abc import Sequence
from typing import Any, cast
import os
from huggingface_hub import login, hf_hub_download

import gradio as gr
import numpy as np
import pillow_heif
import spaces
import torch
from gradio_image_annotation import image_annotator
from gradio_imageslider import ImageSlider
from PIL import Image
from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
from refiners.fluxion.utils import no_grad
from refiners.solutions import BoxSegmenter
from transformers import GroundingDinoForObjectDetection, GroundingDinoProcessor
from diffusers import FluxPipeline

BoundingBox = tuple[int, int, int, int]

pillow_heif.register_heif_opener()
pillow_heif.register_avif_opener()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# HF 토큰 설정
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN is None:
    raise ValueError("Please set the HF_TOKEN environment variable")

try:
    login(token=HF_TOKEN)
except Exception as e:
    raise ValueError(f"Failed to login to Hugging Face: {str(e)}")

# 모델 초기화
segmenter = BoxSegmenter(device="cpu")
segmenter.device = device
segmenter.model = segmenter.model.to(device=segmenter.device)

gd_model_path = "IDEA-Research/grounding-dino-base"
gd_processor = GroundingDinoProcessor.from_pretrained(gd_model_path)
gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_dtype=torch.float32)
gd_model = gd_model.to(device=device)
assert isinstance(gd_model, GroundingDinoForObjectDetection)

# FLUX 파이프라인 초기화
pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    torch_dtype=torch.bfloat16,
    use_auth_token=HF_TOKEN
)
pipe.load_lora_weights(
    hf_hub_download(
        "ByteDance/Hyper-SD",
        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
        use_auth_token=HF_TOKEN
    )
)
pipe.fuse_lora(lora_scale=0.125)
pipe.to(device="cuda", dtype=torch.bfloat16)

class timer:
    def __init__(self, method_name="timed process"):
        self.method = method_name
    def __enter__(self):
        self.start = time.time()
        print(f"{self.method} starts")
    def __exit__(self, exc_type, exc_val, exc_tb):
        end = time.time()
        print(f"{self.method} took {str(round(end - self.start, 2))}s")

def bbox_union(bboxes: Sequence[list[int]]) -> BoundingBox | None:
    if not bboxes:
        return None
    for bbox in bboxes:
        assert len(bbox) == 4
        assert all(isinstance(x, int) for x in bbox)
    return (
        min(bbox[0] for bbox in bboxes),
        min(bbox[1] for bbox in bboxes),
        max(bbox[2] for bbox in bboxes),
        max(bbox[3] for bbox in bboxes),
    )

def corners_to_pixels_format(bboxes: torch.Tensor, width: int, height: int) -> torch.Tensor:
    x1, y1, x2, y2 = bboxes.round().to(torch.int32).unbind(-1)
    return torch.stack((x1.clamp_(0, width), y1.clamp_(0, height), x2.clamp_(0, width), y2.clamp_(0, height)), dim=-1)

def gd_detect(img: Image.Image, prompt: str) -> BoundingBox | None:
    inputs = gd_processor(images=img, text=f"{prompt}.", return_tensors="pt").to(device=device)
    with no_grad():
        outputs = gd_model(**inputs)
    width, height = img.size
    results: dict[str, Any] = gd_processor.post_process_grounded_object_detection(
        outputs,
        inputs["input_ids"],
        target_sizes=[(height, width)],
    )[0]
    assert "boxes" in results and isinstance(results["boxes"], torch.Tensor)
    bboxes = corners_to_pixels_format(results["boxes"].cpu(), width, height)
    return bbox_union(bboxes.numpy().tolist())

def apply_mask(img: Image.Image, mask_img: Image.Image, defringe: bool = True) -> Image.Image:
    assert img.size == mask_img.size
    img = img.convert("RGB")
    mask_img = mask_img.convert("L")
    if defringe:
        rgb, alpha = np.asarray(img) / 255.0, np.asarray(mask_img) / 255.0
        foreground = cast(np.ndarray[Any, np.dtype[np.uint8]], estimate_foreground_ml(rgb, alpha))
        img = Image.fromarray((foreground * 255).astype("uint8"))
    result = Image.new("RGBA", img.size)
    result.paste(img, (0, 0), mask_img)
    return result

def generate_background(prompt: str, width: int, height: int) -> Image.Image:
    """배경 이미지 생성 함수"""
    try:
        with timer("Background generation"):
            image = pipe(
                prompt=prompt,
                width=width,
                height=height,
                num_inference_steps=8,
                guidance_scale=4.0,
            ).images[0]
        return image
    except Exception as e:
        raise gr.Error(f"Background generation failed: {str(e)}")

def combine_with_background(foreground: Image.Image, background: Image.Image) -> Image.Image:
    """전경과 배경 합성 함수"""
    background = background.resize(foreground.size)
    return Image.alpha_composite(background.convert('RGBA'), foreground)

@spaces.GPU
def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Image.Image, BoundingBox | None, list[str]]:
    time_log: list[str] = []
    if isinstance(prompt, str):
        t0 = time.time()
        bbox = gd_detect(img, prompt)
        time_log.append(f"detect: {time.time() - t0}")
        if not bbox:
            print(time_log[0])
            raise gr.Error("No object detected")
    else:
        bbox = prompt
    t0 = time.time()
    mask = segmenter(img, bbox)
    time_log.append(f"segment: {time.time() - t0}")
    return mask, bbox, time_log

def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str | None = None) -> tuple[tuple[Image.Image, Image.Image, Image.Image], gr.DownloadButton]:
    if img.width > 2048 or img.height > 2048:
        orig_res = max(img.width, img.height)
        img.thumbnail((2048, 2048))
        if isinstance(prompt, tuple):
            x0, y0, x1, y1 = (int(x * 2048 / orig_res) for x in prompt)
            prompt = (x0, y0, x1, y1)

    mask, bbox, time_log = _gpu_process(img, prompt)
    masked_alpha = apply_mask(img, mask, defringe=True)

    if bg_prompt:
        try:
            background = generate_background(bg_prompt, img.width, img.height)
            combined = combine_with_background(masked_alpha, background)
        except Exception as e:
            raise gr.Error(f"Background processing failed: {str(e)}")
    else:
        combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)

    thresholded = mask.point(lambda p: 255 if p > 10 else 0)
    bbox = thresholded.getbbox()
    to_dl = masked_alpha.crop(bbox)

    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    to_dl.save(temp, format="PNG")
    temp.close()

    return (img, combined, masked_alpha), gr.DownloadButton(value=temp.name, interactive=True)

def process_bbox(img: Image.Image, box_input: str) -> tuple[list[Image.Image], str]:
    try:
        if img is None or box_input.strip() == "":
            raise gr.Error("Please provide both image and bounding box coordinates")
        
        # Parse box coordinates
        try:
            coords = eval(box_input)
            if not isinstance(coords, list) or len(coords) != 4:
                raise ValueError("Invalid box format")
            bbox = tuple(int(x) for x in coords)
        except:
            raise gr.Error("Invalid box format. Please provide [xmin, ymin, xmax, ymax]")
        
        # Process the image
        results, download_path = _process(img, bbox)
        
        # Convert results to list for gallery
        gallery_images = list(results)
        
        return gallery_images, download_path
    except Exception as e:
        raise gr.Error(str(e))


def on_change_bbox(prompts: dict[str, Any] | None):
    return gr.update(interactive=prompts is not None)


def on_change_prompt(img: Image.Image | None, prompt: str | None, bg_prompt: str | None = None):
    return gr.update(interactive=bool(img and prompt))

def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None) -> tuple[list[Image.Image], str]:
    try:
        if img is None or prompt.strip() == "":
            raise gr.Error("Please provide both image and prompt")
        
        # Process the image
        results, download_path = _process(img, prompt, bg_prompt)
        
        # Convert results to list for gallery
        gallery_images = list(results)
        
        return gallery_images, download_path
    except Exception as e:
        raise gr.Error(str(e))


def update_process_button(img, prompt):
    return gr.Button.update(
        interactive=bool(img and prompt),
        variant="primary" if bool(img and prompt) else "secondary"
    )
    
def update_box_button(img, box_input):
    try:
        if img and box_input:
            coords = eval(box_input)
            if isinstance(coords, list) and len(coords) == 4:
                return gr.Button.update(interactive=True, variant="primary")
        return gr.Button.update(interactive=False, variant="secondary")
    except:
        return gr.Button.update(interactive=False, variant="secondary")


# 맨 앞부분에 CSS 정의 추가
css = """
footer {display: none}
.main-title {
    text-align: center;
    margin: 2em 0;
    padding: 1em;
    background: #f7f7f7;
    border-radius: 10px;
}
.main-title h1 {
    color: #2196F3;
    font-size: 2.5em;
    margin-bottom: 0.5em;
}
.main-title p {
    color: #666;
    font-size: 1.2em;
}
.container {
    max-width: 1200px;
    margin: auto;
    padding: 20px;
}
.tabs {
    margin-top: 1em;
}
.input-group {
    background: white;
    padding: 1em;
    border-radius: 8px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.output-group {
    background: white;
    padding: 1em;
    border-radius: 8px;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
button.primary {
    background: #2196F3;
    border: none;
    color: white;
    padding: 0.5em 1em;
    border-radius: 4px;
    cursor: pointer;
    transition: background 0.3s ease;
}
button.primary:hover {
    background: #1976D2;
}
"""

# UI 부분만 수정
# Main Gradio app
with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
    gr.HTML("""
        <div class="main-title">
            <h1>🎨 Image Object Extractor</h1>
            <p>Extract objects from images using text prompts or bounding boxes</p>
        </div>
    """)

    with gr.Tabs(selected=0):
        # Text-based extraction tab
        with gr.TabItem("Extract by Text"):
            with gr.Row():
                with gr.Column(scale=1):
                    input_image = gr.Image(
                        type="pil",
                        label="Upload Image",
                        interactive=True
                    )
                    text_prompt = gr.Textbox(
                        label="Object to Extract",
                        placeholder="Enter what you want to extract...",
                        interactive=True
                    )
                    bg_prompt = gr.Textbox(
                        label="Background Prompt (optional)",
                        placeholder="Describe the background...",
                        interactive=True
                    )
                    process_btn = gr.Button(
                        "Process",
                        variant="primary",
                        interactive=False
                    )

                with gr.Column(scale=1):
                    output_display = gr.Gallery(
                        label="Results",
                        show_download_button=False,
                        visible=True
                    )
                    download_btn = gr.DownloadButton(
                        "Download Result",
                        visible=True
                    )

        # Box-based extraction tab
        with gr.TabItem("Extract by Box"):
            with gr.Row():
                with gr.Column(scale=1):
                    box_image = gr.Image(
                        type="pil",
                        label="Upload Image for Box",
                        interactive=True
                    )
                    box_input = gr.Textbox(
                        label="Bounding Box (xmin, ymin, xmax, ymax)",
                        placeholder="Enter coordinates as [x1, y1, x2, y2]",
                        interactive=True
                    )
                    box_btn = gr.Button(
                        "Extract Selection",
                        variant="primary",
                        interactive=False
                    )

                with gr.Column(scale=1):
                    box_output = gr.Gallery(
                        label="Results",
                        show_download_button=False,
                        visible=True
                    )
                    box_download = gr.DownloadButton(
                        "Download Result",
                        visible=True
                    )

    # Event bindings
    input_image.change(
        fn=update_process_button,
        inputs=[input_image, text_prompt],
        outputs=process_btn,
        queue=False
    )
    
    text_prompt.change(
        fn=update_process_button,
        inputs=[input_image, text_prompt],
        outputs=process_btn,
        queue=False
    )

    process_btn.click(
        fn=process_prompt,
        inputs=[input_image, text_prompt, bg_prompt],
        outputs=[output_display, download_btn],
        queue=True
    )

    box_image.change(
        fn=update_box_button,
        inputs=[box_image, box_input],
        outputs=box_btn,
        queue=False
    )

    box_input.change(
        fn=update_box_button,
        inputs=[box_image, box_input],
        outputs=box_btn,
        queue=False
    )

    box_btn.click(
        fn=process_bbox,
        inputs=[box_image, box_input],
        outputs=[box_output, box_download],
        queue=True
    )

demo.queue(max_size=30, api_open=False)
demo.launch()