from dotenv import load_dotenv
import json
import os
import base64
import asyncio
import re
import time
import requests
from pathlib import Path
from agents import Agent, trace, Runner, function_tool, gen_trace_id
from agents.model_settings import ModelSettings
from openai.types.responses import ResponseInputImageParam, ResponseInputTextParam
from openai.types.responses.response_input_item_param import Message

# Brand DNA loader function
@function_tool
def load_brand_dna(brand_dna_path: str = "Brand_DNA/brand_DNA.json") -> dict:
    """
    Load brand DNA guidelines from JSON file.

    Args:
        brand_dna_path: Path to the brand_DNA.json file

    Returns:
        Dictionary containing brand DNA guidelines

    Raises:
        FileNotFoundError: If the brand DNA file doesn't exist
        json.JSONDecodeError: If the JSON file is malformed
    """

    if not os.path.exists(brand_dna_path):
        raise FileNotFoundError(
            f"Brand DNA file not found at: {brand_dna_path}")

    with open(brand_dna_path, 'r', encoding='utf-8') as file:
        brand_dna = json.load(file)

    return brand_dna

# AI Agent prompt generator

INSTRUCTIONS = """You are an expert food photography prompt engineer specializing in Bap Kitchen's brand aesthetic.

Your task: Generate 4 diverse scene prompts for photographing a Korean dish while maintaining strict brand consistency with Bap Kitchen's HEYTEA-inspired minimalist style rooted in natural materiality.

INPUTS:
1. Raw product images (the actual dish - DO NOT alter the food itself)
2. Campaign brief
3. Brand DNA (load using load_brand_dna tool)

CRITICAL RULES:
- Keep the dish EXACTLY as shown in the raw image..... # and so on

# Image processing function
def load_product_images(product_folder: str = "Raw_products") -> list:
  
    if not os.path.exists(product_folder):
        raise FileNotFoundError(
            f"Product folder not found at: {product_folder}")

    # Supported image formats
    supported_formats = ('.jpg', '.jpeg', '.png', '.webp')

    # Find all image files
    image_files = [f for f in os.listdir(product_folder)
                   if f.lower().endswith(supported_formats)]

    if not image_files:
        raise ValueError(f"No images found in folder: {product_folder}")

    images_data = []

    for image_file in image_files:
        image_path = os.path.join(product_folder, image_file)

        try:
            # Read and encode image
            with open(image_path, 'rb') as img_file:
                image_bytes = img_file.read()
                base64_image = base64.b64encode(image_bytes).decode('utf-8')

            # Determine MIME type
            extension = os.path.splitext(image_file)[1].lower()
            mime_type = {
                '.jpg': 'image/jpeg',
                '.jpeg': 'image/jpeg',
                '.png': 'image/png',
                '.webp': 'image/webp'
            }.get(extension, 'image/jpeg')

            # Store simplified data
            image_data = {
                "filename": image_file,
                "base64_image": base64_image,
                "mime_type": mime_type
            }

            images_data.append(image_data)

        except Exception as e:
            print(f"Warning: Failed to load image {image_file}: {e}")
            continue

    if not images_data:
        raise ValueError(
            f"Failed to load any images from folder: {product_folder}")

    return images_data

# Wavespeed API batch processing
ef generate_campaign_image(prompt: str, product_image_path: str, size: str = "2227*3183") -> dict:

    API_KEY = os.getenv("WAVESPEED_API_KEY")
    if not API_KEY:
        raise ValueError(
            "WAVESPEED_API_KEY not found in environment variables")

    if not os.path.exists(product_image_path):
        raise FileNotFoundError(
            f"Product image not found at: {product_image_path}")

    # Read and encode the product image
    with open(product_image_path, 'rb') as img_file:
        image_bytes = img_file.read()
        base64_image = base64.b64encode(image_bytes).decode('utf-8')

    # Determine MIME type
    extension = os.path.splitext(product_image_path)[1].lower()
    mime_type = {
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.png': 'image/png',
        '.webp': 'image/webp'
    }.get(extension, 'image/jpeg')

    # Create data URL
    image_data_url = f"data:{mime_type};base64,{base64_image}"

    # Submit generation request
    url = "https://api.wavespeed.ai/api/v3/bytedance/seedream-v4/edit"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}",
    }
    payload = {
        "enable_base64_output": False,
        "enable_sync_mode": False,
        "images": [image_data_url],
        "prompt": prompt,
        "size": size
    }

    begin = time.time()
    response = requests.post(url, headers=headers, data=json.dumps(payload))

    if response.status_code != 200:
        return {
            "status": "failed",
            "error": f"API Error: {response.status_code} - {response.text}",
            "generation_time": time.time() - begin
        }
      [...]

# Gallery organization function
```

**5. That's it!** More or less. The workflow is fully automated and acts like a complete campaign production studio — taking seconds instead of hours.

**6. Then I built a simple frontend** where users can:
- Upload their raw product photo
- Enter their campaign prompt ("spring collection, botanical vibes")
- Click generate
- Get back 4-5 campaign-ready images

All done!

## The Complete Workflow
```
User Input
├── Raw Product Image (blank canvas)
├── Campaign Prompt (e.g., "winter campaign, fluffy snow")
└── Brand DNA Context
    ↓
AI Agent (Prompt Generator)
├── Takes: Campaign Prompt + Brand DNA
└── Outputs: 4 Different Scene Prompts
    ├── Prompt 1: Snowy Mountain Scene
    ├── Prompt 2: Cabin Interior
    ├── Prompt 3: Winter Forest
    └── Prompt 4: Cozy Fireplace
        ↓
Wavespeed API (Batch Requests)
├── Request 1: Product Image + Prompt 1 → seedream-v4/edit
├── Request 2: Product Image + Prompt 2 → seedream-v4/edit
├── Request 3: Product Image + Prompt 3 → seedream-v4/edit
└── Request 4: Product Image + Prompt 4 → seedream-v4/edit
    ↓
Generated Campaign Images (4x)