#!/usr/bin/env python3
"""
Gemini i2i (Image-to-Image) 앵글/색상 편집 테스트.

Nano Banana Pro 방식: 원본 이미지 + 텍스트 프롬프트 → Gemini generateContent로 편집.
카메라 다이어그램 없이도, 텍스트로 앵글/조명 변화를 지시할 수 있음.

실행: cd backend && .venv/bin/python tests/test_gemini_i2i_angle.py
"""

import base64
import json
import os
import sys
import time
import urllib.request
import urllib.error
from pathlib import Path
from PIL import Image
import io
import math

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
# 2026-04-27 hotfix: prod URL setdefault 제거. 단독 실행 시 prod DB가 의도하지 않은
# 변경에 노출될 위험이 있어 sqlite default로 변경. prod DB가 필요하면 명시적으로
# export DATABASE_URL=postgresql://... 후 실행할 것.
import tempfile as _tmp_default
os.environ.setdefault("DATABASE_URL", f"sqlite:///{_tmp_default.gettempdir()}/theroad_standalone.sqlite")

from app.core.config import settings

OUTPUT_DIR = Path(settings.projects_dir) / "gemini_i2i_test"
GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"


def create_camera_diagram(angle_deg: int, elevation_deg: int = 0, size: int = 512) -> bytes:
    """카메라 앵글 다이어그램 이미지 생성 (PIL로 그리기)."""
    img = Image.new("RGB", (size, size), (30, 30, 40))
    from PIL import ImageDraw, ImageFont
    draw = ImageDraw.Draw(img)

    cx, cy = size // 2, size // 2
    r = size // 3

    # 원 그리기
    draw.ellipse([cx - r, cy - r, cx + r, cy + r], outline=(200, 200, 200), width=2)

    # 0도, 90도 표시
    draw.text((cx + r + 10, cy - 10), "0°", fill=(255, 100, 100))
    draw.text((cx - 15, cy + r + 10), "90°", fill=(255, 100, 100))

    # IMAGE 표시 (중앙)
    draw.rectangle([cx - 30, cy - 8, cx + 30, cy + 8], fill=(80, 80, 80), outline=(200, 200, 200))
    draw.text((cx - 20, cy - 6), "IMAGE", fill=(255, 255, 255))

    # 카메라 위치 계산 (각도 기준)
    rad = math.radians(angle_deg)
    cam_x = cx + int(r * 0.7 * math.cos(rad))
    cam_y = cy + int(r * 0.7 * math.sin(rad))

    # 카메라 → 이미지 화살표
    draw.line([(cam_x, cam_y), (cx, cy)], fill=(255, 200, 0), width=3)

    # 카메라 아이콘
    draw.rectangle([cam_x - 12, cam_y - 10, cam_x + 12, cam_y + 10], fill=(180, 180, 180), outline=(255, 255, 255))
    draw.text((cam_x - 25, cam_y + 15), "CAMERA", fill=(200, 200, 200))

    # 각도 텍스트
    draw.text((10, 10), f"Angle: {angle_deg}°", fill=(100, 200, 255))
    if elevation_deg != 0:
        draw.text((10, 30), f"Elevation: {elevation_deg}°", fill=(100, 200, 255))

    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return buf.getvalue()


def gemini_i2i(prompt: str, input_images: list[bytes], aspect_ratio: str = "16:9") -> bytes:
    """Gemini generateContent로 이미지 편집 (i2i).

    원본 이미지 + 프롬프트 → 편집된 이미지.
    """
    parts = [{"text": prompt}]

    for img_bytes in input_images:
        parts.append({
            "inline_data": {
                "mime_type": "image/png",
                "data": base64.b64encode(img_bytes).decode("ascii"),
            }
        })

    body = {
        "contents": [{"parts": parts}],
        "generationConfig": {
            "responseModalities": ["TEXT", "IMAGE"],
            "imageConfig": {
                "aspectRatio": aspect_ratio,
                "imageSize": "2K",
            },
        },
    }

    url = GEMINI_URL.format(model=settings.gemini_image_model, api_key=settings.gemini_api_key)
    req = urllib.request.Request(url, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})

    with urllib.request.urlopen(req, timeout=120) as resp:
        data = json.loads(resp.read())

    for candidate in data.get("candidates", []):
        for part in candidate.get("content", {}).get("parts", []):
            inline = part.get("inlineData") or part.get("inline_data")
            if inline and inline.get("data"):
                return base64.b64decode(inline["data"])

    # 에러 정보 추출
    block = data.get("promptFeedback", {}).get("blockReason", "")
    if block:
        raise RuntimeError(f"Blocked: {block}")
    raise RuntimeError(f"No image: {json.dumps(data)[:300]}")


