import threading
import gradio as gr
import os

# ==========================================
# 🌐 1. تشغيل واجهة Gradio لإبقاء السيرفر نَشِطاً
# ==========================================
def start_gradio_ui():
    demo = gr.Interface(
        fn=lambda x: "Bot is running online 24/7!",
        inputs="text",
        outputs="text",
        title="Manga Cleaner Discord Bot Status"
    )
    demo.launch(server_name="0.0.0.0", server_port=7860)

threading.Thread(target=start_gradio_ui, daemon=True).start()

# ==========================================
# 🛡️ 2. كود تخطي حماية PyTorch 2.6
# ==========================================
import torch
original_load = torch.load
def safe_load(*args, **kwargs):
    kwargs['weights_only'] = False 
    return original_load(*args, **kwargs)
torch.load = safe_load

# ==========================================
# 📦 3. استيراد باقي المكتبات الأساسية
# ==========================================
import nest_asyncio
nest_asyncio.apply()

import re
import gc
import time
import asyncio
import shutil
import datetime
import zipfile
import aiohttp
import requests
import discord
from discord import app_commands
from discord.ext import commands
import cv2
import numpy as np
from PIL import Image, ImageFile
from ultralytics import YOLO
from simple_lama_inpainting import SimpleLama
from bs4 import BeautifulSoup
import patoolib

ImageFile.LOAD_TRUNCATED_IMAGES = True

# ==========================================
# 📥 4. تنزيل وإعداد نماذج الذكاء الاصطناعي
# ==========================================
intents = discord.Intents.default()
intents.message_content = True
bot = commands.Bot(command_prefix="!", intents=intents)

# ==========================================
# 📥 تحميل نماذج الذكاء الاصطناعي على الـ CPU بأمان
# ==========================================
import torch

manga_model = None
lama = None
MODEL_PATH = "manga_model.pt"

HF_MODEL_URL = "https://huggingface.co/ogkalu/comic-text-segmenter-yolov8m/resolve/main/comic-text-segmenter.pt"

print("⏳ Checking & Loading Models on CPU...")
if not os.path.exists(MODEL_PATH):
    print("📥 Downloading model...")
    os.system(f'wget -O {MODEL_PATH} "{HF_MODEL_URL}"')

try:
    manga_model = YOLO(MODEL_PATH)
    print("✅ YOLO Manga Model loaded!")
except Exception as e:
    print(f"❌ Error loading YOLO: {e}")

try:
    # اجبار Lama على العمل باستخدام المعالج CPU حصرياً
    lama = SimpleLama(device="cpu")
    print("✅ Inpainting Model loaded on CPU successfully!")
except Exception as e:
    print(f"❌ Error loading Inpainting model: {e}")

# ==========================================
# 📥 5. دوال التحميل والتنزيل الذكي الشاملة
# ==========================================
def ultimate_sort_key(filepath):
    filename = os.path.basename(filepath)
    name_without_ext = os.path.splitext(filename)[0]
    return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', name_without_ext)]

