#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import gradio as gr
import requests
import base64
import json
import os
import gc
import numpy as np
import cv2
from PIL import Image
from io import BytesIO

# ========== 可选依赖导入（云端API）==========
try:
    import dashscope
    from dashscope import ImageSynthesis
    DASHSCOPE_AVAILABLE = True
except ImportError:
    DASHSCOPE_AVAILABLE = False

try:
    import google.generativeai as genai
    GOOGLE_API_AVAILABLE = True
except ImportError:
    GOOGLE_API_AVAILABLE = False

# ================== 配置 ==================
SILICONFLOW_API_BASE = "https://api.siliconflow.cn/v1"
# 图像编辑模型（SiliconFlow）
DEFAULT_EDIT_MODEL_SF = "Qwen/Qwen-Image-Edit"
# 文本图像生成模型（通义千问）
DEFAULT_TEXT2IMG_MODEL_TONGYI = "wanx-v1"
# 图像生成模型（Gemini）
DEFAULT_TEXT2IMG_MODEL_GEMINI = "gemini-1.5-pro"  # Gemini 实际上通过 generate_content 生成图像，需具体实现

BASE_NEGATIVE = "worst quality, low quality, ugly, deformed, blurry, messy, watermark, signature"

# ================== 风格预设（与之前一致）==================
def load_style_presets():
    default_presets = {
        "自定义": {"positive": "", "negative": BASE_NEGATIVE},
        "🎨 二次元动漫": {"positive": "anime style, manga, vibrant colors, cel shading, studio ghibli, high quality, detailed background", "negative": "photorealistic, realistic, 3d render, blurry, low quality, ugly"},
        "📸 写实摄影": {"positive": "photorealistic, 8k uhd, professional photography, sharp focus, natural lighting, high detail, canon r5", "negative": "anime, cartoon, illustration, painting, blurry, low quality"},
        "🌆 赛博朋克": {"positive": "cyberpunk city, neon lights, dark alley, rain, glowing signs, intricate details, blade runner style", "negative": "bright daylight, rural, low quality, ugly"},
        "🖌️ 油画风格": {"positive": "oil painting, thick brushstrokes, canvas texture, impressionist, rich colors, artistic", "negative": "photorealistic, smooth, digital art, 3d render, blurry"},
        "💧 水彩风格": {"positive": "watercolor painting, soft edges, paper texture, flowing colors, artistic, delicate", "negative": "oil painting, thick strokes, high contrast, ugly"},
        "🔮 梦幻奇幻": {"positive": "fantasy art, ethereal, glowing magical atmosphere, dreamlike, detailed, intricate, soft colors", "negative": "realistic, modern, city, ugly, low quality"},
        "🔥 火焰/灯光特效增强": {"positive": "dramatic lighting, cinematic light, glowing fire, warm rim light, high contrast, volumetric light", "negative": "flat lighting, no shadows, dull, low contrast"},
    }
    json_path = os.path.join(os.path.dirname(__file__), "prompts.json")
    if os.path.exists(json_path):
        try:
            with open(json_path, "r", encoding="utf-8") as f:
                user = json.load(f)
                default_presets.update(user)
        except Exception as e:
            print(f"加载 prompts.json 失败: {e}")
    return default_presets

STYLE_PRESETS = load_style_presets()

# ================== 工具函数 ==================
def encode_image_to_base64(pil_image, format="PNG"):
    buffered = BytesIO()
    pil_image.save(buffered, format=format)
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def pil_to_data_url(pil_image):
    return f"data:image/png;base64,{encode_image_to_base64(pil_image)}"

