#!/usr/bin/env python3
"""Gemini UI Assets Generator CLI Tool.
Tự động hóa toàn bộ quy trình tạo UI/UX Assets, 3D Icons, Game Spritesheets, App Badges bằng Google Imagen 4.

Tính năng:
1. Prompt Engineering Chuẩn Quán Quân: Ép Nền Đen Phẳng Pure Black (#000000) + High Contrast Rim Lighting.
2. Sinh Ảnh Qua Google Imagen 4 REST API (imagen-4.0-generate-001).
3. Bóc Tách Nền Trong Suốt Tự Động Với Thuật Toán Border-Connected Flood-Fill Protection:
   - Xóa sạch 100% viền xám dư thừa (Edge Halo).
   - Bảo vệ 100% chi tiết màu tối bên trong vật thể (mắt đen gấu baby, đai đen, giáp tối, profile glyphs).
   - Chiết xuất bóng phản chiếu mềm (Soft Reflection Shadow) kênh Alpha.
"""

from __future__ import annotations

import argparse
import base64
import json
import os
import subprocess
import sys
import time
import urllib.request
from pathlib import Path

CREDENTIALS_PATH = os.path.expanduser("~/.gemini/antigravity/.credentials.json")
MODEL_GEMINI = "imagen-4.0-generate-001"
SCRIPT_REMOVE_CHROMA = Path(__file__).parent / "remove_chroma_key.py"

STYLE_PRESETS = {
    "3d-glass": (
        "3D glassmorphism glossy icon render, vibrant gradient colors, frosted acrylic refraction, soft claymation touch."
    ),
    "claymation": (
        "Handcrafted soft polymer clay texture, tactile organic indents, smooth matte finish, cute chibi aesthetic."
    ),
    "game-spritesheet": (
        "Uniform 2D/3D game animation spritesheet strip, crisp character silhouettes, game-ready asset."
    ),
    "magic-fx": (
        "Vibrant glowing 3D particle magic spell effect, intense self-illuminated aura, particle embers."
    ),
    "isometric": (
        "Clean low-poly 3D render, vibrant pastel colors, sharp geometric isometric angle."
    ),
    "luxury-gold": (
        "Photorealistic polished gold metal, lustrous white pearl reflections, high specular jewelry shine."
    ),
}


def get_gemini_key() -> str:
    key = os.getenv("GEMINI_API_KEY")
    if not key and os.path.exists(CREDENTIALS_PATH):
        try:
            with open(CREDENTIALS_PATH, "r", encoding="utf-8") as f:
                key = json.load(f).get("gemini_api_key")
        except Exception:
            pass
    return key or ""


def compile_imagen4_prompt(user_prompt: str, style_name: str = "3d-glass", grid: str | None = None) -> str:
    style_desc = STYLE_PRESETS.get(style_name, STYLE_PRESETS["3d-glass"])
    
    grid_prefix = ""
    if grid:
        grid_prefix = f"A {grid} grid layout array of "

    black_formula = (
        f"{grid_prefix}{user_prompt}.\n"
        f"STYLE: {style_desc}\n"
        f"BACKGROUND: 100% solid flat pure pitch black backdrop (#000000) with zero texture and zero wall ambient light.\n"
        f"LIGHTING: Clean high contrast rim lighting framing all shapes sharply off the dark background.\n"
        f"ENVIRONMENT: Floating in zero-gravity digital space with ZERO floor plane, ZERO floor reflection, ZERO cast shadow, ZERO ambient fog."
    )
    return black_formula


def generate_imagen4_image(prompt: str, out_path: Path, gemini_key: str) -> bool:
    url = f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL_GEMINI}:predict?key={gemini_key}"
    payload = {
        "instances": [{"prompt": prompt}],
        "parameters": {"sampleCount": 1, "aspectRatio": "1:1"},
    }
    req = urllib.request.Request(
        url, data=json.dumps(payload).encode("utf-8"), headers={"Content-Type": "application/json"}
    )
    try:
        with urllib.request.urlopen(req, timeout=90) as resp:
            data = json.loads(resp.read().decode("utf-8"))
            b64_data = data["predictions"][0].get("bytesBase64Encoded")
            if b64_data:
                out_path.parent.mkdir(parents=True, exist_ok=True)
                out_path.write_bytes(base64.b64decode(b64_data))
                return True
    except Exception as e:
        print(f"❌ Error generating image: {e}", file=sys.stderr)
    return False