class DownloadManager:
    @staticmethod
    def get_drive_real_name(url):
        try:
            r = requests.get(url, timeout=10)
            soup = BeautifulSoup(r.text, 'html.parser')
            title = soup.title.string
            if title:
                clean_title = title.replace(" - Google Drive", "").strip()
                if clean_title and clean_title != "Google Drive":
                    return clean_title
            return "Downloaded_Chapter"
        except:
            return "Downloaded_Chapter"

    @staticmethod
    def extract_archives(folder):
        for root, _, files in os.walk(folder):
            for file in files:
                fp = os.path.join(root, file)
                try:
                    if file.lower().endswith(".zip"):
                        with zipfile.ZipFile(fp, 'r') as z: z.extractall(root)
                        os.remove(fp)
                    elif file.lower().endswith((".rar", ".cbr", ".cbz")):
                        patoolib.extract_archive(fp, outdir=root, verbosity=-1)
                        os.remove(fp)
                except: pass

    @staticmethod
    def _scrape_web_selenium_sync(url, output_folder):
        from selenium import webdriver
        from selenium.webdriver.chrome.options import Options
        from selenium.webdriver.chrome.service import Service
        from webdriver_manager.chrome import ChromeDriverManager
        from selenium.webdriver.common.by import By
        from io import BytesIO

        chrome_options = Options()
        chrome_options.add_argument("--headless=new")
        chrome_options.add_argument("--no-sandbox")
        chrome_options.add_argument("--disable-dev-shm-usage")
        
        try:
            service = Service(ChromeDriverManager().install())
            driver = webdriver.Chrome(service=service, options=chrome_options)
        except Exception as e:
            print(f"Selenium Error: {e}")
            return False

        try:
            driver.get(url)
            time.sleep(3)
            
            last_height = driver.execute_script("return document.body.scrollHeight")
            for _ in range(20):
                driver.execute_script("window.scrollBy(0, 800);")
                time.sleep(0.5)
                new_h = driver.execute_script("return window.pageYOffset + window.innerHeight")
                if new_h >= last_height:
                    break
                last_height = driver.execute_script("return document.body.scrollHeight")
            time.sleep(2)

            elements = driver.find_elements(By.TAG_NAME, 'img')
            urls = [e.get_attribute('src') or e.get_attribute('data-src') for e in elements]
            urls = list(set([u for u in urls if u and u.startswith('http') and not any(x in u.lower() for x in ['logo', 'icon', 'avatar'])]))

            req_headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64)'}
            downloaded_count = 0
            
            for idx, u in enumerate(urls):
                try:
                    r = requests.get(u, headers=req_headers, timeout=10).content
                    im = Image.open(BytesIO(r)).convert('RGB')
                    if im.width > 150 and im.height > 150:
                        path = os.path.join(output_folder, f'page_{idx:03d}.jpg')
                        im.save(path)
                        downloaded_count += 1
                except:
                    continue

            if downloaded_count == 0: return False

            title = driver.title.split('|')[0].strip() if driver.title else "Web_Chapter"
            title = re.sub(r'[\\/*?:"<>|]', "", title)
            return title
        except Exception as e:
            print(f"Scraping Error: {e}")
            return False
        finally:
            try:
                driver.quit()
            except: pass

    @staticmethod
    async def scrape_web_selenium_async(url, output_folder):
        loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(None, DownloadManager._scrape_web_selenium_sync, url, output_folder)
        return result

    @staticmethod
    async def download_gdown(url, input_dir):
        # يفضل وضع الـ API Key الخاص بك هنا لو أردت تحميل سريع، أو الاعتماد على الـ gdown العادي
        import gdown
        loop = asyncio.get_running_loop()
        
        if "drive.google.com" in url:
            if "folder" in url or "folders" in url:
                await loop.run_in_executor(None, lambda: gdown.download_folder(url=url, output=input_dir, quiet=False, use_cookies=False))
                return True
            else:
                file_id = None
                match = re.search(r"([-\w]{25,})", url)
                if match: file_id = match.group(1)
                
                if file_id:
                    target_path = os.path.join(input_dir, "downloaded.zip")
                    await loop.run_in_executor(None, lambda: gdown.download(id=file_id, output=target_path, quiet=False))
                else:
                    target_path = os.path.join(input_dir, "downloaded.zip")
                    await loop.run_in_executor(None, lambda: gdown.download(url=url, output=target_path, quiet=False, fuzzy=True))
                return True
        else:
            target_path = os.path.join(input_dir, "downloaded_file")
            async with aiohttp.ClientSession() as session:
                async with session.get(url) as response:
                    if response.status == 200:
                        with open(target_path, 'wb') as f:
                            f.write(await response.read())
                        return True
        return False

# ==========================================
# 🖌️ 6. المعالجة الآمنة والمتزامنة مع كارت الشاشة
# ==========================================
def get_combined_mask(results, img_shape):
    mask = np.zeros((img_shape[0], img_shape[1]), dtype=np.uint8)
    if hasattr(results[0], 'masks') and results[0].masks is not None:
        for m in results[0].masks.xy:
            contour = np.array(m, dtype=np.int32)
            cv2.drawContours(mask, [contour], -1, 255, -1) 
    if hasattr(results[0], 'boxes') and results[0].boxes is not None:
        for box in results[0].boxes.xyxy:
            x1, y1, x2, y2 = map(int, box[:4])
            cv2.rectangle(mask, (x1, y1), (x2, y2), 255, -1)
    return mask

def process_slice(slice_cv_img):
    img_h, img_w = slice_cv_img.shape[:2]
    results = manga_model.predict(source=slice_cv_img, imgsz=1024, conf=0.15, verbose=False)
    
    bubble_count = len(results[0].boxes) if hasattr(results[0], 'boxes') and results[0].boxes is not None else 0
    if bubble_count == 0: return slice_cv_img, 0

    mask_np = get_combined_mask(results, (img_h, img_w))
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (9, 9)) 
    dilated_mask_np = cv2.dilate(mask_np, kernel, iterations=4)
    
    slice_pil = Image.fromarray(cv2.cvtColor(slice_cv_img, cv2.COLOR_BGR2RGB))
    mask_pil = Image.fromarray(dilated_mask_np).convert("L")
    
    cleaned_pil = lama(slice_pil, mask_pil)
    cleaned_cv = cv2.cvtColor(np.array(cleaned_pil), cv2.COLOR_RGB2BGR)
    return cleaned_cv, bubble_count