# ================== 云端API调用函数 ==================
def siliconflow_image_edit(content_img, prompt, negative_prompt, api_key, strength=0.3):
    """调用 SiliconFlow Qwen-Image-Edit API"""
    if not api_key:
        return None, "请提供 SiliconFlow API Key"
    if content_img is None:
        return None, "请上传内容图像"
    try:
        img_url = pil_to_data_url(content_img)
        headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
        payload = {
            "model": DEFAULT_EDIT_MODEL_SF,
            "image": img_url,
            "prompt": prompt,
            "negative_prompt": negative_prompt if negative_prompt else "",
            "num_inference_steps": 30,
            "guidance_scale": 7.0,
            "strength": strength,
            "output_format": "url"
        }
        response = requests.post(f"{SILICONFLOW_API_BASE}/images/edits", headers=headers, json=payload, timeout=120)
        if response.status_code == 200:
            data = response.json()
            img_url = data['images'][0]['url']
            img_resp = requests.get(img_url, timeout=60)
            if img_resp.status_code == 200:
                return Image.open(BytesIO(img_resp.content)), "SiliconFlow 重绘成功"
            else:
                return None, f"下载图片失败: {img_resp.status_code}"
        else:
            return None, f"API 错误 {response.status_code}: {response.text}"
    except Exception as e:
        return None, f"异常: {str(e)}"

def tongyi_text2img(prompt, negative_prompt, api_key):
    """通义千问文生图（无法编辑原图，仅根据提示词生成新图）"""
    if not api_key:
        return None, "请提供通义千问 API Key"
    if not DASHSCOPE_AVAILABLE:
        return None, "请安装 dashscope: pip install dashscope"
    try:
        dashscope.api_key = api_key
        response = ImageSynthesis.call(model=DEFAULT_TEXT2IMG_MODEL_TONGYI,
                                        prompt=prompt,
                                        negative_prompt=negative_prompt,
                                        n=1,
                                        size='1024*1024')
        if response.status_code == 200:
            img_url = response.output.results[0].url
            img_resp = requests.get(img_url, timeout=60)
            if img_resp.status_code == 200:
                return Image.open(BytesIO(img_resp.content)), "通义千问生成成功"
            else:
                return None, f"下载图片失败: {img_resp.status_code}"
        else:
            return None, f"API 错误: {response.message}"
    except Exception as e:
        return None, f"异常: {str(e)}"

def gemini_text2img(prompt, negative_prompt, api_key):
    """Gemini 文生图（需要支持图像输出的模型）"""
    if not api_key:
        return None, "请提供 Gemini API Key"
    if not GOOGLE_API_AVAILABLE:
        return None, "请安装 google-generativeai: pip install google-generativeai"
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel(DEFAULT_TEXT2IMG_MODEL_GEMINI)
        # Gemini 目前图像生成需要特殊模型，此处简单占位
        # 实际上 Gemini 图像生成需使用 Imagen 模型，这里简化处理
        response = model.generate_content(f"Generate an image: {prompt}")
        # 解析响应获取图像（假设返回图片base64）
        # 实际情况需根据Gemini API最新文档调整
        return None, "Gemini 文生图暂未完全实现，请使用 SiliconFlow 或通义千问"
    except Exception as e:
        return None, f"异常: {str(e)}"

def cloud_relight(service, content_img, prompt, negative_prompt, api_key, strength):
    """根据服务选择调用相应API"""
    if service == "SiliconFlow":
        return siliconflow_image_edit(content_img, prompt, negative_prompt, api_key, strength)
    elif service == "通义千问":
        if content_img is not None:
            # 通义千问暂不支持图生图，可提示用户未实现
            return None, "通义千问 API 目前仅支持文生图，不会保留原图结构。建议使用 SiliconFlow。"
        else:
            return tongyi_text2img(prompt, negative_prompt, api_key)
    elif service == "Gemini":
        return gemini_text2img(prompt, negative_prompt, api_key)
    else:
        return None, "未知服务"

