import os
import sys
import numpy as np
import PIL.Image
import torch
import torchvision.transforms as T
from huggingface_hub import hf_hub_download
import gradio as gr
import time
import cv2

# افزودن مسیر مورد نیاز برای ماژول‌های CelebAMask-HQ
celebamask_path = "/home/user/app/CelebAMask-HQ"
face_parsing_path = os.path.join(celebamask_path, "face_parsing")
sys.path.insert(0, celebamask_path)
sys.path.insert(0, face_parsing_path)

print("Python path:", sys.path)
print("CelebAMask path exists:", os.path.exists(celebamask_path))
print("Face parsing path exists:", os.path.exists(face_parsing_path))

# ایمپورت ماژول‌های مورد نیاز
try:
    from unet import unet
    from utils import generate_label
    IMPORT_SUCCESS = True
    print("✅ Successfully imported CelebAMask-HQ modules")
except ImportError as e:
    IMPORT_SUCCESS = False
    print(f"❌ Failed to import CelebAMask-HQ modules: {e}")

# تنظیمات دستگاه
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# تنظیم مسیرهای کش
os.environ["HF_HOME"] = "/home/user/app/hf_cache"

# تعریف transform
transform = T.Compose([
    T.Resize((512, 512)),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# کلاس‌های Face Parsing
CELEBA_CLASSES = [
    'background', 'skin', 'l_brow', 'r_brow', 'l_eye', 'r_eye', 'eye_g', 'l_ear', 'r_ear', 'ear_r',
    'nose', 'mouth', 'u_lip', 'l_lip', 'neck', 'neck_l', 'cloth', 'hair', 'hat'
]

class FaceParsingModel:
    def __init__(self):
        self.model = None
        self.device = device
        self.load_model()
    
    def load_model(self):
        """لود مدل Face Parsing"""
        try:
            print("📥 Downloading model...")
            model_path = hf_hub_download(
                repo_id="public-data/CelebAMask-HQ-Face-Parsing",
                filename="models/model.pth",
                cache_dir="/home/user/app/hf_cache"
            )
            print(f"✅ Model downloaded to: {model_path}")
            
            # ایجاد مدل با معماری صحیح
            self.model = unet(
                feature_scale=4,
                n_classes=19,
                is_deconv=True,
                in_channels=3,
                is_batchnorm=True
            )
            
            # لود state dict
            state_dict = torch.load(model_path, map_location="cpu")
            
            # اگر state dict از DataParallel باشد، module. را حذف می‌کنیم
            new_state_dict = {}
            for k, v in state_dict.items():
                if k.startswith('module.'):
                    k = k[7:]
                new_state_dict[k] = v
            
            self.model.load_state_dict(new_state_dict)
            self.model.eval()
            self.model.to(self.device)
            
            print("✅ Model loaded successfully")
            
        except Exception as e:
            print(f"❌ Failed to load model: {e}")
            import traceback
            traceback.print_exc()
            self.model = None
    
    def predict(self, image):
        """پردازش تصویر و تولید ماسک"""
        if self.model is None:
            raise ValueError("Model not loaded properly")
        
        # تبدیل به PIL Image اگر لازم است
        if isinstance(image, str):
            image = PIL.Image.open(image).convert('RGB')
        elif isinstance(image, np.ndarray):
            image = PIL.Image.fromarray(image)
        
        # ذخیره تصویر اصلی
        original_image = image.copy()
        
        # پیش‌پردازش
        data = transform(image)
        data = data.unsqueeze(0).to(self.device)

        # پیش‌بینی
        with torch.no_grad():
            out = self.model(data)
            label_out = generate_label(out, 512)
            mask = label_out[0].cpu().numpy()
            
            # رنگ‌آمیزی ماسک
            colored_mask = self.colorize_mask(mask)
            
        # ترکیب تصویر اصلی با ماسک
        resized_image = np.asarray(original_image.resize((512, 512)))
        blended = cv2.addWeighted(resized_image, 0.7, colored_mask, 0.3, 0)

        return colored_mask, blended

    def colorize_mask(self, mask):
        """رنگ‌آمیزی ماسک بر اساس کلاس‌ها"""
        # پالت رنگ برای 19 کلاس (متفاوت برای تشخیص بهتر)
        palette = [
            [0, 0, 0],        # background - سیاه
            [255, 200, 200],  # skin - پوست
            [0, 255, 0],      # l_brow - سبز
            [0, 200, 0],      # r_brow - سبز تیره
            [255, 0, 0],      # l_eye - قرمز
            [200, 0, 0],      # r_eye - قرمز تیره
            [255, 255, 0],    # eye_g - زرد
            [0, 0, 255],      # l_ear - آبی
            [0, 0, 200],      # r_ear - آبی تیره
            [128, 0, 128],    # ear_r - بنفش
            [255, 165, 0],    # nose - نارنجی
            [255, 0, 255],    # mouth - صورتی
            [200, 0, 200],    # u_lip - صورتی تیره
            [165, 42, 42],    # l_lip - قهوه‌ای
            [0, 255, 255],    # neck - فیروزه‌ای
            [0, 200, 200],    # neck_l - فیروزه‌ای تیره
            [128, 128, 128],  # cloth - خاکستری
            [255, 255, 255],  # hair - سفید
            [255, 215, 0]     # hat - طلایی
        ]
        
        colored = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
        for i in range(len(palette)):
            colored[mask == i] = palette[i]
        
        return colored

def initialize_app():
    """Initialize application"""
    print("===== Application Startup at {} =====".format(time.strftime("%Y-%m-%d %H:%M:%S")))
    
    print("[Info] PYTHONPATH:", os.environ.get("PYTHONPATH"))
    print("[Info] CelebAMask-HQ path exists:", os.path.exists(celebamask_path))
    print("[Info] face_parsing folder exists:", os.path.exists(face_parsing_path))
    print("[Info] Module import success:", IMPORT_SUCCESS)
    
    try:
        face_parser = FaceParsingModel()
        success = face_parser.model is not None
        status_msg = "Model loaded successfully" if success else "Model failed to load"
        return success, status_msg, face_parser
    except Exception as e:
        print(f"[Error] Initialization failed: {e}")
        return False, f"Initialization failed: {e}", None

# Initialize the application
success, status_msg, face_parser = initialize_app()

def process_image(input_image):
    """پردازش تصویر ورودی"""
    if input_image is None:
        return None, None, "لطفاً یک تصویر آپلود کنید"
    
    if not success or face_parser is None:
        return None, None, "❌ مدل لود نشده است. لطفاً دوباره تلاش کنید."
    
    try:
        # پردازش تصویر
        mask, blended = face_parser.predict(input_image)
        
        # اطلاعات پردازش
        if isinstance(input_image, str):
            original_img = PIL.Image.open(input_image)
            img_size = original_img.size
        else:
            img_size = input_image.size if hasattr(input_image, 'size') else input_image.shape[:2][::-1]
        
        info_text = f"""
        ✅ پردازش انجام شد!
        - اندازه تصویر ورودی: {img_size}
        - اندازه خروجی: 512x512
        - کلاس‌های تشخیص: {len(CELEBA_CLASSES)}
        - دستگاه پردازش: {device}
        """
        
        return blended, mask, info_text
        
    except Exception as e:
        error_msg = f"❌ خطا در پردازش تصویر: {str(e)}"
        print(error_msg)
        import traceback
        traceback.print_exc()
        return None, None, error_msg

def create_legend():
    """ایجاد لیجند برای کلاس‌ها"""
    import matplotlib.pyplot as plt
    
    legend_html = """
    <div style='max-height: 300px; overflow-y: auto; border: 1px solid #ccc; padding: 10px; border-radius: 5px;'>
    <h4>🎨 Legend - کلاس‌های Face Parsing:</h4>
    """
    
    colors = plt.get_cmap('tab20', len(CELEBA_CLASSES))
    
    for i, class_name in enumerate(CELEBA_CLASSES):
        color = colors(i)
        color_hex = '#%02x%02x%02x' % (int(color[0]*255), int(color[1]*255), int(color[2]*255))
        text_color = 'white' if color[0] * 0.299 + color[1] * 0.587 + color[2] * 0.114 < 0.5 else 'black'
        legend_html += f"""
        <div style='margin: 2px; padding: 5px; background-color: {color_hex}; color: {text_color}; border-radius: 3px;'>
        <strong>{i}:</strong> {class_name}
        </div>
        """
    
    legend_html += "</div>"
    return legend_html

# ایجاد اینترفیس Gradio
with gr.Blocks(title="CelebAMask-HQ Face Parsing", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎭 CelebAMask-HQ Face Parsing Demo
    **آپلود یک تصویر صورت و دریافت خروجی Face Parsing**
    
    این مدل صورت را به 19 بخش مختلف تقسیم می‌کند (پوست، چشم، ابرو، بینی، دهان، مو و ...)
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="📷 تصویر ورودی",
                type="filepath",
                sources=["upload"],
                height=300
            )
            process_btn = gr.Button("🚀 پردازش تصویر", variant="primary", size="lg")
            
            with gr.Accordion("ℹ️ وضعیت برنامه", open=False):
                status_display = gr.Markdown(f"""
                **وضعیت:**
                - 🎯 مدل: {'✅ لود شده' if success else '❌ خطا در لود'}
                - 💻 دستگاه: `{device}`
                - 📦 ماژول‌ها: {'✅ ایمپورت شده' if IMPORT_SUCCESS else '❌ خطا در ایمپورت'}
                - 🗂️ کلاس‌ها: {len(CELEBA_CLASSES)}
                """)
        
        with gr.Column():
            output_blended = gr.Image(
                label="🎨 نتیجه ترکیبی (تصویر + ماسک)",
                height=300
            )
            output_mask = gr.Image(
                label="🎭 ماسک سگمنتیشن",
                height=300
            )
    
    with gr.Row():
        info_output = gr.Textbox(
            label="📊 اطلاعات پردازش",
            lines=3,
            max_lines=6
        )
    
    with gr.Row():
        gr.HTML(create_legend())
    
    # اتصال رویدادها
    process_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_blended, output_mask, info_output]
    )
    
    input_image.upload(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_blended, output_mask, info_output]
    )

if __name__ == "__main__":
    print("🚀 Starting Face Parsing Application...")
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )