src/design_generation.py

import base64
import json
import mimetypes
import os
import uuid
from pathlib import Path

import replicate
from dotenv import load_dotenv
from openai import OpenAI
import zipfile
from prompt import (
    JEWELLERY_ANALYSIS_PROMPT,
    build_design_directions_prompt,
    build_design_preview_prompt,
    create_final_image_prompts,
)

load_dotenv()

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
replicate_client = replicate.Client(api_token=os.getenv("REPLICATE_API_TOKEN"))
OUTPUT_DIR = Path("generated_outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

def save_uploaded_file(uploaded_file) -> str:
    image_path = OUTPUT_DIR / uploaded_file.name

    with open(image_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    return str(image_path)


def save_uploaded_files(uploaded_files) -> list[str]:
    image_paths = []

    for uploaded_file in uploaded_files:
        image_path = OUTPUT_DIR / uploaded_file.name

        with open(image_path, "wb") as f:
            f.write(uploaded_file.getbuffer())

        image_paths.append(str(image_path))

    return image_paths

def create_zip_from_images(images: list[dict], zip_path: str) -> str:
    with zipfile.ZipFile(zip_path, "w") as zip_file:
        for img in images:
            image_path = img["path"]
            image_name = Path(image_path).name
            zip_file.write(image_path, arcname=image_name)

    return zip_path

def image_to_data_uri(image_path: str) -> str:
    mime_type, _ = mimetypes.guess_type(image_path)
    mime_type = mime_type or "image/jpeg"

    image_b64 = base64.b64encode(Path(image_path).read_bytes()).decode("utf-8")
    return f"data:{mime_type};base64,{image_b64}"


def analyze_jewellery(image_path: str) -> str:
    image_uri = image_to_data_uri(image_path)

    response = client.responses.create(
        model="gpt-5.5",
        input=[
            {
                "role": "user",
                "content": [
                    {
                        "type": "input_text",
                        "text": JEWELLERY_ANALYSIS_PROMPT,
                    },
                    {
                        "type": "input_image",
                        "image_url": image_uri,
                    },
                ],
            }
        ],
    )

    return response.output_text


def generate_design_directions(analysis: str) -> list[dict]:
    prompt = build_design_directions_prompt(analysis)

    response = client.responses.create(
        model="gpt-5.5",
        input=prompt,
        text={
            "format": {
                "type": "json_schema",
                "name": "jewellery_design_directions",
                "schema": {
                    "type": "object",
                    "properties": {
                        "directions": {
                            "type": "array",
                            "minItems": 6,
                            "maxItems": 6,
                            "items": {
                                "type": "object",
                                "properties": {
                                    "name": {"type": "string"},
                                    "design_philosophy": {"type": "string"},
                                    "what_to_keep": {"type": "string"},
                                    "what_to_change": {"type": "string"},
                                    "material_changes": {"type": "string"},
                                    "gemstone_changes": {"type": "string"},
                                    "image_generation_prompt": {"type": "string"},
                                },
                                "required": [
                                    "name",
                                    "design_philosophy",
                                    "what_to_keep",
                                    "what_to_change",
                                    "material_changes",
                                    "gemstone_changes",
                                    "image_generation_prompt",
                                ],
                                "additionalProperties": False,
                            },
                        }
                    },
                    "required": ["directions"],
                    "additionalProperties": False,
                },
                "strict": True,
            }
        },
    )

    data = json.loads(response.output_text)

    if "directions" not in data or not data["directions"]:
        raise ValueError("No design directions generated.")

    return data["directions"]


def generate_image(reference_images: list[str], prompt: str, output_path: str) -> str:
    image_inputs = [image_to_data_uri(path) for path in reference_images]

    output = replicate_client.run(
        "google/nano-banana-2",
        input={
            "prompt": prompt,
            "resolution": "2K",
            "image_input": image_inputs,
            "aspect_ratio": "1:1",
            "image_search": False,
            "google_search": False,
            "output_format": "jpg",
        },
    )

    with open(output_path, "wb") as file:
        file.write(output.read())

    return output_path

def generate_six_campaign_images(
    reference_images: list[str],
    analysis: str,
    directions: list[dict],
) -> list[dict]:
    if not directions:
        raise ValueError("directions is empty. generate_campaign_directions() returned None or [].")

    designs = []
    run_id = uuid.uuid4().hex[:8]

    for i, direction in enumerate(directions, start=1):
        prompt = build_design_preview_prompt(direction, analysis)
        output_path = str(OUTPUT_DIR / f"campaign_option_{run_id}_{i}.jpg")

        generate_image(
            reference_images=reference_images,
            prompt=prompt,
            output_path=output_path,
        )

        designs.append(
            {
                "name": direction.get("name", f"Design {i}"),
                "path": output_path,
                "prompt": prompt,
                "direction": direction,
            }
        )

    return designs


def generate_final_images(
    reference_images: list[str],
    selected_campaign: dict,
    analysis: str,
    user_prompt: str = "",
) -> list[dict]:
    final_prompts = create_final_image_prompts(
        selected_campaign=selected_campaign,
        analysis=analysis,
        user_prompt=user_prompt,
    )

    final_images = []
    run_id = uuid.uuid4().hex[:8]

    for i, item in enumerate(final_prompts, start=1):
        image_type = item["type"].lower().replace(" ", "_")
        output_path = str(OUTPUT_DIR / f"final_{run_id}_{i}_{image_type}.jpg")

        generate_image(
            reference_images=reference_images,
            prompt=item["prompt"],
            output_path=output_path,
        )

        final_images.append(
            {
                "type": item["type"],
                "path": output_path,
                "prompt": item["prompt"],
            }
        )

    return final_images