# ================== 保留原有本地处理功能 ==================
def lighting_transfer_advanced(image, exposure, temp, contrast, intensity):
    if image is None:
        return None, "无图像"
    img = np.array(image).astype(np.float32) / 255.0
    lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
    L, a, b = cv2.split(lab)
    if exposure != 0:
        gamma = 1.0 / (1.0 + exposure * 0.8) if exposure >= 0 else 1.0 / (1.0 - exposure * 0.6)
        L = np.power(L, gamma)
        L = np.clip(L, 0, 1)
    if contrast != 1.0:
        L_mid = 0.5
        L = (L - L_mid) * contrast + L_mid
        L = np.clip(L, 0, 1)
    if temp != 0:
        a = a + temp * 0.1
        b = b - temp * 0.05
        a = np.clip(a, -128, 127)
        b = np.clip(b, -128, 127)
    lab_adjusted = cv2.merge([L, a, b])
    rgb_adjusted = cv2.cvtColor(lab_adjusted, cv2.COLOR_LAB2RGB)
    rgb_adjusted = np.clip(rgb_adjusted, 0, 1)
    if intensity < 1.0:
        rgb_adjusted = cv2.addWeighted(img, 1-intensity, rgb_adjusted, intensity, 0)
    result = (rgb_adjusted * 255).astype(np.uint8)
    return Image.fromarray(result), f"滑块调整完成 (曝光{exposure:+.1f}, 色温{temp:+.1f}, 对比度{contrast:.2f}, 强度{intensity:.2f})"

def apply_lighting_effect(image, lighting_type):
    if image is None:
        return None, "无图像"
    img = np.array(image)
    if lighting_type == "电影级灯光":
        img = cv2.convertScaleAbs(img, alpha=1.2, beta=5)
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
        hsv[:,:,1] = np.clip(hsv[:,:,1] * 1.3, 0, 255)
        img = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
    elif lighting_type == "柔和漫射光":
        img = cv2.convertScaleAbs(img, alpha=0.9, beta=20)
        img = cv2.GaussianBlur(img, (3,3), 0.5)
    elif lighting_type == "边缘轮廓光":
        lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
        l, a, b = cv2.split(lab)
        l = cv2.addWeighted(l, 1.2, np.zeros_like(l), 0, 10)
        lab = cv2.merge([l, a, b])
        img = cv2.cvtColor(lab, cv2.COLOR_LAB2RGB)
    elif lighting_type == "霓虹光效":
        hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV).astype(np.float32)
        hsv[:,:,0] = (hsv[:,:,0] + 30) % 180
        hsv[:,:,1] = np.clip(hsv[:,:,1] * 1.5, 0, 255)
        hsv[:,:,2] = np.clip(hsv[:,:,2] * 0.7, 0, 255)
        img = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2RGB)
    elif lighting_type == "黄金时刻光":
        b, g, r = cv2.split(img)
        r = np.clip(r * 1.3, 0, 255).astype(np.uint8)
        g = np.clip(g * 1.1, 0, 255).astype(np.uint8)
        b = (b * 0.8).astype(np.uint8)
        img = cv2.merge([r, g, b])
    elif lighting_type == "影棚标准光":
        img = cv2.convertScaleAbs(img, alpha=1.0, beta=10)
    return Image.fromarray(img), f"已应用滤镜: {lighting_type}"

def color_transfer(content_img, style_img, intensity=0.8):
    if content_img is None or style_img is None:
        return None, "请同时上传内容图和参考图"
    content_np = np.array(content_img).astype(np.float32) / 255.0
    style_np = np.array(style_img.resize(content_img.size)).astype(np.float32) / 255.0
    content_lab = cv2.cvtColor(content_np, cv2.COLOR_RGB2LAB)
    style_lab = cv2.cvtColor(style_np, cv2.COLOR_RGB2LAB)
    for i in range(1,3):
        c_mean = content_lab[:,:,i].mean()
        c_std = content_lab[:,:,i].std()
        s_mean = style_lab[:,:,i].mean()
        s_std = style_lab[:,:,i].std()
        content_lab[:,:,i] = (content_lab[:,:,i] - c_mean) * (s_std / (c_std + 1e-8)) + s_mean
    transferred = cv2.cvtColor(content_lab, cv2.COLOR_LAB2RGB)
    result = cv2.addWeighted(content_np, 1-intensity, transferred, intensity, 0)
    result = np.clip(result * 255, 0, 255).astype(np.uint8)
    return Image.fromarray(result), f"颜色迁移完成 (强度={intensity})"

