import streamlit as st
# تمت إزالة استيراد components
# import streamlit.components.v1 as components # لاستقبال الرسائل
from PIL import Image
import torch
from torchvision import transforms
from transformers import AutoModelForImageSegmentation
import io
import os
import requests
from io import BytesIO

# --- 1. إعدادات الصفحة (تبقى في الأعلى) ---
st.set_page_config(
    page_title="SkyData - Background Removal",
    layout="wide"
)

# --- <<< إزالة كاملة لخطوة المصادقة والتحقق من النطاق ---

# تهيئة متغير "المصادقة" في حالة الجلسة
if 'authenticated' not in st.session_state:
    # **التعديل هنا:** تعيين القيمة إلى True بشكل افتراضي لتجاوز التحقق
    st.session_state.authenticated = True 

# --- <<< هذا هو الجزء الذي تم تعديله (تمت إزالته) ---
# كود HTML/JavaScript الذي سيتم حقنه في Streamlit
# auth_listener_html = """
# <script>
# // --- الجزء الأول: الاستماع لرسالة المصادقة من الأب ---

# // النطاقات المسموحة (للرسالة القادمة)
# const allowedOrigins = [
#     "https://www.skydata.kozow.com",
#     "https://skydata.kozow.com"
# ];
# const allowedDomains = [
#     "www.skydata.kozow.com",
#     "skydata.kozow.com"
# ];

# // الاستماع للرسائل القادمة من النافذة "الأب"
# window.addEventListener('message', (event) => {
    
#     // التحقق إذا كان مصدر الرسالة ضمن القائمة المسموحة
#     if (!allowedOrigins.includes(event.origin)) {
#         console.warn("تم رفض رسالة من نطاق غير مصرح به:", event.origin);
#         return;
#     }

#     try {
#         const data = JSON.parse(event.data);
        
#         // التحقق من محتوى الرسالة (لزيادة الأمان)
#         if (data.source === 'skydata-auth' && allowedDomains.includes(data.domain)) {
#             // إذا نجحت المصادقة، أرسل "True" إلى Streamlit
#             Streamlit.setComponentValue({ "authenticated": true });
#         }
#     } catch (e) {
#         console.error("خطأ في معالجة الرسالة:", e);
#     }
# }, false);


# // --- الجزء الثاني (الجديد): إرسال رسالة "أنا جاهز" إلى الأب ---
# // (هذا الكود سيعمل فوراً عند تحميل هذا المكوّن)

# const readyMessage = JSON.stringify({ source: 'streamlit-ready' });

# // يجب أن نرسل الرسالة إلى النطاقات المسموحة للأب
# allowedOrigins.forEach(origin => {
#     try {
#         // إرسال الرسالة إلى النافذة "الأب"
#         window.parent.postMessage(readyMessage, origin);
#     } catch (e) {
#         console.error("فشل إرسال رسالة الاستعداد إلى الأب:", origin, e);
#     }
# });

# </script>
# """
# --- <<< نهاية الجزء الذي تم تعديله (تمت إزالته) ---

# تشغيل "المستمع" كعنصر HTML غير مرئي (تمت إزالته)
# auth_result = components.html(auth_listener_html, height=0, width=0)

# (الكود الذي أصلحناه سابقاً - تمت إزالته لعدم الحاجة)
# if auth_result and isinstance(auth_result, dict) and auth_result.get("authenticated") == True:
#     st.session_state.authenticated = True
# --- نهاية خطوة المصادقة (تمت إزالة الأجزاء غير الضرورية) ---


