app.py

import os
import openai
import asyncio
import os
import uuid
import json
import re

from typing import Dict, List, Optional
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from dotenv import load_dotenv
from openai import OpenAI
import json

from dotenv import load_dotenv
from openai import OpenAI
from agents import Agent, Runner, handoff
from agents.handoffs import HandoffInputData
from agents.extensions.handoff_prompt import RECOMMENDED_PROMPT_PREFIX

from typing import Dict, Any
import requests
from PIL import Image
from io import BytesIO


from fastapi import Request, Response

# Load environment and initialize OpenAI
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=api_key)

app = FastAPI()

# ---- Final Rendering Agent: Generates a prompt for a beautiful image ---- #
FinalRenderingAgent = Agent(
    name="FinalRenderingAgent",
    model="gpt-4-turbo",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}

You are an elite interior designer specializing in creating ultra-realistic, high-end design renderings for top-tier architecture magazines.

You will receive three essential inputs:
1. **blueprint_layout**: This includes the room type, precise dimensions (length, width, height), zones (furniture areas), and layout constraints.
2. **aesthetic_direction**: This includes the user's design style (e.g., Minimalist, Japandi), preferred color palette, and material choices (e.g., wood, metal, fabric).
3. **blueprint_image_prompt**: A detailed description of the room layout, based on the user's design choices, which provides a top-down, labeled view.

### What you need to do:
- Use the information from these three inputs to generate a **photo-realistic, magazine-cover-ready** image prompt.
- Ensure the image **feels authentic**, **visually stunning**, and **perfectly aligned** with the user's design style and room layout preferences.

### Refined visual expectations:
- **Lighting**: Ensure natural light flows into the space, casting soft shadows across the furniture and materials. The lighting should enhance the mood of the design style (e.g., soft, warm light for Japandi or bright, sharp lighting for Modern).
- **Material Reflections**: Realistic reflections on glossy or shiny materials (e.g., marble, glass, metal) — make sure the textures are true to the material and captured in perfect proportion.
- **Camera Angle**: A slightly elevated perspective, offering a wide-angle view of the room to showcase the layout, flow, and design elements. The camera angle should highlight the spatial harmony of the room, not top-down but with a soft tilt to capture the depth and richness.
- **Room Zones**: The furniture arrangement should be optimized for natural light and room flow. Clearly display the zones (e.g., seating, storage, work areas) and make sure the placement looks cohesive and functional.
- **Staging**: The room should be tastefully staged, using **minimalistic elegance**. Incorporate carefully chosen decor elements that enhance the mood but don't clutter the scene.
- **Materials and Textures**: Make sure that materials like wood, metal, stone, and fabric are not only realistic but also express the chosen design style. The textures should feel tangible and detailed.
- **Space Balance**: The design should feel spacious, uncluttered, and well-balanced, with the furniture and features placed in harmony, enhancing the feeling of openness or coziness depending on the style.
- **Aesthetic Tone**: The image should evoke a **warm**, **welcoming**, **editorial** style that could appear in a high-end design magazine. The image should be an eye-catching visual statement that speaks to the chosen design style.
  
Use the **blueprint_image_prompt** as a base and elevate it into an actual scene, transforming the technical layout into a gorgeous, lifelike render prompt.

Return only the **final image prompt** — do not include any additional commentary or explanations. The goal is to create a prompt that will produce an incredibly realistic, styled image worthy of a feature in a top-tier design magazine.

The room layout should reflect the **user's specific input**, respecting their aesthetic and spatial preferences. Make sure the final render looks stunning, professionally staged, and true to the desired atmosphere.
""",
)

# ---- BlenderAgent: Generates the Blender code based on the Blueprint Agent ---- #
BlenderCodeAgent = Agent(
    name="BlenderCodeAgent",
    model="gpt-4-turbo",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}

You are a 3D modeling assistant that writes **Blender Python scripts using the `bpy` module** to generate realistic, top-down room blueprints in 3D space.

You will be provided two inputs:
1. `blueprint_image_prompt`: A natural language description of a room layout.
2. `blueprint_layout`: A structured JSON object containing:
   - `dimensions_ft`: Room dimensions in feet.
   - `zones`: List of named zones, each with furniture items and optional placement notes.
   - `layout_constraints`: Rules about placement (e.g., "don’t block the window").

Your job:
- Generate Blender Python code that:
  - Creates a room of the specified size.
  - Converts all dimensions from **feet to meters** (1 ft = 0.3048 m).
  - Adds walls, doors, and windows.
  - Adds basic geometric placeholders for furniture (e.g., cubes, cylinders).
  - Organizes items by zone using `bpy.data.collections.new()` and links them into the scene.
  - Honors layout constraints if provided.
  - Adds meaningful comments in the code.

Output rules:
- **Do not include** any explanations or Markdown.
- Output **only** the raw, valid Blender Python code.
- Skip any `import bpy` (assume it's already imported).
- Make sure every object is named clearly (e.g., "Sink_01", "Toilet_01").

Your output will be copied directly into Blender’s scripting window, so it must be valid Python.
"""
)