def relight_by_reference(content_img, style_img, strength=0.8, preserve_luminance=True):
    if content_img is None:
        return None, "无内容图像"
    if style_img is None:
        return content_img, "未提供参考图，返回原图"
    content_np = np.array(content_img).astype(np.float32) / 255.0
    style_np = np.array(style_img.resize(content_img.size)).astype(np.float32) / 255.0
    content_lab = cv2.cvtColor(content_np, cv2.COLOR_RGB2LAB)
    style_lab = cv2.cvtColor(style_np, cv2.COLOR_RGB2LAB)
    for i in range(1,3):
        c_mean = content_lab[:,:,i].mean()
        c_std = content_lab[:,:,i].std()
        s_mean = style_lab[:,:,i].mean()
        s_std = style_lab[:,:,i].std()
        content_lab[:,:,i] = (content_lab[:,:,i] - c_mean) * (s_std / (c_std + 1e-8)) + s_mean
    if not preserve_luminance:
        L_content = content_lab[:,:,0]
        L_style = style_lab[:,:,0]
        Lc = (L_content - L_content.min()) / (L_content.max() - L_content.min() + 1e-8)
        Ls = (L_style - L_style.min()) / (L_style.max() - L_style.min() + 1e-8)
        hist_c, bins = np.histogram(Lc.flatten(), 256, [0,1])
        hist_s, _ = np.histogram(Ls.flatten(), 256, [0,1])
        cdf_c = hist_c.cumsum() / hist_c.sum()
        cdf_s = hist_s.cumsum() / hist_s.sum()
        mapping = np.interp(cdf_c, cdf_s, np.linspace(0,1,256))
        L_matched = np.interp(Lc.flatten(), bins[:-1], mapping).reshape(Lc.shape)
        content_lab[:,:,0] = L_matched * (L_content.max() - L_content.min()) + L_content.min()
    else:
        content_lab[:,:,0] = content_lab[:,:,0]
    transferred = cv2.cvtColor(content_lab, cv2.COLOR_LAB2RGB)
    result = cv2.addWeighted(content_np, 1-strength, transferred, strength, 0)
    result = np.clip(result * 255, 0, 255).astype(np.uint8)
    return Image.fromarray(result), f"参考图迁移完成 (强度={strength:.2f})"

def clear_cache():
    gc.collect()
    return "内存已清理"

def update_prompts_from_style(style_name):
    preset = STYLE_PRESETS.get(style_name, STYLE_PRESETS["自定义"])
    return preset["positive"], preset["negative"]

