import gradio as gr
import cv2
import numpy as np
import os
import zipfile
from pathlib import Path
import tempfile
import shutil
from typing import List, Tuple, Optional

class MangaWebSplitter:
    def __init__(self):
        self.temp_dir = None
        
    def detect_separators(self, image: np.ndarray, 
                         white_threshold: int = 240, 
                         black_threshold: int = 15,
                         min_separator_height: int = 15) -> List[int]:
        """
        Phát hiện các vùng separator (trắng hoặc đen) để cắt
        """
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()
            
        height, width = gray.shape
        split_points = []
        
        # Phân tích từng hàng pixel (kiểm tra mỗi 3 pixel để tăng tốc)
        for y in range(0, height, 3):
            row = gray[y, :]
            
            # Đếm pixel trắng và đen
            white_pixels = np.sum(row > white_threshold)
            black_pixels = np.sum(row < black_threshold)
            white_ratio = white_pixels / width
            black_ratio = black_pixels / width
            
            # Kiểm tra xem có phải là vùng separator không
            is_separator = white_ratio > 0.85 or black_ratio > 0.85
            
            if is_separator:
                # Kiểm tra chiều cao vùng separator
                separator_height = 0
                separator_type = "white" if white_ratio > black_ratio else "black"
                
                for check_y in range(y, min(y + 50, height)):
                    check_row = gray[check_y, :]
                    if separator_type == "white":
                        check_ratio = np.sum(check_row > white_threshold) / width
                    else:
                        check_ratio = np.sum(check_row < black_threshold) / width
                    
                    if check_ratio > 0.85:
                        separator_height += 1
                    else:
                        break
                
                # Nếu vùng separator đủ cao
                if separator_height >= min_separator_height:
                    split_points.append(y + separator_height // 2)
        
        # Loại bỏ các điểm quá gần nhau
        filtered_points = []
        for point in split_points:
            if not filtered_points or point - filtered_points[-1] > 100:
                filtered_points.append(point)
        
        return filtered_points
    
    def split_manga(self, 
                   image_file, 
                   max_height: int = None,
                   white_threshold: int = 240,
                   black_threshold: int = 15,
                   min_separator_height: int = 15,
                   show_preview: bool = False,
                   auto_height: bool = True) -> Tuple[str, str, str]:
        """
        Cắt ảnh manga và trả về kết quả
        """
        try:
            # Đọc ảnh
            image = cv2.imread(image_file)
            if image is None:
                return "❌ Không thể đọc file ảnh! Vui lòng kiểm tra định dạng.", "", ""
            
            height, width = image.shape[:2]
            
            # Tự động điều chỉnh chiều cao nếu bật auto mode
            if auto_height or max_height is None or max_height == 0:
                if height <= 2000:
                    calculated_height = max(height // 2, 800)
                elif height <= 4000:
                    calculated_height = max(height // 3, 1000)
                elif height <= 6000:
                    calculated_height = 2000
                elif height <= 10000:
                    calculated_height = 2500
                else:
                    calculated_height = height // (height // 2000)
                
                max_height = calculated_height
                auto_mode_text = f" (Tự động tối ưu)"
            else:
                auto_mode_text = " (Thủ công)"
            
            # Tìm điểm cắt
            split_points = self.detect_separators(
                image, white_threshold, black_threshold, min_separator_height
            )
            
            # Nếu không tìm thấy điểm cắt, dùng chiều cao đã tính
            if not split_points:
                split_points = list(range(max_height, height, max_height))
                info_msg = f"🔄 Không tìm thấy vùng separator, cắt theo chiều cao {max_height}px{auto_mode_text}"
            else:
                info_msg = f"✅ Tìm thấy {len(split_points)} vùng separator tự động{auto_mode_text}"
            
            # Thêm điểm đầu và cuối
            all_points = [0] + split_points + [height]
            all_points = sorted(list(set(all_points)))
            
            # Tạo thư mục tạm
            self.temp_dir = tempfile.mkdtemp()
            output_files = []
            base_name = Path(image_file).stem
            
            # Cắt ảnh
            valid_parts = 0
            for i in range(len(all_points) - 1):
                start_y = all_points[i]
                end_y = all_points[i + 1]
                
                # Bỏ qua phần quá nhỏ
                if end_y - start_y < 100:
                    continue
                
                valid_parts += 1
                cropped = image[start_y:end_y, :]
                
                # Lưu file
                output_file = os.path.join(self.temp_dir, f"{base_name}_part_{valid_parts:03d}.png")
                cv2.imwrite(output_file, cropped)
                output_files.append(output_file)
            
            # Tạo file zip
            zip_path = os.path.join(self.temp_dir, f"{base_name}_split.zip")
            with zipfile.ZipFile(zip_path, 'w') as zipf:
                for file in output_files:
                    zipf.write(file, os.path.basename(file))
            
            result_msg = f"""
📊 **Kết quả cắt ảnh:**
- Kích thước gốc: {width}×{height}px
- Chiều cao mỗi phần: {max_height}px{auto_mode_text}
- Số phần đã cắt: {valid_parts}
- Điểm cắt: {split_points if split_points else 'Theo chiều cao cố định'}

{info_msg}
            """.strip()
            
            # Tạo preview nếu yêu cầu
            preview_path = ""
            if show_preview and split_points:
                preview_image = image.copy()
                for point in split_points:
                    cv2.line(preview_image, (0, point), (width, point), (0, 0, 255), 3)
                preview_path = os.path.join(self.temp_dir, "preview.png")
                cv2.imwrite(preview_path, preview_image)
            
            return result_msg, zip_path, preview_path
            
        except Exception as e:
            return f"❌ Lỗi: {str(e)}", "", ""

# Khởi tạo splitter
splitter = MangaWebSplitter()

def process_manga(image_file, auto_height, max_height, white_threshold, black_threshold, min_separator_height, show_preview):
    """Xử lý ảnh manga"""
    if image_file is None:
        return "⚠️ Vui lòng upload ảnh trước!", None, None
    
    result_msg, zip_path, preview_path = splitter.split_manga(
        image_file, max_height, white_threshold, black_threshold, min_separator_height, show_preview, auto_height
    )
    
    return result_msg, zip_path if zip_path else None, preview_path if preview_path else None

# CSS tùy chỉnh
css = """
.container {
    max-width: 1200px;
    margin: 0 auto;
}
.header {
    text-align: center;
    background: linear-gradient(45deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 2rem;
    border-radius: 10px;
    margin-bottom: 2rem;
}
.upload-area {
    border: 2px dashed #667eea;
    border-radius: 10px;
    padding: 2rem;
    text-align: center;
    background: #f8f9ff;
}
"""

# Tạo giao diện Gradio
with gr.Blocks(css=css, title="🎨 Manga Splitter Pro") as app:
    
    gr.HTML("""
    <div class="header">
        <h1>🎨 MANGA SPLITTER PRO</h1>
        <p>Công cụ cắt ảnh manga thông minh - Tự động phát hiện vùng trắng & đen</p>
        <p>Tránh cắt trùng vào bóng thoại • Hỗ trợ OCR tối ưu</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML("<h3>📤 Upload & Cài đặt</h3>")
            
            # Upload file
            image_input = gr.File(
                label="🖼️ Chọn ảnh manga",
                file_types=[".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tiff"],
                type="filepath"
            )
            
            # Cài đặt
            with gr.Accordion("⚙️ Cài đặt nâng cao", open=False):
                auto_height = gr.Checkbox(
                    label="🤖 Tự động điều chỉnh chiều cao tối ưu",
                    value=True,
                    info="Tự động tính chiều cao phù hợp dựa trên kích thước ảnh"
                )
                
                max_height = gr.Slider(
                    minimum=500, maximum=5000, value=2000, step=100,
                    label="📏 Chiều cao tối đa mỗi phần (px)",
                    visible=False
                )
                
                white_threshold = gr.Slider(
                    minimum=200, maximum=255, value=240, step=5,
                    label="⚪ Ngưỡng pixel trắng"
                )
                
                black_threshold = gr.Slider(
                    minimum=0, maximum=50, value=15, step=5,
                    label="⚫ Ngưỡng pixel đen"
                )
                
                min_separator_height = gr.Slider(
                    minimum=5, maximum=50, value=15, step=5,
                    label="📐 Chiều cao tối thiểu vùng separator"
                )
                
                show_preview = gr.Checkbox(
                    label="👁️ Hiển thị preview điểm cắt",
                    value=True
                )
            
            # Nút xử lý
            process_btn = gr.Button(
                "🚀 Cắt ảnh manga", 
                variant="primary",
                size="lg"
            )
        
        with gr.Column(scale=1):
            gr.HTML("<h3>📊 Kết quả</h3>")
            
            # Kết quả
            result_text = gr.Markdown(
                value="👋 Chào mừng! Upload ảnh manga để bắt đầu cắt.",
                label="📈 Thông tin"
            )
            
            # Download
            download_file = gr.File(
                label="💾 Tải xuống file ZIP",
                visible=False
            )
            
            # Preview
            preview_image = gr.Image(
                label="👁️ Preview điểm cắt",
                visible=False
            )
    
    # Thông tin hướng dẫn
    with gr.Accordion("📚 Hướng dẫn sử dụng", open=False):
        gr.Markdown("""
        ### 🎯 Cách sử dụng:
        1. **Upload ảnh** manga (JPG, PNG) vào khung bên trái
        2. **Điều chỉnh cài đặt** nếu cần (mặc định đã tối ưu)
        3. **Bấm "Cắt ảnh manga"** để xử lý
        4. **Tải xuống** file ZIP chứa các phần đã cắt
        
        ### ⚙️ Tham số:
        - **Chiều cao tối đa**: Giới hạn chiều cao mỗi phần cắt
        - **Ngưỡng trắng**: Độ trắng để nhận diện vùng separator (240 = rất trắng)
        - **Ngưỡng đen**: Độ đen để nhận diện vùng dramatic (15 = rất đen) 
        - **Chiều cao separator**: Vùng separator phải cao ít nhất bao nhiêu pixel
        
        ### 🎨 Hoạt động với:
        - ✅ Manga truyền thống (nền trắng)
        - ✅ Action manga (vùng đen dramatic)
        - ✅ Webtoon dài
        - ✅ Raw scan chất lượng cao
        """)
    
    # Xử lý sự kiện
    def toggle_manual_height(auto_mode):
        return gr.update(visible=not auto_mode, value=2000)
    
    auto_height.change(
        toggle_manual_height,
        inputs=[auto_height],
        outputs=[max_height]
    )
    
    process_btn.click(
        fn=process_manga,
        inputs=[image_input, auto_height, max_height, white_threshold, black_threshold, min_separator_height, show_preview],
        outputs=[result_text, download_file, preview_image]
    ).then(
        # Hiển thị kết quả
        lambda zip_file, preview_img: (
            gr.update(visible=zip_file is not None), 
            gr.update(visible=preview_img is not None)
        ),
        inputs=[download_file, preview_image],
        outputs=[download_file, preview_image]
    )

if __name__ == "__main__":
    app.launch(
        server_name="0.0.0.0",  # Cho phép truy cập từ mạng local
        server_port=7860,
        share=False,  # Đặt True nếu muốn public link
        debug=True
    )