# !pip install google-genai pillow scipy

import io
import base64
import colorsys
from google import genai
from PIL import Image, ImageFilter, ImageMorph
import numpy as np

# Khởi tạo Gemini client
client = genai.Client()

# Model sử dụng để tạo ảnh
MODEL_ID = "gemini-3-pro-image-preview"

def decode_image(base64_data: str) -> Image.Image:
    """Giải mã dữ liệu ảnh base64 sang đối tượng PIL Image."""
    image_bytes = base64.b64decode(base64_data)
    return Image.open(io.BytesIO(image_bytes))


def rgb_to_hsv_array(rgb_array: np.ndarray) -> np.ndarray:
    """Chuyển đổi mảng RGB sang mảng HSV một cách hiệu quả."""
    # Chuẩn hóa RGB về khoảng 0-1
    rgb_normalized = rgb_array.astype(np.float32) / 255.0

    r, g, b = rgb_normalized[:, :, 0], rgb_normalized[:, :, 1], rgb_normalized[:, :, 2]

    max_c = np.maximum(np.maximum(r, g), b)
    min_c = np.minimum(np.minimum(r, g), b)
    delta = max_c - min_c

    # Tính toán Hue (Sắc độ)
    h = np.zeros_like(max_c)

    # Khi max == r
    mask_r = (max_c == r) & (delta != 0)
    h[mask_r] = (60 * ((g[mask_r] - b[mask_r]) / delta[mask_r]) + 360) % 360

    # Khi max == g
    mask_g = (max_c == g) & (delta != 0)
    h[mask_g] = (60 * ((b[mask_g] - r[mask_g]) / delta[mask_g]) + 120)

    # Khi max == b
    mask_b = (max_c == b) & (delta != 0)
    h[mask_b] = (60 * ((r[mask_b] - g[mask_b]) / delta[mask_b]) + 240)

    # Tính toán Saturation (Độ bão hòa)
    s = np.zeros_like(max_c)
    s[max_c != 0] = delta[max_c != 0] / max_c[max_c != 0]

    # Value (Giá trị độ sáng) chính là max_c
    v = max_c

    return np.stack([h, s * 100, v * 100], axis=-1)


def remove_green_screen_hsv(
    image: Image.Image,
    hue_center: float = 120,    # Tâm màu xanh lá trong HSV thường là 120 độ
    hue_range: float = 25,     # Khoảng lệch sắc độ
    min_saturation: float = 75, # Độ bão hòa tối thiểu
    min_value: float = 70,      # Độ sáng tối thiểu
    dilation_iterations: int = 2,
    erosion_iterations: int = 0
) -> Image.Image:
    """
    Xóa phông xanh sử dụng không gian màu HSV để nhận diện tốt hơn.

    HSV vượt trội trong việc nhận diện dải màu vì nó tách biệt
    sắc độ (hue) khỏi độ bão hòa (saturation) và giá trị độ sáng (value).
    """
    # Chuyển sang RGBA nếu chưa đúng định dạng
    if image.mode != 'RGBA':
        image = image.convert('RGBA')

    # Chuyển sang mảng numpy
    data = np.array(image)
    rgb = data[:, :, :3]

    # Chuyển sang HSV
    hsv = rgb_to_hsv_array(rgb)
    h, s, v = hsv[:, :, 0], hsv[:, :, 1], hsv[:, :, 2]

    # Tính toán khoảng cách sắc độ (tính đến tính chất vòng tròn của màu sắc)
    hue_diff = np.abs(h - hue_center)
    hue_diff = np.minimum(hue_diff, 360 - hue_diff)

    # Tạo mask (mặt nạ) cho các điểm ảnh màu xanh
    # Là màu xanh nếu: sắc độ nằm trong khoảng VÀ độ bão hòa đủ cao VÀ độ sáng đủ cao
    green_mask = (
        (hue_diff < hue_range) &
        (s > min_saturation) &
        (v > min_value)
    )

    # Áp dụng làm sạch hình thái học để loại bỏ lỗi răng cưa ở viền
    if dilation_iterations > 0 or erosion_iterations > 0:
        from scipy import ndimage

        # Giãn nở (Dilate) mask để bắt các điểm ảnh bị mờ ở viền (anti-aliased)
        if dilation_iterations > 0:
            green_mask = ndimage.binary_dilation(green_mask, iterations=dilation_iterations)

        # Tùy chọn xói mòn (Erode) lại (giúp loại bỏ nhiễu cô lập)
        if erosion_iterations > 0:
            green_mask = ndimage.binary_erosion(green_mask, iterations=erosion_iterations)

    # Gán độ trong suốt cho các điểm ảnh màu xanh
    alpha = data[:, :, 3].copy()
    alpha[green_mask] = 0
    data[:, :, 3] = alpha

    return Image.fromarray(data)