def find_scene_images(limit: int = 3) -> list:
    """기존 씬 이미지 찾기."""
    projects_dir = Path(settings.projects_dir)
    return sorted(projects_dir.rglob("*/scene/*.png"))[:limit]


def run():
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    print("=== Gemini i2i 앵글/색상 편집 테스트 ===")
    print(f"모델: {settings.gemini_image_model}")
    print()

    scene_images = find_scene_images(3)
    if not scene_images:
        print("❌ 씬 이미지 없음")
        return

    print(f"테스트 이미지: {len(scene_images)}개\n")

    # 앵글 변형 (텍스트 프롬프트 + 카메라 다이어그램)
    angle_edits = [
        {
            "name": "우측 45도",
            "prompt": "Create a new angle of this scene, rotating the camera 45 degrees to the right. Keep all characters, objects, and lighting the same. Only change the viewpoint.",
            "diagram_angle": 45,
        },
        {
            "name": "위에서 내려다봄",
            "prompt": "Create a new angle of this scene, shot from a high angle looking down at about 30 degrees. Keep all characters and objects the same. Bird's eye perspective.",
            "diagram_angle": 0,
            "diagram_elevation": 30,
        },
        {
            "name": "낮은 앵글 (올려다봄)",
            "prompt": "Create a new angle of this scene, shot from a low angle looking up. The camera is near ground level, making characters appear more imposing. Keep everything else the same.",
            "diagram_angle": 0,
            "diagram_elevation": -20,
        },
    ]

    # 색상/조명 변형
    color_edits = [
        {
            "name": "따뜻한 석양 톤",
            "prompt": "Re-light this exact scene with warm golden sunset lighting. Orange and amber tones. Long shadows. Keep the composition, characters, and objects exactly the same. Only change the lighting and color temperature.",
        },
        {
            "name": "차가운 청색 야간",
            "prompt": "Re-light this exact scene as a cold blue night scene. Moonlight and cool tones. Keep the composition, characters, and objects exactly the same. Only change the lighting to nighttime blue.",
        },
        {
            "name": "네온 사이버펑크",
            "prompt": "Re-light this exact scene with neon cyberpunk lighting. Pink and cyan neon reflections. Dark shadows with bright neon highlights. Keep composition and characters the same.",
        },
    ]

    results = []

    for img_idx, img_path in enumerate(scene_images):
        print(f"[{img_idx + 1}/{len(scene_images)}] {img_path.name}")
        original_bytes = img_path.read_bytes()

        img_result = {
            "source": str(img_path),
            "source_name": img_path.name,
            "edits": [],
        }

        # 앵글 변형
        for edit in angle_edits:
            print(f"  앵글: {edit['name']}...")

            # 카메라 다이어그램 생성
            diagram_bytes = create_camera_diagram(
                edit.get("diagram_angle", 0),
                edit.get("diagram_elevation", 0),
            )

            try:
                result_bytes = gemini_i2i(
                    prompt=edit["prompt"],
                    input_images=[original_bytes, diagram_bytes],
                )
                out_name = f"img{img_idx}_angle_{edit['name'].replace(' ', '_')}.png"
                out_path = OUTPUT_DIR / out_name
                out_path.write_bytes(result_bytes)
                print(f"    ✅ ({len(result_bytes) // 1024}KB)")
                img_result["edits"].append({
                    "name": edit["name"],
                    "type": "angle",
                    "status": "success",
                    "path": str(out_path),
                })
            except Exception as e:
                print(f"    ❌ {str(e)[:100]}")
                img_result["edits"].append({
                    "name": edit["name"],
                    "type": "angle",
                    "status": "error",
                    "error": str(e)[:200],
                })
            time.sleep(3)

        # 색상 변형
        for edit in color_edits:
            print(f"  색상: {edit['name']}...")
            try:
                result_bytes = gemini_i2i(
                    prompt=edit["prompt"],
                    input_images=[original_bytes],
                )
                out_name = f"img{img_idx}_color_{edit['name'].replace(' ', '_')}.png"
                out_path = OUTPUT_DIR / out_name
                out_path.write_bytes(result_bytes)
                print(f"    ✅ ({len(result_bytes) // 1024}KB)")
                img_result["edits"].append({
                    "name": edit["name"],
                    "type": "color",
                    "status": "success",
                    "path": str(out_path),
                })
            except Exception as e:
                print(f"    ❌ {str(e)[:100]}")
                img_result["edits"].append({
                    "name": edit["name"],
                    "type": "color",
                    "status": "error",
                    "error": str(e)[:200],
                })
            time.sleep(3)

        results.append(img_result)

    # HTML 생성
    html = build_html(results)
    html_path = OUTPUT_DIR / "gemini_i2i_comparison.html"
    html_path.write_text(html, encoding="utf-8")
    print(f"\n✅ 비교 페이지: {html_path}")
    print(f"   open {html_path}")