# --- التحقق الرئيسي: لا تقم بتشغيل أي شيء إلا بعد المصادقة ---
# **تم ترك هذا التحقق، لكن بما أن القيمة هي True افتراضياً، سيتم تشغيل التطبيق**
if st.session_state.authenticated:
    
    # --- 2. كود التصميم (CSS) ---
    CUSTOM_CSS = """
    <style>
        /* ... (نفس كود CSS بالكامل) ... */
        body, [data-testid="stAppViewContainer"], [data-testid="stHeader"] {
            background-color: #121212 !important; color: #e0e0e0 !important;
        }
        [data-testid="stHeader"] {
            background-color: #1e1e1e !important; border-bottom: 1px solid #333;
        }
        h1, h2, h3, h4, h5, h6 { color: #ffffff !important; }
        p, label, [data-testid="stMarkdownContainer"] { color: #e0e0e0 !important; }

        /* ... (باقي كود CSS الخاص بك) ... */
        [data-testid="stSidebar"][aria-expanded="true"] {
            background-color: #1e1e1e !important; border-right: 1px solid #333 !important;
        }
        [data-testid="stSidebar"][aria-expanded="true"] [data-testid="stRadio"] label { 
            background-color: #2a2a2a; padding: 10px; border-radius: 5px;
            margin: 5px 0; transition: all 0.3s ease;
        }
        [data-testid="stSidebar"][aria-expanded="true"] [data-testid="stRadio"] label:hover {
            background-color: #333; color: #007bff;
        }
        [data-testid="stSidebar"][aria-expanded="true"] [data-testid="stRadio"] input:checked + div {
            color: #007bff; 
        }
        [data-testid="stDownloadButton"] > button {
            background-color: #007bff !important; color: #ffffff !important;
            border: none !important; border-radius: 5px !important;
            font-weight: bold !important; padding: 12px 20px !important;
            transition: all 0.3s ease !important; width: 100%; 
        }
        [data-testid="stDownloadButton"] > button:hover {
            background-color: #0056b3 !important; transform: scale(1.03) !important;
        }
        [data-testid="stFileUploader"] section {
            background-color: #1e1e1e !important; border: 1px dashed #333 !important;
            border-radius: 8px !important;
        }
        [data-testid="stFileUploader"] section > button { 
            background-color: #007bff !important; color: #ffffff !important;
            border: none !important; border-radius: 5px !important;
            font-weight: bold !important; transition: all 0.3s ease !important;
        }
        [data-testid="stFileUploader"] section > button:hover {
            background-color: #0056b3 !important; transform: scale(1.03) !important;
        }
        [data-testid="stFileUploader"] label { color: #ffffff !important; }
        [data-testid="stFileUploader"] section [data-testid="stMarkdownContainer"] p {
            color: #aaa !important;
        }
        [data-testid="stImage"] figcaption { color: #aaa !important; }
        [data-testid="stSpinner"] > div { color: #e0e0e0 !important; }
        [data-testid="stTextInput"] input { 
            background-color: #1e1e1e !important; color: #e0e0e0 !important;
            border: 1px solid #333 !important; border-radius: 5px !important;
        }
    </style>
    """
    st.markdown(CUSTOM_CSS, unsafe_allow_html=True)

    # --- 3. قاموس الترجمة ---
    localization = {
        'en': {
            'brand_title': "", 'lang_select': "Select Language", 'input_method': "Select Input Method",
            'tab_upload': "Image Upload", 'tab_url': "URL Input", 'tab_file': "File Output",
            'upload_prompt': "Upload an image", 'upload_prompt_file': "Upload an image for file output",
            'url_prompt': "Paste an image URL", 'processed_caption': "Processed Image",
            'download_button': "Download PNG", 'error_fetching': "Error fetching image from URL:",
            'error_processing': "Error processing image:", 'generic_error': "An error occurred:",
            'page_title': "SkyData - Background Removal", 'spinner_text': "Processing image, please wait...",
        },
        'ar': {
            'brand_title': "", 'lang_select': "اختر اللغة", 'input_method': "اختر طريقة الإدخال",
            'tab_upload': "رفع صورة", 'tab_url': "رابط صورة", 'tab_file': "تصدير ملف",
            'upload_prompt': "ارفع صورة", 'upload_prompt_file': "ارفع صورة لتصدير الملف",
            'url_prompt': "الصق رابط الصورة", 'processed_caption': "الصورة المعالجة",
            'download_button': "تحميل بصيغة PNG", 'error_fetching': "خطأ في جلب الصورة من الرابط:",
            'error_processing': "خطأ في معالجة الصورة:", 'generic_error': "حدث خطأ:",
            'page_title': "SkyData - أداة إزالة الخلفية", 'spinner_text': "جاري معالجة الصورة، يرجى الانتظار...",
        }
    }

    # --- 4. تحميل النموذج ---
    torch.set_float32_matmul_precision(["high", "highest"][0])
    use_cuda = torch.cuda.is_available()
    device = "cuda" if use_cuda else "cpu"

    @st.cache_resource
    def load_model():
        model = AutoModelForImageSegmentation.from_pretrained("ZhengPeng7/BiRefNet", trust_remote_code=True)
        model.to(device)
        if use_cuda:
            model = model.half()
        return model

    birefnet = load_model()

    transform_image = transforms.Compose([
        transforms.Resize((1024, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])

    # --- 5. دالة المعالجة الرئيسية ---
    @st.cache_data
    def process(image):
        image_size = image.size
        input_images = transform_image(image).unsqueeze(0).to(device)
        if use_cuda:
            input_images = input_images.half()

        with torch.no_grad():
            preds = birefnet(input_images)[-1].sigmoid().cpu()
        
        pred = preds[0].squeeze()
        pred_pil = transforms.ToPILImage()(pred)
        mask = pred_pil.resize(image_size)
        image.putalpha(mask)
        
        img_bytes = io.BytesIO()
        image.save(img_bytes, format="PNG")
        img_bytes = img_bytes.getvalue()
        
        return image, img_bytes

    # --- 6. واجهة التطبيق (Sidebar) ---
    lang_choice = st.sidebar.radio(
        label="Select Language / اختر اللغة", 
        options=["English", "العربية"], 
        horizontal=True
    )
    lang_code = 'ar' if lang_choice == 'العربية' else 'en'
    loc = localization[lang_code]

    # --- 7. واجهة التطبيق الرئيسية (Main App) ---
    LOGO_URL = "https://i.ibb.co/v4vwvcGq/skydatafull.webp"
    col1, col2 = st.columns([1, 6]) 
    with col1:
        st.image(LOGO_URL, width=80) 
    with col2:
        st.title(loc['page_title'])

    st.sidebar.title(loc['brand_title'])
    selected_tab = st.sidebar.radio(loc['input_method'], [loc['tab_upload'], loc['tab_url'], loc['tab_file']])

    # --- 8. منطق التبويبات ---
    if selected_tab == loc['tab_upload']:
        uploaded_file = st.file_uploader(loc['upload_prompt'], type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            try:
                image = Image.open(uploaded_file).convert("RGB")
                with st.spinner(loc['spinner_text']):
                    processed_image, file_bytes = process(image)
                st.image(processed_image, caption=loc['processed_caption'])
                st.download_button(
                    label=loc['download_button'], data=file_bytes,
                    file_name=f"{uploaded_file.name.rsplit('.', 1)[0]}.png", mime="image/png",
                )
            except Exception as e:
                st.error(f"{loc['generic_error']} {e}")

    elif selected_tab == loc['tab_url']:
        image_url = st.text_input(loc['url_prompt'])
        if image_url:
            try:
                response = requests.get(image_url, stream=True)
                response.raise_for_status()
                image = Image.open(BytesIO(response.content)).convert("RGB")
                with st.spinner(loc['spinner_text']):
                    processed_image, file_bytes = process(image)
                st.image(processed_image, caption=loc['processed_caption'])
                try:
                    file_name = image_url.split('/')[-1].rsplit('.', 1)[0] + ".png"
                except Exception:
                    file_name = "processed_image.png"
                st.download_button(
                    label=loc['download_button'], data=file_bytes,
                    file_name=file_name, mime="image/png",
                )
            except requests.exceptions.RequestException as e:
                st.error(f"{loc['error_fetching']} {e}")
            except Exception as e:
                st.error(f"{loc['error_processing']} {e}")

    elif selected_tab == loc['tab_file']:
        uploaded_file = st.file_uploader(loc['upload_prompt_file'], type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            try:
                image = Image.open(uploaded_file).convert("RGB")
                with st.spinner(loc['spinner_text']):
                    processed_image, file_bytes = process(image)
                st.image(processed_image, caption=loc['processed_caption'])
                st.download_button(
                    label=loc['download_button'], data=file_bytes,
                    file_name=f"{uploaded_file.name.rsplit('.', 1)[0]}.png", mime="image/png",
                )
            except Exception as e:
                st.error(f"{loc['generic_error']} {e}")


# --- إذا فشلت المصادقة ---
# **تمت إزالة الكود الذي يعرض رسالة "الوصول مرفوض"**
else:
    pass 
    # st.error("🔒 الوصول مرفوض.")
    # st.error("لا يمكن تشغيل هذه الأداة إلا من خلال النطاق المعتمد: https://www.skydata.kozow.com/")
    # st.warning("Access Denied. This tool can only be run when embedded on https://www.skydata.kozow.com/")