from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from PIL import Image
import io
import base64
import os
import httpx
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

app = FastAPI(title="AI Renovation Inference Service")

STABILITY_API_KEY = os.getenv("STABILITY_API_KEY")
if not STABILITY_API_KEY:
    print("⚠️ Warning: STABILITY_API_KEY not set. Will use placeholder images.")

def generate_prompt(style: str, layout_id: str, room_type: str = "living_room") -> tuple:
    """根据风格和房间类型生成 Stability AI 提示词和负面提示词"""
    
    # 房间类型映射
    room_prompts = {
        "living_room": "living room, sofa, coffee table, TV cabinet",
        "bedroom": "bedroom, bed, nightstand, wardrobe",
        "kitchen": "kitchen, cabinets, countertops, sink",
        "bathroom": "bathroom, bathtub, shower, vanity",
        "dining_room": "dining room, dining table, chairs"
    }
    room_desc = room_prompts.get(room_type, room_prompts["living_room"])
    
    # 风格提示词（极大差异化）
    # 支持中英文 key
    style_prompts = {
        "现代简约": f"ultra modern minimalist {room_type}, floor-to-ceiling windows, pure white walls, no decorations, sleek glass coffee table, hidden LED strip lighting, open concept, IKEA-style furniture, concrete floor, high ceiling, 8K realistic photo",
        "modern minimalist": f"ultra modern minimalist {room_type}, floor-to-ceiling windows, pure white walls, no decorations, sleek glass coffee table, hidden LED strip lighting, open concept, IKEA-style furniture, concrete floor, high ceiling, 8K realistic photo",
        
        "北欧风": f"authentic Scandinavian {room_type}, light birch wood floor, sheepskin rug, hygge atmosphere, pastel color palette (mint, blush, sky blue), pendant lamp with fabric shade, minimalist bookshelf, Monstera plant, white walls with wood accents, cozy textiles, 8K realistic photo",
        "Scandinavian": f"authentic Scandinavian {room_type}, light birch wood floor, sheepskin rug, hygge atmosphere, pastel color palette (mint, blush, sky blue), pendant lamp with fabric shade, minimalist bookshelf, Monstera plant, white walls with wood accents, cozy textiles, 8K realistic photo",
        
        "中式风格": f"traditional Chinese {room_type}, dark rosewood furniture with intricate carvings, red lacquer cabinets, golden dragon motifs, silk embroidered cushions, paper lanterns hanging from ceiling, moon gate entrance, bonsai tree, calligraphy wall art, terra cotta floor tiles, 8K realistic photo",
        "Chinese traditional": f"traditional Chinese {room_type}, dark rosewood furniture with intricate carvings, red lacquer cabinets, golden dragon motifs, silk embroidered cushions, paper lanterns hanging from ceiling, moon gate entrance, bonsai tree, calligraphy wall art, terra cotta floor tiles, 8K realistic photo",
        
        "美式风格": f"classic American {room_type}, warm beige walls, overstuffed sectional sofa, brick fireplace with wooden mantel, oak hardwood floor, vintage rug, family photos on wall, ceiling fan with lights, bay window with curtains, 8K realistic photo",
        "American classic": f"classic American {room_type}, warm beige walls, overstuffed sectional sofa, brick fireplace with wooden mantel, oak hardwood floor, vintage rug, family photos on wall, ceiling fan with lights, bay window with curtains, 8K realistic photo",
        
        "工业风": f"authentic industrial loft {room_type}, exposed red brick walls, polished concrete floor, black metal frame furniture, Edison bulb pendant lights, steel beams on ceiling, leather sofa with metal legs, vintage factory cart coffee table, pipe shelving, 8K realistic photo",
        "industrial": f"authentic industrial loft {room_type}, exposed red brick walls, polished concrete floor, black metal frame furniture, Edison bulb pendant lights, steel beams on ceiling, leather sofa with metal legs, vintage factory cart coffee table, pipe shelving, 8K realistic photo",
        
        "日式风格": f"zen Japanese {room_type}, tatami mat floor, low wooden table, shoji sliding paper doors, bonsai tree, neutral earth tones (beige, brown, gray), minimalist futon, bamboo accents, rock garden view through window, 8K realistic photo",
        "Japanese zen": f"zen Japanese {room_type}, tatami mat floor, low wooden table, shoji sliding paper doors, bonsai tree, neutral earth tones (beige, brown, gray), minimalist futon, bamboo accents, rock garden view through window, 8K realistic photo",
        
        "轻奢风": f"luxury modern {room_type}, Italian marble floor, gold-plated furniture legs, crystal chandelier, velvet sofa in emerald green, mirrored accent wall, abstract art painting, smart home automation visible, floor-to-ceiling curtains, 8K realistic photo",
        "luxury modern": f"luxury modern {room_type}, Italian marble floor, gold-plated furniture legs, crystal chandelier, velvet sofa in emerald green, mirrored accent wall, abstract art painting, smart home automation visible, floor-to-ceiling curtains, 8K realistic photo",
        
        "地中海风格": f"authentic Mediterranean {room_type}, white stucco walls with blue accents, arched doorways and windows, terracotta floor tiles, mosaic tile borders in blue and white, exposed wooden ceiling beams, wicker furniture with blue cushions, large clay pots with olive trees, sea view through window, wrought iron details, 8K realistic photo",
        "Mediterranean": f"authentic Mediterranean {room_type}, white stucco walls with blue accents, arched doorways and windows, terracotta floor tiles, mosaic tile borders in blue and white, exposed wooden ceiling beams, wicker furniture with blue cushions, large clay pots with olive trees, sea view through window, wrought iron details, 8K realistic photo",
        "mediterranean": f"authentic Mediterranean {room_type}, white stucco walls with blue accents, arched doorways and windows, terracotta floor tiles, mosaic tile borders in blue and white, exposed wooden ceiling beams, wicker furniture with blue cushions, large clay pots with olive trees, sea view through window, wrought iron details, 8K realistic photo"
    }
    base_prompt = style_prompts.get(style, style_prompts["现代简约"])
    
    # 组合提示词
    prompt = f"{base_prompt}, {room_desc}, professional architectural photography, canon EOS R5, 24mm lens, f/8, high detail"
    
    # 负面提示词（防止风格混淆）
    negative_prompt = "cartoon, drawing, sketch, low quality, blurry, distorted, deformed, ugly, bad anatomy, watermark, text, signature, other styles mixed, messy, cluttered"
    
    return prompt, negative_prompt