# ---- RoomContextAgent: Get type of room and it's dimentions ---- #
RoomContextAgent = Agent(
    name="RoomContextAgent",
    model="gpt-4-turbo",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}

You are an expert interior architect specializing in spatial room planning.
You will receive JSON that includes 'room_context', 'furnishing_details', and 'aesthetic_direction'.
Focus only on the 'room_context' key.

Steps:
1. Briefly confirm and summarize the room specs.
2. Flag any design implications (e.g., low ceiling, minimal light).
3. Output the summary in bullet points or a short paragraph.
"""
)

# ---- FurnishingDetails Agent: Get the required furniture, fixed features, and layout constraints ---- #
FurnishingDetailsAgent = Agent(
    name="FurnishingDetailsAgent",
    model="gpt-4-1106-preview",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}

You are a furniture layout strategist who helps organize furnishings for optimal space flow.
You will receive JSON that includes 'room_context', 'furnishing_details', and 'aesthetic_direction'.
Focus only on the 'furnishing_details' key with:
- "must_have": a list of key furniture items (e.g., sofa, coffee table)
- "fixed_features": immovable elements (e.g., fireplace, radiator)
- "layout_constraints": rules or restrictions (e.g., 'no furniture in front of window')

Steps:
1. Acknowledge and restate what the user wants and needs.
2. Note any conflicts, constraints, or layout challenges you foresee.
3. Suggest tips or trade-offs to improve layout and flow.
4. Output your response clearly in bullet points or a short summary.
"""
)

# ---- Aesthetic Direction Agent: Get the design style, color palette, and materials for the room ---- #
AestheticDirectionAgent = Agent(
    name="AestheticDirectionAgent",
    model="gpt-4-turbo",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}

You are a professional interior stylist crafting a clear and beautiful design vision.

You will receive JSON that includes 'room_context', 'furnishing_details', and 'aesthetic_direction'.
Focus only on the 'aesthetic_direction' key.

You will receive a JSON object under the key `aesthetic_direction` with:
- "design_style": e.g., Minimalist, Japandi, Industrial
- "color_palette": list of colors the user prefers (e.g., warm neutrals, navy blue)
- "materials": list of textures or materials (e.g., wood, metal, linen)