# ================== Gradio 界面 ==================
with gr.Blocks(title="AI 光影工坊 - 云端增强版", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 AI 光影工坊 - 云端重绘增强版")
    gr.Markdown("**保留所有原有本地功能，新增「云端AI重绘」模块，支持白天转黑夜、专业灯光效果。**")
    with gr.Row():
        with gr.Column():
            content_img = gr.Image(label="📷 内容图像", type="pil")
            style_img = gr.Image(label="🎨 风格参考图（用于参考图迁移）", type="pil")
            style_selector = gr.Dropdown(choices=list(STYLE_PRESETS.keys()), value="自定义", label="✨ 风格预设")
            prompt_text = gr.Textbox(label="📝 正提示词", lines=2)
            negative_text = gr.Textbox(label="🚫 负提示词", lines=2, value=BASE_NEGATIVE)
            # 原有本地滑块调整
            gr.Markdown("### 🌟 手动滑块调整 (结构不变)")
            with gr.Row():
                exposure_slider = gr.Slider(label="曝光补偿", minimum=-0.8, maximum=0.8, value=0.0, step=0.05)
                temp_slider = gr.Slider(label="色温", minimum=-0.8, maximum=0.8, value=0.0, step=0.05)
            with gr.Row():
                contrast_slider = gr.Slider(label="对比度", minimum=0.6, maximum=1.8, value=1.0, step=0.05)
                intensity_slider = gr.Slider(label="整体强度", minimum=0.0, maximum=1.0, value=0.8, step=0.05)
            btn_local = gr.Button("🌟 应用滑块调整", variant="secondary")
            # 参考图迁移
            gr.Markdown("### 🎨 参考图光影迁移")
            ref_strength = gr.Slider(label="迁移强度", minimum=0.0, maximum=1.0, value=0.8, step=0.05)
            preserve_lum = gr.Checkbox(label="保留原图亮度结构", value=True)
            btn_reference = gr.Button("🎨 参考图光影迁移", variant="secondary")
            # 快速滤镜和颜色迁移放在右侧列？
        with gr.Column():
            output_img = gr.Image(label="结果", type="pil", interactive=False, height=450)
            lighting_type = gr.Dropdown(choices=["电影级灯光","柔和漫射光","边缘轮廓光","霓虹光效","黄金时刻光","影棚标准光"], value="电影级灯光", label="💡 快速滤镜")
            btn_filter = gr.Button("🎬 应用快速滤镜", variant="secondary")
            color_intensity = gr.Slider(label="颜色迁移强度", minimum=0.0, maximum=1.0, value=0.7, step=0.05)
            btn_color = gr.Button("🎨 纯颜色迁移", variant="secondary")
            btn_clear = gr.Button("🗑️ 清理缓存", variant="secondary")
            # 云端AI重绘区域（新增）
            gr.Markdown("---\n### 🌐 云端AI重绘（保留结构 + 专业灯光）")
            cloud_service = gr.Radio(choices=["SiliconFlow", "通义千问", "Gemini"], label="选择服务", value="SiliconFlow")
            with gr.Column(visible=True) as sf_key_col:
                sf_key = gr.Textbox(label="SiliconFlow API Key", type="password", placeholder="sk-...")
            with gr.Column(visible=False) as tongyi_key_col:
                tongyi_key = gr.Textbox(label="通义千问 API Key", type="password", placeholder="sk-...")
            with gr.Column(visible=False) as gemini_key_col:
                gemini_key = gr.Textbox(label="Gemini API Key", type="password", placeholder="...")
            cloud_strength = gr.Slider(label="重绘强度（仅SiliconFlow有效）", minimum=0.1, maximum=0.6, value=0.3, step=0.01,
                                       info="值越低越保留原图结构")
            btn_cloud = gr.Button("☁️ 开始云端重绘", variant="primary")
            output_info = gr.Textbox(label="状态", interactive=False)
    
    # 动态显示API Key输入框
    def update_cloud_visibility(service):
        return {
            sf_key_col: gr.update(visible=(service == "SiliconFlow")),
            tongyi_key_col: gr.update(visible=(service == "通义千问")),
            gemini_key_col: gr.update(visible=(service == "Gemini"))
        }
    cloud_service.change(fn=update_cloud_visibility, inputs=cloud_service, outputs=[sf_key_col, tongyi_key_col, gemini_key_col])
    
    # 事件绑定
    style_selector.change(fn=update_prompts_from_style, inputs=style_selector, outputs=[prompt_text, negative_text])
    btn_local.click(fn=lighting_transfer_advanced,
                    inputs=[content_img, exposure_slider, temp_slider, contrast_slider, intensity_slider],
                    outputs=[output_img, output_info])
    btn_reference.click(fn=relight_by_reference,
                        inputs=[content_img, style_img, ref_strength, preserve_lum],
                        outputs=[output_img, output_info])
    btn_filter.click(fn=apply_lighting_effect, inputs=[content_img, lighting_type], outputs=[output_img, output_info])
    btn_color.click(fn=color_transfer, inputs=[content_img, style_img, color_intensity], outputs=[output_img, output_info])
    btn_clear.click(fn=clear_cache, inputs=[], outputs=[output_info])
    
    # 云端重绘：根据服务选择对应的API Key
    def cloud_relight_router(service, content_img, prompt, negative, sf_key, tongyi_key, gemini_key, strength):
        if service == "SiliconFlow":
            api_key = sf_key
        elif service == "通义千问":
            api_key = tongyi_key
        elif service == "Gemini":
            api_key = gemini_key
        else:
            return None, "未知服务"
        return cloud_relight(service, content_img, prompt, negative, api_key, strength)
    
    btn_cloud.click(fn=cloud_relight_router,
                    inputs=[cloud_service, content_img, prompt_text, negative_text, sf_key, tongyi_key, gemini_key, cloud_strength],
                    outputs=[output_img, output_info])
    
    demo.load(fn=lambda: update_prompts_from_style("自定义"), outputs=[prompt_text, negative_text])

if __name__ == "__main__":
    demo.launch()