async def process_single_image_concurrently(input_path, output_path, gpu_semaphore):
    async with gpu_semaphore:
        original_img_pil = Image.open(input_path).convert("RGB")
        W, H = original_img_pil.size
        num_slices = max(1, H // 2000)
        slice_h = H // num_slices
        overlap = 150 
        
        cv_original = cv2.cvtColor(np.array(original_img_pil), cv2.COLOR_RGB2BGR)
        loop = asyncio.get_running_loop()

        tasks = []
        slice_coords = []
        
        for i in range(num_slices):
            start_y = max(0, i * slice_h - (overlap if i > 0 else 0))
            end_y = min(H, (i + 1) * slice_h + (overlap if i < num_slices - 1 else 0))
            slice_cv_img = cv_original[start_y:end_y, 0:W]
            slice_coords.append((start_y, end_y))
            
            task = loop.run_in_executor(None, process_slice, slice_cv_img)
            tasks.append(task)

        results = await asyncio.gather(*tasks)
        
        cleaned_slices_data = []
        total_bubbles = 0
        
        for idx, (cleaned_slice_cv, bubbles) in enumerate(results):
            total_bubbles += bubbles
            cleaned_slices_data.append((cleaned_slice_cv, slice_coords[idx][0], slice_coords[idx][1]))

        final_img = Image.new("RGB", (W, H))
        for i, (cleaned_cv_slice, start_y, end_y) in enumerate(cleaned_slices_data):
            cleaned_pil_slice = Image.fromarray(cv2.cvtColor(cleaned_cv_slice, cv2.COLOR_BGR2RGB))
            if i == 0:
                final_img.paste(cleaned_pil_slice.crop((0, 0, W, cleaned_pil_slice.height - overlap // 2)), (0, 0))
            elif i == num_slices - 1:
                final_img.paste(cleaned_pil_slice.crop((0, overlap // 2, W, cleaned_pil_slice.height)), (0, start_y + overlap // 2))
            else:
                final_img.paste(cleaned_pil_slice.crop((0, overlap // 2, W, cleaned_pil_slice.height - overlap // 2)), (0, start_y + overlap // 2))
                
        final_img.save(output_path, quality=95, subsampling=0)
        
        del original_img_pil, cv_original, final_img
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        return total_bubbles

# ==========================================
# 📊 7. واجهة ديسكورد
# ==========================================
class ProgressEmbed(discord.Embed):
    def __init__(self, filename, start_time, user: discord.Member):
        super().__init__(title="⏳ Cleaning Webtoon...", color=discord.Color.dark_theme())
        self.filename = filename
        self.start_time = start_time
        
        self.add_field(name="📄 File/Folder", value=f"`{filename}`", inline=True)
        self.add_field(name="⏱️ Elapsed", value="`0.0s`", inline=True)
        self.add_field(name="💭 Elements Detected", value="`0 Cleaned`", inline=False)
        self.add_field(name="📈 Progress", value=self._make_bar(0), inline=False)
        
        current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M")
        self.set_footer(text=f"Requested by {user.display_name} • {current_time}")

    def _make_bar(self, percent):
        size = 15
        filled = int(size * percent / 100)
        bar = "▓" * filled + "░" * (size - filled)
        return f"`[{bar}]` **{percent}%**"

    def update_progress(self, percent, bubbles, status="Processing..."):
        elapsed = round(time.time() - self.start_time, 1)
        self.set_field_at(1, name="⏱️ Elapsed", value=f"`{elapsed}s`", inline=True)
        self.set_field_at(2, name="💭 Elements Detected", value=f"`{bubbles} Cleaned`", inline=False)
        self.set_field_at(3, name="📈 Progress", value=f"{self._make_bar(percent)} ({status})", inline=False)

# ==========================================
# 🤖 8. أمر البوت الشامل
# ==========================================
@bot.event
async def on_ready():
    print(f"🚀 Bot {bot.user} is online and running on Hugging Face!")
    await bot.tree.sync()

@bot.tree.command(name="clean", description="Clean Manhwa/Webtoon from Drive, Web, or Attachment")
@app_commands.describe(
    image="Attach an image file (Optional)",
    url="Provide a direct URL, Web Chapter link, or Google Drive link (Optional)"
)
async def clean_manhwa(interaction: discord.Interaction, image: discord.Attachment = None, url: str = None):
    if manga_model is None or lama is None:
        return await interaction.response.send_message("❌ Models are not loaded.", ephemeral=True)
    if not image and not url:
        return await interaction.response.send_message("❌ Provide an image or a URL!", ephemeral=True)

    start_time = time.time()
    await interaction.response.defer(thinking=True)
    
    chapter_name = "Downloaded_Chapter"
    
    if url:
        if "drive.google.com" in url:
            chapter_name = DownloadManager.get_drive_real_name(url)
        else:
            chapter_name = "Web_Chapter"
    elif image:
        chapter_name = os.path.splitext(image.filename)[0]
        
    safe_chapter_name = "".join(x for x in chapter_name if x.isalnum() or x in " -_")
    
    embed = ProgressEmbed(safe_chapter_name, start_time, interaction.user)
    progress_msg = await interaction.followup.send(embed=embed)

    temp_dir = f"./temp_{int(start_time)}"
    input_dir = os.path.join(temp_dir, "input")
    output_dir = os.path.join(temp_dir, "output")
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)

    try:
        embed.update_progress(5, 0, "Downloading files...")
        await progress_msg.edit(embed=embed)
        
        if image:
            await image.save(os.path.join(input_dir, image.filename))
        elif url:
            if "drive.google.com" in url:
                success = await DownloadManager.download_gdown(url, input_dir)
                if not success: raise Exception("Failed to download from Google Drive.")
            else:
                web_title = await DownloadManager.scrape_web_selenium_async(url, input_dir)
                if web_title:
                    safe_chapter_name = "".join(x for x in web_title if x.isalnum() or x in " -_")
                else:
                    raise Exception("Failed to scrape images from the provided URL.")

            DownloadManager.extract_archives(input_dir)

        images_to_process = []
        for root, _, files in os.walk(input_dir):
            for file in files:
                if file.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
                    images_to_process.append(os.path.join(root, file))

        images_to_process.sort(key=ultimate_sort_key)
        if not images_to_process: raise Exception("No valid image files found in chapter.")

        total_files = len(images_to_process)
        
        embed.update_progress(10, 0, f"Processing {total_files} pages...")
        await progress_msg.edit(embed=embed)

        # في Hugging Face الـ CPU قوي بس برضه هنخليه يعالج 4 صفحات بالتزامن
        gpu_semaphore = asyncio.Semaphore(4) 
        
        processing_tasks = []
        for img_path in images_to_process:
            base_name = os.path.basename(img_path)
            out_path = os.path.join(output_dir, f"cleaned_{base_name}")
            task = process_single_image_concurrently(img_path, out_path, gpu_semaphore)
            processing_tasks.append(task)

        results = await asyncio.gather(*processing_tasks)
        total_bubbles = sum(results)

        embed.update_progress(95, total_bubbles, "Zipping files...")
        await progress_msg.edit(embed=embed)

        elapsed = round(time.time() - start_time, 1)
        success_embed = discord.Embed(
            title="✅ Cleaning Complete!",
            description=f"Processed **{total_files}** pages in **{elapsed}s**.",
            color=discord.Color.green()
        )
        success_embed.add_field(name="💭 Cleaned Elements", value=f"`{total_bubbles} Items`", inline=True)
        
        zip_filename = f"{safe_chapter_name}_cleaned.zip"
        zip_path = os.path.join(temp_dir, zip_filename)
        
        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            for root, _, files in os.walk(output_dir):
                for file in files:
                    zipf.write(os.path.join(root, file), file)

        success_embed.add_field(name="🔗 Download", value=f"File: `{zip_filename}`", inline=False)
        
        if os.path.getsize(zip_path) / (1024 * 1024) < 25:
            await interaction.followup.send(embed=success_embed, file=discord.File(zip_path))
        else:
            success_embed.description += "\n⚠️ File is too large for Discord (>25MB)."
            await interaction.followup.send(embed=success_embed)

        await progress_msg.delete()

    except Exception as e:
        error_embed = discord.Embed(title="❌ Error", description=str(e), color=discord.Color.red())
        await interaction.followup.send(embed=error_embed, ephemeral=True)
        try: await progress_msg.delete()
        except: pass
    finally:
        if os.path.exists(temp_dir): shutil.rmtree(temp_dir)

# ==========================================
# ▶️ 9. تشغيل البوت الأساسي
# ==========================================
BOT_TOKEN = os.environ.get("BOT_TOKEN")

if __name__ == "__main__":
    if not BOT_TOKEN:
        print("❌ Error: BOT_TOKEN is missing in Secrets!")
    else:
        bot.run(BOT_TOKEN)