import gradio as gr
import torch
import numpy as np
from PIL import Image
import os

# 모델 import들
try:
    from diffusers import ControlNetModel, StableDiffusionControlNetInpaintPipeline, UniPCMultistepScheduler
    from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
    from controlnet_aux import MLSDdetector
    MLSD_AVAILABLE = True
except ImportError as e:
    print(f"일부 라이브러리 로딩 실패: {e}")
    MLSD_AVAILABLE = False

class SpacelyFurnitureDesigner:
    def __init__(self):
        self.pipe = None
        self.seg_processor = None
        self.segmentor = None
        self.mlsd_processor = None
        
        # 회사용 프롬프트 템플릿
        self.office_templates = {
            "개인사무실": "modern private office with executive desk, ergonomic chair, bookshelf, and professional lighting",
            "회의실": "professional conference room with large meeting table, comfortable chairs, whiteboard, and corporate lighting", 
            "휴게실": "corporate break room with comfortable seating, coffee table, plants, and relaxing atmosphere",
            "오픈오피스": "open office space with multiple workstations, modern desks, office chairs, and collaborative areas",
            "리셉션": "elegant reception area with reception desk, waiting chairs, company logo wall, and welcoming ambiance",
            "CEO실": "luxury executive office with premium wooden desk, leather chair, awards display, and elegant lighting"
        }
        
        self.quality_suffix = "professional interior design, corporate style, clean, modern, functional, well-lit, 4K, high quality"
        
    def load_models(self):
        """모델 지연 로딩"""
        if self.pipe is None:
            print("🔄 AI 모델 로딩 중...")
            
            try:
                if MLSD_AVAILABLE and torch.cuda.is_available():
                    # 전체 ControlNet 설정
                    controlnet = [
                        ControlNetModel.from_pretrained("BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16),
                        ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-mlsd", torch_dtype=torch.float16)
                    ]
                    self.mlsd_processor = MLSDdetector.from_pretrained("lllyasviel/Annotators")
                else:
                    # 세그멘테이션만 사용
                    controlnet = ControlNetModel.from_pretrained("BertChristiaens/controlnet-seg-room", torch_dtype=torch.float16)
                
                # 메인 파이프라인
                self.pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
                    "SG161222/Realistic_Vision_V3.0_VAE",
                    controlnet=controlnet,
                    safety_checker=None,
                    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
                )
                
                self.pipe.scheduler = UniPCMultistepScheduler.from_config(self.pipe.scheduler.config)
                
                if torch.cuda.is_available():
                    self.pipe = self.pipe.to("cuda")
                    try:
                        self.pipe.enable_xformers_memory_efficient_attention()
                    except:
                        pass
                
                # 세그멘테이션 모델
                self.seg_processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b5-finetuned-ade-640-640")
                self.segmentor = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b5-finetuned-ade-640-640")
                
                print("✅ 모델 로딩 완료!")
                
            except Exception as e:
                print(f"❌ 모델 로딩 실패: {e}")
                return False
                
        return True
    
    def resize_dimensions(self, dimensions, target_size=768):
        """비율 유지하며 리사이즈"""
        width, height = dimensions
        if width < target_size and height < target_size:
            return dimensions
        if width > height:
            aspect_ratio = height / width
            return (target_size, int(target_size * aspect_ratio))
        else:
            aspect_ratio = width / height
            return (int(target_size * aspect_ratio), target_size)
    
    def create_simple_mask(self, image):
        """간단한 마스크 생성"""
        # 전체 이미지를 변경 대상으로 설정
        mask = Image.new('RGB', image.size, (255, 255, 255))
        return image, mask
    
    def design_space(self, input_image, space_type, custom_prompt="", num_steps=30, guidance_scale=12):
        """공간 디자인 생성 메인 함수"""
        
        if input_image is None:
            return None, "❌ 이미지를 업로드해주세요!"
        
        # 모델 로딩
        if not self.load_models():
            return None, "❌ AI 모델 로딩에 실패했습니다."
        
        try:
            # 프롬프트 구성
            if custom_prompt.strip():
                base_prompt = custom_prompt
            else:
                base_prompt = self.office_templates.get(space_type, self.office_templates["개인사무실"])
            
            full_prompt = f"{base_prompt}, {self.quality_suffix}"
            
            # 이미지 전처리
            orig_w, orig_h = input_image.size
            new_width, new_height = self.resize_dimensions(input_image.size, 768)
            resized_image = input_image.resize((new_width, new_height))
            
            # 마스크 생성
            seg_image, mask_image = self.create_simple_mask(resized_image)
            
            # ControlNet 이미지 준비
            if MLSD_AVAILABLE and self.mlsd_processor:
                mlsd_image = self.mlsd_processor(resized_image)
                mlsd_image = mlsd_image.resize(resized_image.size)
                control_images = [seg_image, mlsd_image]
                controlnet_conditioning_scale = [0.4, 0.2]
                control_guidance_start = [0, 0.1] 
                control_guidance_end = [0.5, 0.25]
            else:
                control_images = seg_image
                controlnet_conditioning_scale = 0.4
                control_guidance_start = 0
                control_guidance_end = 0.5
            
            # AI 이미지 생성
            result = self.pipe(
                prompt=full_prompt,
                negative_prompt="lowres, watermark, blurry, unprofessional, cluttered, outdated furniture, bad quality",
                num_inference_steps=num_steps,
                strength=0.8,
                guidance_scale=guidance_scale,
                image=resized_image,
                mask_image=mask_image,
                control_image=control_images,
                controlnet_conditioning_scale=controlnet_conditioning_scale,
                control_guidance_start=control_guidance_start,
                control_guidance_end=control_guidance_end
            ).images[0]
            
            # 원본 크기로 복원
            final_image = result.resize((orig_w, orig_h), Image.Resampling.LANCZOS)
            
            success_msg = f"✅ {space_type} 디자인 완료!\n📝 사용된 프롬프트: {full_prompt[:100]}..."
            
            return final_image, success_msg
            
        except Exception as e:
            error_msg = f"❌ 디자인 생성 중 오류 발생: {str(e)}"
            print(error_msg)
            return None, error_msg