def img_to_b64(path) -> str:
    return base64.b64encode(Path(path).read_bytes()).decode()


def build_html(results: list) -> str:
    cards = ""
    for r in results:
        orig_b64 = img_to_b64(r["source"])

        angle_imgs = ""
        color_imgs = ""
        for edit in r["edits"]:
            if edit["status"] == "success":
                b64 = img_to_b64(edit["path"])
                block = f'''
                <div class="var">
                    <h4>{edit["name"]}</h4>
                    <img src="data:image/png;base64,{b64}">
                </div>'''
            else:
                block = f'''
                <div class="var fail">
                    <h4>{edit["name"]}</h4>
                    <div class="err">{edit.get("error", "실패")[:100]}</div>
                </div>'''

            if edit["type"] == "angle":
                angle_imgs += block
            else:
                color_imgs += block

        cards += f'''
        <div class="card">
            <h3>원본: {r["source_name"]}</h3>
            <div class="orig"><img src="data:image/png;base64,{orig_b64}"></div>

            <h3 class="section">앵글 변형 (Gemini i2i + 카메라 다이어그램)</h3>
            <div class="grid">{angle_imgs}</div>

            <h3 class="section">색상/조명 변형 (Gemini i2i)</h3>
            <div class="grid">{color_imgs}</div>
        </div>'''

    return f'''<!DOCTYPE html>
<html lang="ko"><head><meta charset="utf-8">
<title>Gemini i2i 앵글/색상 편집 비교</title>
<style>
* {{ box-sizing: border-box; margin: 0; padding: 0; }}
body {{ font-family: -apple-system, sans-serif; background: #0f1117; color: #e4e6ed; padding: 24px; }}
h2 {{ margin-bottom: 8px; }}
.sub {{ color: #8b90a0; margin-bottom: 24px; font-size: 14px; }}
.card {{ background: #1e2230; border-radius: 12px; padding: 20px; margin-bottom: 24px; }}
.card h3 {{ margin-bottom: 12px; font-size: 15px; }}
.section {{ margin-top: 20px; color: #6386ff; }}
.orig img {{ max-width: 500px; border-radius: 8px; margin-bottom: 16px; }}
.grid {{ display: grid; grid-template-columns: repeat(auto-fill, minmax(300px, 1fr)); gap: 12px; }}
.var {{ background: #252938; border-radius: 8px; padding: 10px; }}
.var h4 {{ font-size: 13px; color: #34d399; margin-bottom: 8px; }}
.var img {{ width: 100%; border-radius: 6px; }}
.fail {{ border: 1px solid rgba(248,113,113,0.3); }}
.fail h4 {{ color: #f87171; }}
.err {{ color: #f87171; font-size: 12px; }}
</style></head><body>
<h2>Gemini i2i 앵글/색상 편집 비교</h2>
<p class="sub">Gemini {settings.gemini_image_model} — 원본 이미지 + 프롬프트 + 카메라 다이어그램으로 편집</p>
{cards}
</body></html>'''


if __name__ == "__main__":
    run()