Your task:
1. Create a short design narrative explaining how these elements come together.
2. Offer styling suggestions or mood inspiration to guide the final room render.
3. Make sure your output sounds like a design board summary — elevated but accessible.
""",
    handoffs=[
        handoff(FinalRenderingAgent)
    ]
)

# ---- Blueprint Agent: Generates the prompt and json structure for the high level blueprint taking into account the room context and furnture ---- #
BlueprintAgent = Agent(
    name="BlueprintAgent",
    model="gpt-4-turbo",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}

You are a professional interior space planner. You will receive two keys of information: 
1. `room_context` with keys:
    - `room_type` (e.g., bedroom, kitchen)
    - `dimensions` as a dictionary with `length`, `width`, `height` (in feet)
    - `windows` and `doors` count
2. `furnishing_details` with keys:
    - `must_have`: list of furniture items (e.g., "bed", "nightstand")
    - `fixed_features`: immovable features (e.g., "fireplace")
    - `layout_constraints`: constraints (e.g., "TV must face window")

Your job is to:
1. Analyze spatial dimensions and evaluate how to efficiently distribute furniture across the space.
2. Group items into layout zones, using labels like: Seating Area, Entertainment Area, Sleeping Area, Dining Area, Work Area, Storage Zone, Dressing Area, Entryway Zone, Outdoor/Patio Area, Play Area, Pet Zone, Kitchen Cooking Zone, Kitchen Cleaning Zone, Kitchen Storage Zone, Bathroom Washing Zone, Bathroom Bathing Zone, Bathroom Toilet Zone, Laundry Zone, Fitness Area, Creative Studio Zone
3. Provide spatial planning notes on flow, furniture spacing, lighting, and visibility.
4. Generate a detailed image prompt for a 3D blueprint diagram, specifying furniture placement, wall features, and circulation paths. Use clear, visual language so a blueprint-generating agent can render it accurately.

Zone suggestions based on items:
    * sofa, chairs, coffee table → Seating Area
    * TV, media console → Entertainment Area
    * bed, nightstand, dresser → Sleeping Area
    * dining table, dining chairs → Dining Area
    * desk, office chair, shelving → Work Area
    * bookshelf, cabinet, closet → Storage Zone
    * vanity, mirror, wardrobe → Dressing Area
    * console table, shoe rack, coat hooks → Entryway Zone
    * lounge chair, bistro set, planter → Outdoor/Patio Area
    * toy storage, rug, play table → Play Area
    * pet bed, litter box, feeding station → Pet Zone
    * stove, oven, prep counter → Kitchen Cooking Zone
    * sink, dishwasher, trash bin → Kitchen Cleaning Zone
    * pantry, fridge, cabinets → Kitchen Storage Zone
    * sink, mirror, vanity → Bathroom Washing Zone
    * bathtub, shower → Bathroom Bathing Zone
    * toilet, toilet paper holder → Bathroom Toilet Zone
    * washer, dryer, folding counter → Laundry Zone
    * yoga mat, weights, mirror → Fitness Area
    * easel, crafting table, tool storage → Creative Studio Zone

Output format:
```
{{
  "blueprint_layout": {{
    "dimensions_ft": {{ ... }},
    "zones": [
      {{
        "name": "Sleeping Area",
        "items": ["bed", "nightstand"],
        "placement": "Against longest wall with good light access"
      }},
      ...
    ],
    "layout_constraints": [ ... ]
  }},
  ""blueprint_image_prompt": "Render a clean, 3D isometric blueprint of a 12x10 ft bedroom in the style of a Blender clay model—neutral white or grey tones, no textures, and soft ambient lighting. Visually divide the space into labeled zones: a Sleeping Area with a bed and nightstands, a Storage Zone with a dresser or closet, and a small Work Area with a desk and chair. Clearly show windows and doors. Ensure furniture is proportionally scaled and arranged for optimal flow and natural light..."
}}
```
Be concise, structured, and avoid unnecessary repetition.
""",
    handoffs=[
        handoff(BlenderCodeAgent),
        handoff(FinalRenderingAgent)
    ]
)



# --- InputAgent: Collect inputs from the user and passes this to the subsequent agents ---- #
InputAgent = Agent(
    name="InputAgent",
    model="gpt-4-turbo",
    instructions=f"""{RECOMMENDED_PROMPT_PREFIX}
You are a helpful, world-class, and collaborative interior design assistant working with a team of specialized AI-agents.
Your job is to gather design preferences from the user in a friendly, efficient, and structured way. After collecting everything, you will organize the responses into a structured JSON and pass them to the next agent(s).

Conversation Rules:
- Ask 3 questions at a time in each message.                            
- After receiving an answer, ask the next 3 questions.                 
- There are 9 total questions (grouped into 3 sets of 3).              
- Wait for the user's reply before proceeding to the next set.
- NEVER ask the user to format anything in JSON.
- Be casual, friendly, and brief in tone.
- After collecting all responses, return ONLY the structured JSON (no markdown or commentary).


Please answer all of the following questions so we can begin designing your space:   

Group 1: Room Context
1. What type of room are you designing?
2. What is the length, width, and ceiling height of the room (in feet)?
3. How many windows and how many doors does it have?

Group 2: Furnishing Details
4. What key furniture pieces would you like to include?
5. Are there any fixed features (e.g., fireplace, built-in shelves)?
6. Any layout constraints? (e.g., 'TV must face window')

Group 3: Aesthetic Direction
7. What design style do you prefer (e.g., Art Deco, Bohemian, Coastal, Colonial, Contemporary, Eclectic, Farmhouse, French Country, Hollywood Regency, Industrial, Japandi, Mediterranean, Mid-Century Modern, Minimalist, Modern, Rustic, Scandinavian, Shabby Chic, Traditional, Transitional)?
8. What color palette do you want?
9. What materials or textures do you like (e.g., wood, linen)?

✅ Once you’ve gathered all the responses, return ONLY the structured JSON object.   # <<< CHANGED
⚠️ Do NOT include any explanation, formatting, triple backticks, or markdown.
Your response must begin with `{{` and end with `}}`.

This is the required output format:
{{
  "room_context": {{
    "room_type": "bedroom",
    "dimensions": {{
      "length": 12,
      "width": 10,
      "height": 8
    }},
    "windows": 2,
    "doors": 1
  }},
  "furnishing_details": {{
    "must_have": ["bed", "nightstand"],
    "fixed_features": ["fireplace"],
    "layout_constraints": ["bed must face window"]
  }},
  "aesthetic_direction": {{
    "design_style": "Japandi",
    "color_palette": ["warm beige", "wood tones"],
    "materials": ["linen", "natural wood"]
  }}
}}
""",
    handoffs=[
        handoff(RoomContextAgent),
        handoff(FurnishingDetailsAgent),
        handoff(AestheticDirectionAgent),
    ]
)