def run_matting(raw_path: Path, trans_path: Path, tolerance: int = 30) -> bool:
    if not SCRIPT_REMOVE_CHROMA.exists():
        print(f"❌ Matting script missing: {SCRIPT_REMOVE_CHROMA}", file=sys.stderr)
        return False

    cmd = [
        sys.executable, str(SCRIPT_REMOVE_CHROMA),
        "--input", str(raw_path),
        "--out", str(trans_path),
        "--key-color", "#000000",
        "--tolerance", str(tolerance),
        "--soft-matte", "--force"
    ]
    proc = subprocess.run(cmd, capture_output=True, text=True)
    if proc.returncode != 0:
        print(f"⚠️ Matting warning: {proc.stderr}", file=sys.stderr)
        return False
    return True


def main():
    parser = argparse.ArgumentParser(description="Gemini UI Assets Generator Tool (Imagen 4)")
    parser.add_argument("--prompt", "-p", required=True, help="Mô tả asset cần tạo (ví dụ: '12 3D glossy icons')")
    parser.add_argument("--style", "-s", choices=list(STYLE_PRESETS.keys()), default="3d-glass", help="Phong cách thiết kế")
    parser.add_argument("--grid", "-g", help="Cấu trúc lưới (ví dụ: 2x2, 3x3, 4x3, 1x4)")
    parser.add_argument("--out", "-o", required=True, help="Đường dẫn file đầu ra (PNG)")
    parser.add_argument("--no-transparent", action="store_true", help="Bỏ qua bước tách nền trong suốt")
    parser.add_argument("--tolerance", type=int, default=30, help="Độ nhạy tách nền (Mặc định: 30)")

    args = parser.parse_args()

    gemini_key = get_gemini_key()
    if not gemini_key:
        print("❌ Lỗi: GEMINI_API_KEY chưa được cấu hình!", file=sys.stderr)
        sys.exit(1)

    out_file = Path(args.out).resolve()
    raw_file = out_file.parent / f"{out_file.stem}_raw.png"

    compiled_prompt = compile_imagen4_prompt(args.prompt, style_name=args.style, grid=args.grid)

    print("=" * 70)
    print("🚀 GEMINI UI ASSET GENERATOR (GOOGLE IMAGEN 4)")
    print("=" * 70)
    print(f"📌 Asset Prompt: {args.prompt}")
    print(f"🎨 Style Preset: {args.style}")
    print(f"📝 Compiled Imagen 4 Prompt:\n{compiled_prompt}")
    print("-" * 70)

    t0 = time.time()
    ok_gen = generate_imagen4_image(compiled_prompt, raw_file, gemini_key)
    if not ok_gen:
        print("❌ Thất bại trong khâu tạo ảnh từ Imagen 4!", file=sys.stderr)
        sys.exit(1)

    print(f"✅ Đã tạo ảnh gốc Imagen 4 ({round(time.time() - t0, 2)}s) -> {raw_file}")

    if args.no_transparent:
        print(f"🎉 Hoàn tất! Ảnh gốc giữ nguyên: {raw_file}")
        return

    print("✂️ Đang chạy thuật toán Border-Connected Flood-Fill Matting...")
    ok_mat = run_matting(raw_file, out_file, tolerance=args.tolerance)
    if ok_mat:
        print("=" * 70)
        print(f"🎉 THÀNH CÔNG! Asset UI Tách Nền Trong Suốt Sắc Nét:\n👉 {out_file}")
        print("=" * 70)
    else:
        print("⚠️ Cảnh báo: Tách nền gặp lỗi, hãy kiểm tra ảnh gốc.")


if __name__ == "__main__":
    main()