def remove_green_screen_aggressive(
    image: Image.Image,
    green_threshold: float = 1.2,
    edge_pixels: int = 0  # Đặt là 0 để tránh lấn vào viền trắng của sticker
) -> Image.Image:
    """
    Xóa xanh kiểu "mạnh tay": nhận diện bất kỳ điểm ảnh nào có màu xanh chiếm ưu thế.

    Cách này bắt được cả các màu xanh đậm/nhạt hơn, hoặc bóng đổ có ám xanh.
    """
    if image.mode != 'RGBA':
        image = image.convert('RGBA')

    data = np.array(image)
    r, g, b = data[:, :, 0].astype(float), data[:, :, 1].astype(float), data[:, :, 2].astype(float)

    # Một pixel được coi là "xanh" nếu kênh Green vượt trội đáng kể so với Red và Blue
    rb_max = np.maximum(r, b) + 1  # +1 để tránh lỗi chia cho 0
    green_ratio = g / rb_max

    # Kiểm tra thêm xem Green có phải là kênh màu chủ đạo không
    green_dominant = (g > r) & (g > b)

    # Kết hợp các điều kiện mặt nạ
    green_mask = (green_ratio > green_threshold) & green_dominant

    # Mở rộng mặt nạ để bắt các điểm ảnh ở rìa
    if edge_pixels > 0:
        from scipy import ndimage
        green_mask = ndimage.binary_dilation(green_mask, iterations=edge_pixels)

    # Áp dụng độ trong suốt
    alpha = data[:, :, 3].copy()
    alpha[green_mask] = 0
    data[:, :, 3] = alpha

    return Image.fromarray(data)


def cleanup_edges(image: Image.Image, threshold: int = 128) -> Image.Image:
    """
    Làm sạch các điểm ảnh bán trong suốt ở viền.

    Giúp loại bỏ hiệu ứng "quầng" (halo) từ các cạnh được khử răng cưa.
    """
    if image.mode != 'RGBA':
        return image

    data = np.array(image)
    alpha = data[:, :, 3]

    # Biến các điểm ảnh bán trong suốt thành trong suốt hoàn toàn hoặc đặc hoàn toàn
    alpha[alpha < threshold] = 0
    alpha[alpha >= threshold] = 255

    data[:, :, 3] = alpha
    return Image.fromarray(data)


def save_transparent_png(image: Image.Image, filename: str):
    """Lưu ảnh dưới dạng PNG và giữ nguyên độ trong suốt."""
    if image.mode != 'RGBA':
        image = image.convert('RGBA')
    image.save(filename, 'PNG')
    print(f"✅ Đã lưu: {filename}")