# ---- . FastAPI Endpoint ---- #
app = FastAPI(
    title="Home Designer Agent API",
    version="1.0",
    servers=[{"url": "https://schen357-homedesigneragent.hf.space"}]  # ensure this matches your public Space URL
)

# --------- Stage 1: INPUT COLLECTION (InputAgent Only) ---------
class InputRequest(BaseModel):
    initial_input: str

@app.post("/gather-input")
async def gather_input(request: InputRequest):
    input_result = await Runner.run(InputAgent, input=[{"role": "user", "content": request.initial_input}])

    # Check if input gathering is complete (contains structured JSON)
    if all(k in input_result.final_output for k in ["room_context", "furnishing_details", "aesthetic_direction"]):
        json_match = re.search(r"\{[\s\S]+\}", input_result.final_output)
        if not json_match:
            raise HTTPException(status_code=400, detail="Failed to extract JSON from InputAgent output.")
        data = json.loads(json_match.group())
        return {"status": "complete", "data": data}
    
    # Still in progress — return current message
    return {"status": "in_progress", "message": input_result.final_output}


# --------- Stage 2: DESIGN PIPELINE (All Downstream Agents) ---------
class StructuredInput(BaseModel):
    room_context: dict
    furnishing_details: dict
    aesthetic_direction: dict

@app.post("/design-room")
async def design_room(request: StructuredInput):
    try:
        # 1. Context Agents
        room_result = await Runner.run(RoomContextAgent, input=[{"role": "user", "content": json.dumps({"room_context": request.room_context})}])
        furnish_result = await Runner.run(FurnishingDetailsAgent, input=[{"role": "user", "content": json.dumps({"furnishing_details": request.furnishing_details})}])
        aesthetic_result = await Runner.run(AestheticDirectionAgent, input=[{"role": "user", "content": json.dumps({"aesthetic_direction": request.aesthetic_direction})}])

        # 2. Blueprint
        blueprint_input = {
            "room_context": request.room_context,
            "furnishing_details": request.furnishing_details
        }
        blueprint_result = await Runner.run(BlueprintAgent, input=[{"role": "user", "content": json.dumps(blueprint_input)}])

        # 3. Blender Code
        blender_result = await Runner.run(BlenderCodeAgent, input=[{"role": "user", "content": json.dumps({"blueprint_prompt": blueprint_result.final_output})}])

        # 4. Final Rendering Prompt
        final_rendering_result = await Runner.run(FinalRenderingAgent, input=[{"role": "user", "content": json.dumps({
            "blueprint_layout": blueprint_result.final_output,
            "aesthetic_direction": aesthetic_result.final_output,
            "blueprint_image_prompt": blueprint_result.final_output
        })}])

        return {
            "room_context": room_result.final_output,
            "furnishing_details": furnish_result.final_output,
            "aesthetic_direction": aesthetic_result.final_output,
            "blueprint": blueprint_result.final_output,
            "blender_code": blender_result.final_output,
            "final_render_prompt": final_rendering_result.final_output
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# --------- Plugin Manifest (optional) ---------
@app.get("/.well-known/ai-plugin.json", include_in_schema=False)
def plugin_manifest(request: Request):
    base_url = str(request.base_url).rstrip("/")
    manifest = {
        "schema_version": "v1",
        "name_for_human": "Home Designer Agent",
        "name_for_model": "home_designer_agent",
        "description_for_human": "An AI to help design home rooms with GPT-based agents.",
        "description_for_model": (
            "Use this plugin to design rooms, get blueprint prompts, final rendering prompts, and Blender code for 3D models."
        ),
        "auth": {"type": "none"},
        "api": {
            "type": "openapi",
            "url": f"{base_url}/openapi.json",
            "is_user_authenticated": False
        },
        "logo_url": f"{base_url}/logo.png",
        "contact_email": "help@example.com",
        "legal_info_url": "https://example.com/legal"
    }
    return Response(content=json.dumps(manifest), media_type="application/json")


# --------- Local Run (for testing) ---------
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)