# 전역 디자이너 인스턴스
designer = SpacelyFurnitureDesigner()

def create_ui():
    """Gradio UI 생성"""
    
    with gr.Blocks(
        title="🏢 Spacely AI 가구 배치 디자이너",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        .title {
            text-align: center;
            color: #2D3748;
            margin-bottom: 20px;
        }
        """
    ) as demo:
        
        gr.HTML("""
        <div class="title">
            <h1>🏢 Spacely AI 가구 배치 디자이너</h1>
            <p>빈 방 사진을 업로드하면 AI가 전문적인 오피스 공간으로 디자인해드립니다</p>
        </div>
        """)
        
        with gr.Row():
            # 입력 컬럼
            with gr.Column(scale=1):
                gr.HTML("<h3>📁 입력 설정</h3>")
                
                input_image = gr.Image(
                    label="빈방 이미지 업로드",
                    type="pil",
                    height=300
                )
                
                space_type = gr.Dropdown(
                    choices=list(designer.office_templates.keys()),
                    label="공간 타입 선택",
                    value="개인사무실"
                )
                
                custom_prompt = gr.Textbox(
                    label="커스텀 프롬프트 (선택사항)",
                    placeholder="예: minimalist CEO office with wooden desk...",
                    lines=3
                )
                
                with gr.Row():
                    num_steps = gr.Slider(
                        minimum=10,
                        maximum=50,
                        value=30,
                        step=5,
                        label="생성 품질 (높을수록 느림)"
                    )
                    
                    guidance_scale = gr.Slider(
                        minimum=5,
                        maximum=20,
                        value=12,
                        step=1,
                        label="프롬프트 반영도"
                    )
                
                generate_btn = gr.Button(
                    "🎨 AI 디자인 생성",
                    variant="primary",
                    size="lg"
                )
            
            # 출력 컬럼  
            with gr.Column(scale=1):
                gr.HTML("<h3>✨ 디자인 결과</h3>")
                
                output_image = gr.Image(
                    label="AI가 디자인한 오피스",
                    height=300
                )
                
                output_text = gr.Textbox(
                    label="생성 결과",
                    lines=4,
                    max_lines=8
                )
                
                gr.HTML("""
                <div style="margin-top: 20px; padding: 15px; background-color: #F7FAFC; border-radius: 10px;">
                    <h4>💡 사용 팁</h4>
                    <ul>
                        <li>깔끔한 빈방 사진일수록 좋은 결과를 얻을 수 있습니다</li>
                        <li>커스텀 프롬프트로 원하는 가구나 스타일을 지정할 수 있습니다</li>
                        <li>생성 품질을 높이면 더 정교한 결과를 얻지만 시간이 오래 걸립니다</li>
                    </ul>
                </div>
                """)
        
        # 이벤트 연결
        generate_btn.click(
            fn=designer.design_space,
            inputs=[input_image, space_type, custom_prompt, num_steps, guidance_scale],
            outputs=[output_image, output_text]
        )
        
        # 예시 이미지들
        gr.HTML("<h3>📸 예시 결과</h3>")
        
        with gr.Row():
            gr.Examples(
                examples=[
                    ["개인사무실", "modern executive office with wooden desk"],
                    ["회의실", "professional conference room for 10 people"],
                    ["휴게실", "comfortable break room with plants and coffee area"],
                ],
                inputs=[space_type, custom_prompt],
                label="빠른 설정 예시"
            )
    
    return demo

# 앱 실행
if __name__ == "__main__":
    demo = create_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        share=True  # 공개 URL 생성
    )