def load_image_as_content(image_path: str) -> dict:
    """
    Tải ảnh từ đường dẫn tệp và trả về dưới dạng khối nội dung (content block) cho API.
    """
    import os
    import mimetypes
    
    # Xác định loại mime từ phần mở rộng tệp
    mime_type, _ = mimetypes.guess_type(image_path)
    if mime_type is None:
        # Mặc định là JPEG nếu không xác định được
        mime_type = "image/jpeg"
    
    # Đọc và mã hóa ảnh sang base64
    with open(image_path, "rb") as f:
        image_data = base64.b64encode(f.read()).decode("utf-8")
    
    return {
        "type": "image",
        "data": image_data,
        "mime_type": mime_type
    }
 
 
def generate_sticker(
    prompt: str, 
    aspect_ratio: str = "1:1",
    image_size: str = "2K",
    input_images: list[str] | None = None
) -> Image.Image:
    """
    Tạo hình ảnh phong cách sticker với nền xanh chromakey.
    """
    # Prompt đã được tối ưu hóa để tách nền chromakey
    enhanced_prompt = f"""Hãy tạo một hình minh họa dạng sticker về: {prompt}
 
CÁC YÊU CẦU QUAN TRỌNG VỀ CHROMAKEY:
1. NỀN (BACKGROUND): Màu xanh chromakey đồng nhất, phẳng. Sử dụng CHÍNH XÁC mã màu hex #00FF00 (RGB 0, 255, 0). 
   Toàn bộ nền phải là một màu xanh thuần khiết duy nhất, KHÔNG biến đổi, KHÔNG gradient, KHÔNG đổ bóng, KHÔNG hiệu ứng ánh sáng.
 
2. VIỀN TRẮNG: Chủ thể PHẢI có một đường viền/mép màu trắng sạch sẽ (rộng 2-3 pixel) ngăn cách nó với nền xanh.
   Viền trắng này giúp ngăn hiện tượng lem màu giữa chủ thể và nền.
 
3. KHÔNG CÓ MÀU XANH TRÊN CHỦ THỂ: Bản thân chủ thể KHÔNG nên chứa bất kỳ màu xanh lá nào để tránh nhầm lẫn với phông nền.
   Nếu chủ thể bắt buộc phải có màu xanh (như lá cây), hãy sử dụng một sắc độ khác biệt rõ rệt như xanh rừng đậm hoặc xanh mòng két (teal).
 
4. CẠNH SẮC NÉT: Chủ thể cần có các cạnh sắc nét, rõ ràng - không có ranh giới mềm hoặc mờ.
 
5. CĂN GIỮA: Chủ thể nên được đặt ở chính giữa với khoảng đệm (padding) ở tất cả các phía.
 
6. PHONG CÁCH: Rực rỡ, sạch sẽ, phong cách sticker hoạt hình/minh họa với màu sắc đậm.
 
Đây là để trích xuất chromakey - nền xanh sẽ được loại bỏ bằng lập trình."""
 
    print(f"🎨 Đang tạo sticker: {prompt}")
    print(f"   Độ phân giải: {image_size}")
    
    # Xây dựng nội dung đầu vào
    # Khi có input_images, tạo danh sách các khối nội dung ảnh trước, sau đó đến văn bản
    if input_images:
        print(f"   Ảnh đầu vào: {len(input_images)} ảnh")
        input_content = []
        for img_path in input_images:
            print(f"   - Đang tải: {img_path}")
            input_content.append(load_image_as_content(img_path))
        # Thêm prompt văn bản vào khối nội dung cuối cùng
        input_content.append({"type": "text", "text": enhanced_prompt})
    else:
        # Nếu không có ảnh đầu vào, chỉ sử dụng trực tiếp prompt văn bản
        input_content = enhanced_prompt
 
    # Gọi Gemini Interactions API
    interaction = client.interactions.create(
        model=MODEL_ID,
        input=input_content,
        generation_config={
            "image_config": {
                "aspect_ratio": aspect_ratio,
                "image_size": image_size  # Sử dụng độ phân giải cao để có đường viền tốt hơn
            }
        }
    )
 
    # Trích xuất hình ảnh được tạo ra
    for output in interaction.outputs:
        if output.type == "image":
            print(f"✅ Đã tạo xong ảnh (mime_type: {output.mime_type})")
            return decode_image(output.data)