async def generate_with_stability(prompt: str, negative_prompt: str = "", size: str = "1024x1024") -> bytes:
    """调用 Stability AI API 生成图片"""
    if not STABILITY_API_KEY:
        # 如果没有 API Key，返回占位图
        return generate_placeholder_image()
    
    try:
        # 解析尺寸
        width, height = map(int, size.split('x'))
        
        async with httpx.AsyncClient() as client:
            response = await client.post(
                "https://api.stability.ai/v1/generation/stable-diffusion-xl-1024-v1-0/text-to-image",
                headers={
                    "Authorization": f"Bearer {STABILITY_API_KEY}",
                    "Content-Type": "application/json",
                    "Accept": "application/json"
                },
                json={
                    "text_prompts": [
                        {"text": prompt, "weight": 1},
                        {"text": negative_prompt, "weight": -1}  # 负面提示词
                    ],
                    "cfg_scale": 7,
                    "clip_guidance_preset": "NONE",
                    "height": height,
                    "width": width,
                    "samples": 1,
                    "steps": 30,
                },
                timeout=60.0
            )
            
            if response.status_code != 200:
                print(f"Stability AI API Error: {response.text}")
                return generate_placeholder_image()
            
            data = response.json()
            # Stability AI 返回 base64 编码的图片
            image_base64 = data["artifacts"][0]["base64"]
            image_bytes = base64.b64decode(image_base64)
            return image_bytes
            
    except Exception as e:
        print(f"Error calling Stability AI API: {e}")
        return generate_placeholder_image()

def generate_placeholder_image(width=512, height=512) -> bytes:
    """生成占位图（当 API 不可用时的后备方案）"""
    img = Image.new('RGB', (width, height), (240, 240, 240))
    img_byte_arr = io.BytesIO()
    img.save(img_byte_arr, format='PNG')
    img_byte_arr.seek(0)
    return img_byte_arr.getvalue()

@app.get("/")
async def root():
    return {
        "message": "AI Renovation Inference Service",
        "endpoints": ["/inference", "/inference/all"],
        "model": "Stable Diffusion XL" if STABILITY_API_KEY else "Placeholder"
    }

@app.post("/inference")
async def inference(style: str = "现代简约", layout_id: str = None, room_type: str = "living_room"):
    """
    根据风格生成装修效果图
    暂时使用 Stability AI（Leonardo AI Key 未配置）
    返回 JSON 格式: {"image_base64": "...", "style": "...", "layout_id": "...", "room_type": "..."}
    """
    prompt, negative_prompt = generate_prompt(style, layout_id, room_type)
    print(f"🔍 调用 Stability AI: style={style}, room_type={room_type}")
    print(f"🔍 提示词: {prompt[:200]}...")
    
    try:
        image_bytes = await generate_with_stability(prompt, negative_prompt)
        # 转换为 base64
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        return {
            "image_base64": image_base64,
            "style": style,
            "layout_id": layout_id,
            "room_type": room_type,
        }
    except Exception as e:
        print(f"❌ Stability AI 调用失败: {e}")
        # 后备：生成占位图
        return {
            "image_base64": base64.b64encode(generate_placeholder_image()).decode('utf-8'),
            "style": style,
            "layout_id": layout_id,
            "room_type": room_type,
            "error": str(e)
        }

@app.post("/inference/all")
async def inference_all(style: str = "现代简约", layout_id: str = None, room_type: str = "living_room"):
    """返回 3 套不同角度的方案（使用 Stability AI）"""
    results = []
    
    views = [
        ("客厅全景", "full room view, wide angle, corner perspective"),
        ("餐厅视角", "dining area view, warm lighting, side angle"),
        ("厨房视角", "kitchen area view, modern appliances, front view")
    ]
    
    for i, (view_name, view_desc) in enumerate(views):
        prompt, negative_prompt = generate_prompt(style, layout_id, room_type)
        full_prompt = f"{prompt}, {view_desc}, view {i+1}"
        
        image_bytes = await generate_with_stability(full_prompt, negative_prompt)
        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
        
        results.append({
            "plan_id": i + 1,
            "style": style,
            "view": view_name,
            "image_base64": image_base64,
            "description": f"{style} - {view_name}（Stable Diffusion XL 生成）"
        })
    
    return {"count": 3, "style": style, "room_type": room_type, "plans": results}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)