def create_sticker(
    prompt: str, 
    output_filename: str,
    aspect_ratio: str = "1:1",
    image_size: str = "2K",
    save_raw: bool = False,
    input_images: list[str] | None = None
) -> Image.Image:
    """
    Quy trình hoàn chỉnh để tạo một sticker trong suốt.
 
    Sử dụng phương pháp tiếp cận đa bước:
    1. Tạo ảnh với prompt chromakey đã tối ưu.
    2. Xóa xanh dựa trên HSV (bắt đúng dải màu).
    3. Xóa xanh chuyên sâu (loại bỏ các sắc xanh còn sót lại).
    4. Làm sạch viền để loại bỏ hiệu ứng quầng (halo).
    """
    import os
 
    # Bước 1: Tạo hình ảnh với phông xanh
    raw_image = generate_sticker(prompt, aspect_ratio, image_size, input_images)
 
    # Tùy chọn lưu ảnh gốc để kiểm tra (debug)
    if save_raw:
        raw_filename = output_filename.replace('.png', '_raw.png')
        raw_image.save(raw_filename)
        print(f"📸 Đã lưu ảnh gốc: {raw_filename}")
 
    # Bước 2: Xóa phông xanh dựa trên HSV
    print("🔧 Bước 1: Đang xóa phông xanh bằng HSV...")
    transparent_image = remove_green_screen_hsv(
        raw_image,
        hue_center=120,       # Sắc độ xanh lá thuần túy
        hue_range=25,         # Khoảng hẹp quanh màu xanh thuần
        min_saturation=75,    # Chỉ chọn màu xanh có độ bão hòa cao (giữ lại các màu xanh khác trên logo)
        min_value=70,         # Chỉ chọn màu xanh sáng
        dilation_iterations=2,  # Bắt các điểm ảnh bị mờ (anti-aliased) ở viền
        erosion_iterations=0
    )
 
    # Bước 3: Bỏ qua xóa chuyên sâu (vô hiệu hóa - vì có thể gây đốm trên chủ thể)
    # transparent_image = remove_green_screen_aggressive(...)
 
    # Bước 4: Làm sạch các lỗi răng cưa bán trong suốt ở viền
    print("✨ Đang làm sạch các đường viền...")
    transparent_image = cleanup_edges(transparent_image, threshold=64)
 
    # Bước 5: Lưu dưới dạng PNG
    save_transparent_png(transparent_image, output_filename)
 
    return transparent_image

# Định nghĩa câu lệnh mô tả sticker
prompt = "một chú mèo con đáng yêu, hạnh phúc với đôi mắt to tròn"
 
# Gọi hàm create_sticker để thực hiện quy trình tự động
sticker1 = create_sticker(
    prompt=prompt,
    output_filename="../assets/cat.png",
    image_size="2K",     # Độ phân giải 2K giúp đường viền sắc nét hơn
    save_raw=True        # Lưu cả ảnh gốc nền xanh để đối chiếu
)

# Prompt mô tả sticker cá nhân hóa
prompt = "Một nhà phát triển mặc áo hoodie Google DeepMind, trông giống như tôi. Sử dụng các hình ảnh đính kèm của tôi và logo Google DeepMind mới."

# Danh sách ảnh đầu vào: ảnh chân dung và ảnh logo
input_images = ["../assets/headshot.png", "../assets/logo.png"]

# Gọi hàm create_sticker để thực hiện quy trình tự động
sticker2 = create_sticker(
    prompt=prompt,
    input_images=input_images,
    output_filename="../assets/developer.png",
    image_size="2K",     # Độ phân giải 2K giúp đường viền sắc nét hơn
    save_raw=True        # Lưu cả ảnh gốc nền xanh để đối chiếu
)